ebm4subjects 0.5.4__tar.gz → 0.5.5__tar.gz

{ebm4subjects-0.5.4 → ebm4subjects-0.5.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ebm4subjects
-Version: 0.5.4
+Version: 0.5.5
 Summary: Embedding Based Matching for Automated Subject Indexing
 Author: Deutsche Nationalbibliothek
 Maintainer-email: Clemens Rietdorf <c.rietdorf@dnb.de>, Maximilian Kähler <m.kaehler@dnb.de>
@@ -14,12 +14,14 @@ Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.10
 Requires-Dist: duckdb>=1.3.0
 Requires-Dist: nltk~=3.9.1
+Requires-Dist: openai>=2.15.0
 Requires-Dist: polars>=1.30.0
 Requires-Dist: pyarrow>=21.0.0
 Requires-Dist: pyoxigraph>=0.4.11
 Requires-Dist: rdflib~=7.1.3
-Requires-Dist: sentence-transformers>=5.0.0
 Requires-Dist: xgboost>=3.0.2
+Provides-Extra: in-process
+Requires-Dist: sentence-transformers>=5.0.0; extra == 'in-process'
 Description-Content-Type: text/markdown
 
 # Embedding Based Matching for Automated Subject Indexing

{ebm4subjects-0.5.4 → ebm4subjects-0.5.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ebm4subjects"
-version = "0.5.4"
+version = "0.5.5"
 description = "Embedding Based Matching for Automated Subject Indexing"
 authors = [
     {name = "Deutsche Nationalbibliothek"},
@@ -29,13 +29,15 @@ requires-python = ">=3.10"
 dependencies = [
     "duckdb>=1.3.0",
     "nltk~=3.9.1",
+    "openai>=2.15.0",
     "polars>=1.30.0",
     "pyarrow>=21.0.0",
     "pyoxigraph>=0.4.11",
     "rdflib~=7.1.3",
-    "sentence-transformers>=5.0.0",
     "xgboost>=3.0.2",
 ]
+[project.optional-dependencies]
+in-process=["sentence-transformers>=5.0.0"]
 
 [build-system]
 requires = ["hatchling"]

{ebm4subjects-0.5.4 → ebm4subjects-0.5.5}/src/ebm4subjects/ebm_logging.py

@@ -39,17 +39,17 @@ class EbmLogger:
         else:
             self.logger.setLevel(logging.NOTSET)
 
-        # Create a file handler to log messages to a file
-        log_file_handler = logging.FileHandler(f"{log_path}/ebm.log")
-        log_file_handler.setFormatter(
-            logging.Formatter(
-                "%(asctime)s %(levelname)s: %(message)s",
-                "%Y-%m-%d %H:%M:%S",
+        # Create a file handler to log messages to a file 
+        if not self.logger.handlers:
+            log_file_handler = logging.FileHandler(f"{log_path}/ebm.log")
+            log_file_handler.setFormatter(
+                logging.Formatter(
+                    "%(asctime)s %(levelname)s: %(message)s",
+                    "%Y-%m-%d %H:%M:%S",
+                )
             )
-        )
 
-        # Add the file handler to the logger
-        self.logger.addHandler(log_file_handler)
+            self.logger.addHandler(log_file_handler)
 
     def get_logger(self) -> logging.Logger:
         """

{ebm4subjects-0.5.4 → ebm4subjects-0.5.5}/src/ebm4subjects/ebm_model.py

@@ -16,7 +16,7 @@ from ebm4subjects.ebm_logging import EbmLogger, NullLogger, XGBLogging
 from ebm4subjects.embedding_generator import (
     EmbeddingGeneratorHuggingFaceTEI,
     EmbeddingGeneratorMock,
-    EmbeddingGeneratorOfflineInference,
+    EmbeddingGeneratorInProcess,
     EmbeddingGeneratorOpenAI,
 )
 
@@ -50,6 +50,7 @@ class EbmModel:
         encode_args_documents: dict | str | None = None,
         log_path: str | None = None,
         logger: logging.Logger | None = None,
+        logging_level: str = "info",
     ) -> None:
         """
         A class representing an Embedding-Based-Matching (EBM) model
@@ -139,7 +140,7 @@ class EbmModel:
         self.train_jobs = int(xgb_jobs)
 
         # Initiliaze logging
-        self.init_logger(log_path, logger)
+        self.init_logger(log_path, logger, logging_level)
 
         # Initialize EBM model
         self.model = None
@@ -180,11 +181,12 @@ class EbmModel:
             None
         """
         if self.generator is None:
-            if self.embedding_model_deployment == "offline-inference":
+            if self.embedding_model_deployment == "in-process":
                 self.logger.info("initializing offline-inference embedding generator")
-                self.generator = EmbeddingGeneratorOfflineInference(
+                self.generator = EmbeddingGeneratorInProcess(
                     model_name=self.embedding_model_name,
                     embedding_dimensions=self.embedding_dimensions,
+                    logger=self.logger,
                     **self.embedding_model_args,
                 )
             elif self.embedding_model_deployment == "mock":
@@ -195,6 +197,7 @@ class EbmModel:
                 self.generator = EmbeddingGeneratorHuggingFaceTEI(
                     model_name=self.embedding_model_name,
                     embedding_dimensions=self.embedding_dimensions,
+                    logger=self.logger,
                     **self.embedding_model_args,
                 )
             elif self.embedding_model_deployment == "OpenAI":
@@ -202,6 +205,7 @@ class EbmModel:
                 self.generator = EmbeddingGeneratorOpenAI(
                     model_name=self.embedding_model_name,
                     embedding_dimensions=self.embedding_dimensions,
+                    logger=self.logger,
                     **self.embedding_model_args,
                 )
             else:
@@ -209,7 +213,10 @@ class EbmModel:
                 raise NotImplementedError
 
     def init_logger(
-        self, log_path: str | None = None, logger: logging.Logger | None = None
+        self,
+        log_path: str | None = None,
+        logger: logging.Logger | None = None,
+        logging_level: str = "info",
     ) -> None:
         """
         Initializes the logging for the EBM model.
@@ -218,7 +225,7 @@ class EbmModel:
             None
         """
         if log_path:
-            self.logger = EbmLogger(log_path, "info").get_logger()
+            self.logger = EbmLogger(log_path, logging_level).get_logger()
             self.xgb_logger = XGBLogging(self.logger, epoch_log_interval=1)
             self.xgb_callbacks = [self.xgb_logger]
         elif logger:
@@ -769,7 +776,4 @@ class EbmModel:
         Returns:
             EbmModel: The loaded EBM model instance.
         """
-        ebm_model = joblib.load(input_path)
-        ebm_model.init_logger()
-
-        return ebm_model
+        return joblib.load(input_path)

{ebm4subjects-0.5.4 → ebm4subjects-0.5.5}/src/ebm4subjects/embedding_generator.py

@@ -1,8 +1,9 @@
+import logging
 import os
 
 import numpy as np
 import requests
-from sentence_transformers import SentenceTransformer
+from openai import BadRequestError, NotFoundError, OpenAI
 from tqdm import tqdm
 
 
@@ -32,25 +33,29 @@ class EmbeddingGenerator:
         pass
 
 
-class EmbeddingGeneratorAPI(EmbeddingGenerator):
+class EmbeddingGeneratorHuggingFaceTEI(EmbeddingGenerator):
     """
-    A base class for API embedding generators.
-
-    Attributes:
-        embedding_dimensions (int): The dimensionality of the generated embeddings.
+    A class for generating embeddings using the HuggingFaceTEI API.
     """
 
     def __init__(
         self,
         model_name: str,
         embedding_dimensions: int,
+        logger: logging.Logger,
         **kwargs,
     ) -> None:
         """
-        Initializes the API EmbeddingGenerator.
+        Initializes the HuggingFaceTEI API EmbeddingGenerator.
 
         Sets the embedding dimensions, and initiliazes and
         prepares a session with the API.
+
+        Args:
+            model_name (str): The name of the SentenceTransformer model.
+            embedding_dimensions (int): The dimensionality of the generated embeddings.
+            logger (Logger): A logger for the embedding generator.
+            **kwargs: Additional keyword arguments to pass to the model.
         """
 
         self.embedding_dimensions = embedding_dimensions
@@ -59,11 +64,36 @@ class EmbeddingGeneratorAPI(EmbeddingGenerator):
         self.api_address = kwargs.get("api_address")
         self.headers = kwargs.get("headers", {"Content-Type": "application/json"})
 
+        self.logger = logger
+        self._test_api()
 
-class EmbeddingGeneratorHuggingFaceTEI(EmbeddingGeneratorAPI):
-    """
-    A class for generating embeddings using the HuggingFaceTEI API.
-    """
+    def _test_api(self):
+        """
+        Tests if the API is working with the given parameters
+        """
+        response = self.session.post(
+            self.api_address,
+            headers=self.headers,
+            json={"inputs": "This is a test request!", "truncate": True},
+        )
+        if response.status_code == 200:
+            self.logger.debug(
+                "API call successful. Everything seems to be working fine."
+            )
+        elif response.status_code == 404:
+            self.logger.error(
+                "API not found under given adress! Please check the corresponding parameter!"
+            )
+            raise RuntimeError(
+                "API not found under given adress! Please check the corresponding parameter!"
+            )
+        else:
+            self.logger.error(
+                "Request to API not possible! Please check the corresponding parameters!"
+            )
+            raise RuntimeError(
+                "Request to API not possible! Please check the corresponding parameters!"
+            )
 
     def generate_embeddings(self, texts: list[str], **kwargs) -> np.ndarray:
         """
@@ -72,8 +102,7 @@ class EmbeddingGeneratorHuggingFaceTEI(EmbeddingGeneratorAPI):
 
         Args:
             texts (list[str]): A list of input texts.
-            **kwargs: Additional keyword arguments to pass to the
-                SentenceTransformer model.
+            **kwargs: Additional keyword arguments to pass to the API.
 
         Returns:
             np.ndarray: A numpy array of shape (len(texts), embedding_dimensions)
@@ -102,19 +131,87 @@ class EmbeddingGeneratorHuggingFaceTEI(EmbeddingGeneratorAPI):
             if response.status_code == 200:
                 embeddings.extend(response.json())
             else:
-                # TODO: write warning to logger
+                self.logger.warn("Call to API NOT successful! Returning 0's.")
                 for _ in batch_texts:
-                    # TODO: ensure same format as true case and truncate dim
-                    embeddings.append([0 for _ in range(self.embedding_dimensions)])
+                    embeddings.append(
+                        [
+                            0
+                            for _ in range(
+                                min(
+                                    self.embedding_dimensions,
+                                    kwargs.get("truncate_prompt_tokens", float("inf")),
+                                ),
+                            )
+                        ]
+                    )
 
         return np.array(embeddings)
 
 
-class EmbeddingGeneratorOpenAI(EmbeddingGeneratorAPI):
+class EmbeddingGeneratorOpenAI(EmbeddingGenerator):
     """
-    A class for generating embeddings using any OpenAI compatibleAPI.
+    A class for generating embeddings using any OpenAI compatible API.
     """
 
+    def __init__(
+        self,
+        model_name: str,
+        embedding_dimensions: int,
+        logger: logging.Logger,
+        **kwargs,
+    ) -> None:
+        """
+        Initializes the OpenAI API EmbeddingGenerator.
+
+        Sets the embedding dimensions, and initiliazes and
+        prepares a session with the API.
+
+        Args:
+            model_name (str): The name of the SentenceTransformer model.
+            embedding_dimensions (int): The dimensionality of the generated embeddings.
+            logger (Logger): A logger for the embedding generator.
+            **kwargs: Additional keyword arguments to pass to the model.
+        """
+
+        self.embedding_dimensions = embedding_dimensions
+        self.model_name = model_name
+
+        if not (api_key := os.environ.get("OPENAI_API_KEY")):
+            api_key = ""
+
+        self.client = OpenAI(api_key=api_key, base_url=kwargs.get("api_address"))
+
+        self.logger = logger
+        self._test_api()
+
+    def _test_api(self):
+        """
+        Tests if the API is working with the given parameters
+        """
+        try:
+            _ = self.client.embeddings.create(
+                input="This is a test request!",
+                model=self.model_name,
+                encoding_format="float",
+            )
+            self.logger.debug(
+                "API call successful. Everything seems to be working fine."
+            )
+        except NotFoundError:
+            self.logger.error(
+                "API not found under given adress! Please check the corresponding parameter!"
+            )
+            raise RuntimeError(
+                "API not found under given adress! Please check the corresponding parameter!"
+            )
+        except BadRequestError:
+            self.logger.error(
+                "Request to API not possible! Please check the corresponding parameters!"
+            )
+            raise RuntimeError(
+                "Request to API not possible! Please check the corresponding parameters!"
+            )
+
     def generate_embeddings(self, texts: list[str], **kwargs) -> np.ndarray:
         """
         Generates embeddings for a list of input texts using a model
@@ -122,8 +219,7 @@ class EmbeddingGeneratorOpenAI(EmbeddingGeneratorAPI):
 
         Args:
             texts (list[str]): A list of input texts.
-            **kwargs: Additional keyword arguments to pass to the
-                SentenceTransformer model.
+            **kwargs: Additional keyword arguments to pass to the API.
 
         Returns:
             np.ndarray: A numpy array of shape (len(texts), embedding_dimensions)
@@ -143,53 +239,54 @@ class EmbeddingGeneratorOpenAI(EmbeddingGeneratorAPI):
 
         for i in tqdm(range(0, len(texts), batch_size), desc="Processing batches"):
             batch_texts = texts[i : i + batch_size]
-            data = {
-                "input": batch_texts,
-                "model": self.model_name,
-                "encoding_format": "float",
-                **kwargs,
-            }
 
-            response = self.session.post(
-                self.api_address, headers=self.headers, json=data
-            )
+            # Try to get embeddings for the batch from the API
+            try:
+                embedding_response = self.client.embeddings.create(
+                    input=batch_texts,
+                    model=self.model_name,
+                    encoding_format="float",
+                    extra_body={**kwargs},
+                )
 
-            # Process all embeddings from the batch response
-            if response.status_code == 200:
-                response_data = response.json()
+                # Process all embeddings from the batch response
                 for i, _ in enumerate(batch_texts):
-                    embedding = response_data["data"][i]["embedding"]
-                    embeddings.append(embedding)
-            else:
-                # TODO: write warning to logger
+                    embeddings.append(embedding_response.data[i].embedding)
+            except (NotFoundError, BadRequestError):
+                self.logger.warn("Call to API NOT successful! Returning 0's.")
                 for _ in batch_texts:
                     embeddings.append([0 for _ in range(self.embedding_dimensions)])
 
         return np.array(embeddings)
 
 
-class EmbeddingGeneratorOfflineInference(EmbeddingGenerator):
+class EmbeddingGeneratorInProcess(EmbeddingGenerator):
     """
     A class for generating embeddings using a given SentenceTransformer model
-    loaded offline with SentenceTransformer.
+    loaded in-process with SentenceTransformer.
 
     Args:
         model_name (str): The name of the SentenceTransformer model.
         embedding_dimensions (int): The dimensionality of the generated embeddings.
+        logger (Logger): A logger for the embedding generator.
         **kwargs: Additional keyword arguments to pass to the model.
-
-    Attributes:
-        model_name (str): The name of the SentenceTransformer model.
-        embedding_dimensions (int): The dimensionality of the generated embeddings.
     """
 
-    def __init__(self, model_name: str, embedding_dimensions: int, **kwargs) -> None:
+    def __init__(
+        self,
+        model_name: str,
+        embedding_dimensions: int,
+        logger: logging.Logger,
+        **kwargs,
+    ) -> None:
         """
-        Initializes the EmbeddingGenerator in offline inference mode.
+        Initializes the EmbeddingGenerator in 'in-process' mode.
 
         Sets the model name, embedding dimensions, and creates a
         SentenceTransformer model instance.
         """
+        from sentence_transformers import SentenceTransformer
+
         self.model_name = model_name
         self.embedding_dimensions = embedding_dimensions
 
@@ -198,6 +295,8 @@ class EmbeddingGeneratorOfflineInference(EmbeddingGenerator):
         self.model = SentenceTransformer(
             model_name, truncate_dim=embedding_dimensions, **kwargs
         )
+        self.logger = logger
+        self.logger.debug(f"SentenceTransfomer model running on {self.model.device}")
 
         # Disabel parallelism for tokenizer
         # Needed because process might be already parallelized