ebm4subjects 0.5.4__py3-none-any.whl → 0.5.6__py3-none-any.whl

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/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,
 )
 
@@ -44,12 +44,13 @@ class EbmModel:
         use_altLabels: bool = True,
         hnsw_index_params: dict | str | None = None,
         embedding_model_name: str | None = None,
-        embedding_model_deployment: str = "offline-inference",
+        embedding_model_deployment: str = "mock",
         embedding_model_args: dict | str | None = None,
         encode_args_vocab: dict | str | None = None,
         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
@@ -100,7 +101,7 @@ class EbmModel:
 
         # Parameters for embedding generator
         self.generator = None
-        self.embedding_model_deployment = embedding_model_deployment
+        self.embedding_model_deployment = embedding_model_deployment.lower()
         self.embedding_model_name = embedding_model_name
         self.embedding_dimensions = int(embedding_dimensions)
         if isinstance(embedding_model_args, str) or not embedding_model_args:
@@ -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,36 +181,41 @@ class EbmModel:
             None
         """
         if self.generator is None:
-            if self.embedding_model_deployment == "offline-inference":
-                self.logger.info("initializing offline-inference embedding generator")
-                self.generator = EmbeddingGeneratorOfflineInference(
+            if self.embedding_model_deployment == "in-process":
+                self.logger.info("initializing in-process embedding generator")
+                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":
                 self.logger.info("initializing mock embedding generator")
                 self.generator = EmbeddingGeneratorMock(self.embedding_dimensions)
-            elif self.embedding_model_deployment == "HuggingFaceTEI":
+            elif self.embedding_model_deployment == "huggingfacetei":
                 self.logger.info("initializing API embedding generator")
                 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":
+            elif self.embedding_model_deployment == "openai":
                 self.logger.info("initializing API embedding generator")
                 self.generator = EmbeddingGeneratorOpenAI(
                     model_name=self.embedding_model_name,
                     embedding_dimensions=self.embedding_dimensions,
+                    logger=self.logger,
                     **self.embedding_model_args,
                 )
             else:
-                self.logger.error("unsupportet API for embedding generator")
-                raise NotImplementedError
+                raise NotImplementedError("Unsupportet API for embedding generator")
 
     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 +224,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:
@@ -663,7 +669,7 @@ class EbmModel:
             )
             self.logger.info("training successful finished")
         except xgb.core.XGBoostError:
-            self.logger.critical(
+            self.logger.warn(
                 "XGBoost can't train with candidates equal to gold standard "
                 "or candidates with no match to gold standard at all - "
                 "Check if your training data and gold standard are correct"
@@ -769,7 +775,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/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,26 @@ 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."
+            )
+        else:
+            raise RuntimeError(
+                "Request to API not possible! Please check the corresponding parameters!"
+            )
 
     def generate_embeddings(self, texts: list[str], **kwargs) -> np.ndarray:
         """
@@ -72,8 +92,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 +121,70 @@ 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
+        """
+        _ = 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.")
+
     def generate_embeddings(self, texts: list[str], **kwargs) -> np.ndarray:
         """
         Generates embeddings for a list of input texts using a model
@@ -122,8 +192,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 +212,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 +268,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

{ebm4subjects-0.5.4.dist-info → ebm4subjects-0.5.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ebm4subjects
-Version: 0.5.4
+Version: 0.5.6
 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.6.dist-info/RECORD

@@ -0,0 +1,12 @@
+ebm4subjects/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ebm4subjects/analyzer.py,sha256=lqX7AF8WsvwIavgtnmoVQ0i3wzBJJSeH47EiEwoLKGg,1664
+ebm4subjects/chunker.py,sha256=HcEFJtKWHFYZL8DmZcHGXLPGEkCqHZhh_0kSqyYVsdE,6764
+ebm4subjects/duckdb_client.py,sha256=8lDIpj2o2VTEtjHC_vTYrI5-RNXZnWMft45bS6z9B_k,13031
+ebm4subjects/ebm_logging.py,sha256=vGMa3xSm6T7ZQ94XeNGJVGCTl3zytt4sbunwXc6qF5U,5987
+ebm4subjects/ebm_model.py,sha256=UTCIv_KCQ4HTJVbcVIAUv4S2j87oq8HXBeN5mfJmclQ,30879
+ebm4subjects/embedding_generator.py,sha256=fk8rRhqBcRCknpCYoFolcXjoCwsx25Qd_UEOt-nUlv8,11774
+ebm4subjects/prepare_data.py,sha256=vQ-BdXkIP3iZJdPXol0WDlY8cRFMHkjzzL7oC7EbouE,3084
+ebm4subjects-0.5.6.dist-info/METADATA,sha256=Dujb7SghFPo3j42yRAgkbqv-VSmwpocJIHW4NgJFhn0,8354
+ebm4subjects-0.5.6.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+ebm4subjects-0.5.6.dist-info/licenses/LICENSE,sha256=RpvAZSjULHvoTR_esTlucJ08-zdQydnoqQLbqOh9Ub8,13826
+ebm4subjects-0.5.6.dist-info/RECORD,,

ebm4subjects-0.5.4.dist-info/RECORD

@@ -1,12 +0,0 @@
-ebm4subjects/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ebm4subjects/analyzer.py,sha256=lqX7AF8WsvwIavgtnmoVQ0i3wzBJJSeH47EiEwoLKGg,1664
-ebm4subjects/chunker.py,sha256=HcEFJtKWHFYZL8DmZcHGXLPGEkCqHZhh_0kSqyYVsdE,6764
-ebm4subjects/duckdb_client.py,sha256=8lDIpj2o2VTEtjHC_vTYrI5-RNXZnWMft45bS6z9B_k,13031
-ebm4subjects/ebm_logging.py,sha256=xkbqeVhSCNuhMwkx2yoIX8_D3z9DcsauZEmHhR1gaS0,5962
-ebm4subjects/ebm_model.py,sha256=lzGx_HLkKyTPVhtU4117DOEDz1rduNdzltvCYSbHQPg,30780
-ebm4subjects/embedding_generator.py,sha256=LKZ_YAe4Th8foI_8-v-3tYFj0KGJ90XJ3OPuMXaqgSQ,9274
-ebm4subjects/prepare_data.py,sha256=vQ-BdXkIP3iZJdPXol0WDlY8cRFMHkjzzL7oC7EbouE,3084
-ebm4subjects-0.5.4.dist-info/METADATA,sha256=OmMMh0pGAdv3YTkTork55wuj2gA0Ac8zV9ad3cDCIks,8274
-ebm4subjects-0.5.4.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-ebm4subjects-0.5.4.dist-info/licenses/LICENSE,sha256=RpvAZSjULHvoTR_esTlucJ08-zdQydnoqQLbqOh9Ub8,13826
-ebm4subjects-0.5.4.dist-info/RECORD,,