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

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,
+    EmbeddingGeneratorOfflineInference,
+    EmbeddingGeneratorOpenAI,
 )
 
 
@@ -43,7 +44,7 @@ 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,
@@ -99,7 +100,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:
@@ -179,19 +180,27 @@ class EbmModel:
             None
         """
         if self.generator is None:
-            if self.embedding_model_type == "offline-inference":
+            if self.embedding_model_deployment == "offline-inference":
                 self.logger.info("initializing offline-inference embedding generator")
                 self.generator = EmbeddingGeneratorOfflineInference(
                     model_name=self.embedding_model_name,
                     embedding_dimensions=self.embedding_dimensions,
                     **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,
+                    **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,
                     **self.embedding_model_args,
                 )

ebm4subjects/embedding_generator.py

@@ -3,6 +3,7 @@ import os
 import numpy as np
 import requests
 from sentence_transformers import SentenceTransformer
+from tqdm import tqdm
 
 
 class EmbeddingGenerator:
@@ -41,6 +42,7 @@ class EmbeddingGeneratorAPI(EmbeddingGenerator):
 
     def __init__(
         self,
+        model_name: str,
         embedding_dimensions: int,
         **kwargs,
     ) -> None:
@@ -52,7 +54,7 @@ class EmbeddingGeneratorAPI(EmbeddingGenerator):
         """
 
         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"})
@@ -85,19 +87,83 @@ 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:
+                # TODO: write warning to logger
+                for _ in batch_texts:
+                    # TODO: ensure same format as true case and truncate dim
+                    embeddings.append([0 for _ in range(self.embedding_dimensions)])
+
+        return np.array(embeddings)
+
+
+class EmbeddingGeneratorOpenAI(EmbeddingGeneratorAPI):
+    """
+    A class for generating embeddings using any OpenAI compatibleAPI.
+    """
+
+    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
+                SentenceTransformer model.
+
+        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]
+            data = {
+                "input": batch_texts,
+                "model": self.model_name,
+                "encoding_format": "float",
+                **kwargs,
+            }
+
+            response = self.session.post(
+                self.api_address, headers=self.headers, json=data
+            )
+
+            # Process all embeddings from the batch response
+            if response.status_code == 200:
+                response_data = response.json()
+                for i, _ in enumerate(batch_texts):
+                    embedding = response_data["data"][i]["embedding"]
+                    embeddings.append(embedding)
             else:
-                embeddings.append([0 for _ in range(self.embedding_dimensions)])
+                # TODO: write warning to logger
+                for _ in batch_texts:
+                    embeddings.append([0 for _ in range(self.embedding_dimensions)])
 
         return np.array(embeddings)
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ebm4subjects
-Version: 0.5.3
+Version: 0.5.4
 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>

{ebm4subjects-0.5.3.dist-info → ebm4subjects-0.5.4.dist-info}/RECORD

@@ -3,10 +3,10 @@ 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/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.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,,
+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,,