Merge pull request #190 from SCAI-BIO/refactor/remove-sqlite

refactor: remove sqlite

Author: mehmetcanay

datastew/__init__.py

@@ -7,11 +7,9 @@
     Concept,
     Mapping,
     Terminology,
-    base,
     model,
     pagination,
     postgresql,
-    sqllite,
 )
 from .visualisation import (
     bar_chart_average_acc_two_distributions,
@@ -26,11 +24,9 @@
     "mapping",
     "ols",
     "parsing",
-    "base",
     "model",
     "pagination",
     "postgresql",
-    "sqllite",
     "EmbeddingModel",
     "GPT4Adapter",
     "HuggingFaceAdapter",

datastew/process/jsonl_adapter.py

@@ -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
 
 
@@ -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.

datastew/process/ols.py

@@ -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:
@@ -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.
 

datastew/repository/__init__.py

@@ -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"]

datastew/repository/base.py

@@ -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,
@@ -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
         )
@@ -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
     ):
@@ -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