Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 0 additions & 4 deletions datastew/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,11 +7,9 @@
Concept,
Mapping,
Terminology,
base,
model,
pagination,
postgresql,
sqllite,
)
from .visualisation import (
bar_chart_average_acc_two_distributions,
Expand All @@ -26,11 +24,9 @@
"mapping",
"ols",
"parsing",
"base",
"model",
"pagination",
"postgresql",
"sqllite",
"EmbeddingModel",
"GPT4Adapter",
"HuggingFaceAdapter",
Expand Down
4 changes: 2 additions & 2 deletions datastew/process/jsonl_adapter.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
from tqdm import tqdm

from datastew.embedding import Vectorizer
from datastew.repository import PostgreSQLRepository, SQLLiteRepository
from datastew.repository import PostgreSQLRepository
from datastew.repository.model import Concept, Mapping, Terminology


Expand Down Expand Up @@ -69,7 +69,7 @@ def _flush_to_file(self, file_path: str):

self._buffer.clear()

def from_repository(self, repository: Union[PostgreSQLRepository, SQLLiteRepository]):
def from_repository(self, repository: PostgreSQLRepository):
"""Export all records from a PostgreSQLRepository to JSONL files

:param repository: Active database repository instace.
Expand Down
4 changes: 2 additions & 2 deletions datastew/process/ols.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@
import requests

from datastew.embedding import Vectorizer
from datastew.repository.base import BaseRepository
from datastew.repository.model import Concept, Mapping, Terminology
from datastew.repository.postgresql import PostgreSQLRepository


class OLSTerminologyImportTask:
Expand Down Expand Up @@ -83,7 +83,7 @@ def __process_page(self, page: int) -> Tuple[List[Concept], List[Mapping]]:
logging.error(f"Failed to fetch concepts and descriptions from OLS for page {page}: {str(e)}")
return [], []

def process_to_repository(self, repository: BaseRepository):
def process_to_repository(self, repository: PostgreSQLRepository):
"""
Fetches concepts and descriptions from the OLS API and stores them in a repository.

Expand Down
3 changes: 1 addition & 2 deletions datastew/repository/__init__.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
from .model import Concept, Mapping, Terminology
from .postgresql import PostgreSQLRepository
from .sqllite import SQLLiteRepository

__all__ = ["Terminology", "Concept", "Mapping", "SQLLiteRepository", "PostgreSQLRepository"]
__all__ = ["Terminology", "Concept", "Mapping", "PostgreSQLRepository"]
124 changes: 0 additions & 124 deletions datastew/repository/base.py

This file was deleted.

44 changes: 41 additions & 3 deletions datastew/repository/postgresql.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,14 +7,14 @@

from datastew.embedding import Vectorizer
from datastew.exceptions import ObjectStorageError
from datastew.repository.base import BaseRepository
from datastew.process.parsing import DataDictionarySource
from datastew.repository.model import Base, Concept, Mapping, MappingResult, Terminology
from datastew.repository.pagination import Page

logger = logging.getLogger(__name__)


class PostgreSQLRepository(BaseRepository):
class PostgreSQLRepository:
def __init__(
self,
connection_string: str,
Expand All @@ -33,7 +33,7 @@ def __init__(
:param pool_timeout: The maximum time (in seconds) to wait for a connection from
the pool before raising an exception.
"""
super().__init__(vectorizer)
self.vectorizer = vectorizer
self.engine = create_engine(
connection_string, pool_size=pool_size, max_overflow=max_overflow, pool_timeout=pool_timeout
)
Expand Down Expand Up @@ -197,6 +197,20 @@ def clear_all(self):
self.session.query(Terminology).delete()
self.session.commit()

def import_data_dictionary(self, data_dictionary: DataDictionarySource, terminology_name: str):
"""Imports a data dictionary, generating concepts and embeddings, and stores them in the database.

:param data_dictionary: Source of variable descriptions and metadata.
:param terminology_name: Name of the terminology being imported.
:raises RuntimeError: If the import or transformation fails.
"""
try:
objects = self._parse_data_dictionary(data_dictionary, terminology_name)
self.store_all(objects)
except Exception as e:
logger.exception("Failed to import data dictionary.")
raise RuntimeError(f"Failed to import data dictionary source: {e}")

def import_from_jsonl(
self, jsonl_path: str, object_type: Literal["terminology", "concept", "mapping"], chunk_size: int = 100
):
Expand Down Expand Up @@ -305,3 +319,27 @@ def _validate_required_fields(
def _initialize_pgvector(self):
with self.engine.begin() as conn:
conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))

def _parse_data_dictionary(
self, data_dictionary: DataDictionarySource, terminology_name: str
) -> List[Union[Concept, Mapping, Terminology]]:
df = data_dictionary.to_dataframe()
descriptions = df["description"].tolist()
vectorizer_name = self.vectorizer.model_name
variable_to_embedding = data_dictionary.get_embeddings(self.vectorizer)

terminology = Terminology(name=terminology_name, id=terminology_name)
objects: List[Union[Concept, Mapping, Terminology]] = [terminology]

for variable, description in zip(variable_to_embedding.keys(), descriptions):
concept_id = f"{terminology_name}:{variable}"
concept = Concept(terminology=terminology, pref_label=variable, concept_identifier=concept_id)
mapping = Mapping(
concept=concept,
text=description,
embedding=variable_to_embedding[variable],
sentence_embedder=vectorizer_name,
)
objects.extend([concept, mapping])

return objects
Loading