ebm4subjects 0.5.3__py3-none-any.whl → 0.5.5__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

@@ -15,8 +15,9 @@ from ebm4subjects.duckdb_client import Duckdb_client
 from ebm4subjects.ebm_logging import EbmLogger, NullLogger, XGBLogging
 from ebm4subjects.embedding_generator import (
     EmbeddingGeneratorHuggingFaceTEI,
-    EmbeddingGeneratorOfflineInference,
     EmbeddingGeneratorMock,
+    EmbeddingGeneratorInProcess,
+    EmbeddingGeneratorOpenAI,
 )
 
 
@@ -43,12 +44,13 @@ class EbmModel:
         use_altLabels: bool = True,
         hnsw_index_params: dict | str | None = None,
         embedding_model_name: str | None = None,
-        embedding_model_type: str = "offline-inference",
+        embedding_model_deployment: str = "offline-inference",
         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
@@ -99,7 +101,7 @@ class EbmModel:
 
         # Parameters for embedding generator
         self.generator = None
-        self.embedding_model_type = embedding_model_type
+        self.embedding_model_deployment = embedding_model_deployment
         self.embedding_model_name = embedding_model_name
         self.embedding_dimensions = int(embedding_dimensions)
         if isinstance(embedding_model_args, str) or not embedding_model_args:
@@ -138,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
@@ -179,20 +181,31 @@ class EbmModel:
             None
         """
         if self.generator is None:
-            if self.embedding_model_type == "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_type == "mock":
+            elif self.embedding_model_deployment == "mock":
                 self.logger.info("initializing mock embedding generator")
                 self.generator = EmbeddingGeneratorMock(self.embedding_dimensions)
-            elif self.embedding_model_type == "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":
+                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:
@@ -200,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.
@@ -209,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:
@@ -760,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/embedding_generator.py

@@ -1,8 +1,10 @@
+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
 
 
 class EmbeddingGenerator:
@@ -31,37 +33,67 @@ 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
-
+        self.model_name = model_name
         self.session = requests.Session()
         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:
         """
@@ -70,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)
@@ -85,45 +116,177 @@ class EmbeddingGeneratorHuggingFaceTEI(EmbeddingGeneratorAPI):
             # If empty, return an empty numpy array with the correct shape
             return np.empty((0, self.embedding_dimensions))
 
-        # process each text
-        for text in texts:
+        # Process in smaller batches to avoid memory overload
+        batch_size = min(32, len(texts))  # HuggingFaceTEI has a limit of 32 as default
+
+        for i in tqdm(range(0, len(texts), batch_size), desc="Processing batches"):
+            batch_texts = texts[i : i + batch_size]
             # send a request to the HuggingFaceTEI API
-            data = {"inputs": text}
+            data = {"inputs": batch_texts, "truncate": True}
             response = self.session.post(
                 self.api_address, headers=self.headers, json=data
             )
 
             # add generated embeddings to return list if request was successfull
             if response.status_code == 200:
-                embeddings.append(response.json()[0])
+                embeddings.extend(response.json())
             else:
-                embeddings.append([0 for _ in range(self.embedding_dimensions)])
+                self.logger.warn("Call to API NOT successful! Returning 0's.")
+                for _ in batch_texts:
+                    embeddings.append(
+                        [
+                            0
+                            for _ in range(
+                                min(
+                                    self.embedding_dimensions,
+                                    kwargs.get("truncate_prompt_tokens", float("inf")),
+                                ),
+                            )
+                        ]
+                    )
+
+        return np.array(embeddings)
+
+
+class EmbeddingGeneratorOpenAI(EmbeddingGenerator):
+    """
+    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
+        via an OpenAI compatible API.
+
+        Args:
+            texts (list[str]): A list of input texts.
+            **kwargs: Additional keyword arguments to pass to the API.
+
+        Returns:
+            np.ndarray: A numpy array of shape (len(texts), embedding_dimensions)
+                containing the generated embeddings.
+        """
+        # prepare list for return
+        embeddings = []
+
+        # Check if the input list is empty
+        if not texts:
+            # If empty, return an empty numpy array with the correct shape
+            return np.empty((0, self.embedding_dimensions))
+
+        # Process in smaller batches to avoid memory overload
+        batch_size = min(200, len(texts))
+        embeddings = []
+
+        for i in tqdm(range(0, len(texts), batch_size), desc="Processing batches"):
+            batch_texts = texts[i : i + batch_size]
+
+            # 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
+                for i, _ in enumerate(batch_texts):
+                    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
 
@@ -132,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

{ebm4subjects-0.5.3.dist-info → ebm4subjects-0.5.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ebm4subjects
-Version: 0.5.3
+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.5.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=mnnRqdO3vlRF0MTNx7JvHgTXPCNg-8YWJm1kOtHinak,30929
+ebm4subjects/embedding_generator.py,sha256=q5HP36q11EMkH_yomduXa176ays7mtRvBvL0f78NFIE,12909
+ebm4subjects/prepare_data.py,sha256=vQ-BdXkIP3iZJdPXol0WDlY8cRFMHkjzzL7oC7EbouE,3084
+ebm4subjects-0.5.5.dist-info/METADATA,sha256=oekB-uWB3p53odPkbtx-CqzxL_AHf6Az3RJcNhw1xhY,8354
+ebm4subjects-0.5.5.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+ebm4subjects-0.5.5.dist-info/licenses/LICENSE,sha256=RpvAZSjULHvoTR_esTlucJ08-zdQydnoqQLbqOh9Ub8,13826
+ebm4subjects-0.5.5.dist-info/RECORD,,

ebm4subjects-0.5.3.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=oVLNQv7IVb7KhhExb8o38z1xS3na_DzL-uoIK2A7IW0,30269
-ebm4subjects/embedding_generator.py,sha256=VXnZ2mqu2emmyIUkW-pw-7I_Zikc2LqsyiGcg2sxMuc,6703
-ebm4subjects/prepare_data.py,sha256=vQ-BdXkIP3iZJdPXol0WDlY8cRFMHkjzzL7oC7EbouE,3084
-ebm4subjects-0.5.3.dist-info/METADATA,sha256=uIuPMpcd4GH4sCCn5mbTPUGkjodQBBoTD0cmBt64_9Q,8274
-ebm4subjects-0.5.3.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-ebm4subjects-0.5.3.dist-info/licenses/LICENSE,sha256=RpvAZSjULHvoTR_esTlucJ08-zdQydnoqQLbqOh9Ub8,13826
-ebm4subjects-0.5.3.dist-info/RECORD,,