ebm4subjects 0.5.2__tar.gz → 0.5.4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ebm4subjects
-Version: 0.5.2
+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.2 → ebm4subjects-0.5.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ebm4subjects"
-version = "0.5.2"
+version = "0.5.4"
 description = "Embedding Based Matching for Automated Subject Indexing"
 authors = [
     {name = "Deutsche Nationalbibliothek"},

{ebm4subjects-0.5.2 → ebm4subjects-0.5.4}/src/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,
-    EmbeddingGeneratorInternal,
     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 = "internal",
+        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 == "internal":
-                self.logger.info("initializing internal embedding generator")
-                self.generator = EmbeddingGeneratorInternal(
+            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,
                 )
@@ -404,7 +413,7 @@ class EbmModel:
             .join(
                 other=gold_standard.with_columns(pl.lit(True).alias("gold")),
                 on=["doc_id", "label_id"],
-                how="outer",
+                how="full",
             )
             # Fill dataframe so that all not suggested labels which are not part of
             # the gold standard and all gold standard labels which where not

{ebm4subjects-0.5.2 → ebm4subjects-0.5.4}/src/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,26 +87,91 @@ 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)
 
 
-class EmbeddingGeneratorInternal(EmbeddingGenerator):
+class EmbeddingGeneratorOfflineInference(EmbeddingGenerator):
     """
-    A class for generating embeddings using a given SentenceTransformer model.
+    A class for generating embeddings using a given SentenceTransformer model
+    loaded offline with SentenceTransformer.
 
     Args:
         model_name (str): The name of the SentenceTransformer model.
@@ -118,7 +185,7 @@ class EmbeddingGeneratorInternal(EmbeddingGenerator):
 
     def __init__(self, model_name: str, embedding_dimensions: int, **kwargs) -> None:
         """
-        Initializes the internal EmbeddingGenerator.
+        Initializes the EmbeddingGenerator in offline inference mode.
 
         Sets the model name, embedding dimensions, and creates a
         SentenceTransformer model instance.

{ebm4subjects-0.5.2 → ebm4subjects-0.5.4}/tests/test_prepare_data.py

@@ -1,8 +1,10 @@
-import polars as pl
 from pathlib import Path
 
+import polars as pl
+
 from ebm4subjects.prepare_data import parse_vocab
 
+
 def test_parse_vocab_reads_ttl_and_returns_dataframe(tmp_path):
     # Copy the sample vocab.ttl to a temp location
     vocab_src = Path(__file__).parent / "data/vocab.ttl"