ebm4subjects 0.5.1__py3-none-any.whl → 0.5.2__py3-none-any.whl

ebm4subjects/ebm_model.py

@@ -13,13 +13,16 @@ from ebm4subjects import prepare_data
 from ebm4subjects.chunker import Chunker
 from ebm4subjects.duckdb_client import Duckdb_client
 from ebm4subjects.ebm_logging import EbmLogger, NullLogger, XGBLogging
-from ebm4subjects.embedding_generator import EmbeddingGenerator
+from ebm4subjects.embedding_generator import (
+    EmbeddingGeneratorHuggingFaceTEI,
+    EmbeddingGeneratorInternal,
+    EmbeddingGeneratorMock,
+)
 
 
 class EbmModel:
     def __init__(
         self,
-        embedding_model_name: str | Any,
         embedding_dimensions: int | str,
         chunk_tokenizer: str | Any,
         max_chunk_count: int | str,
@@ -39,7 +42,9 @@ class EbmModel:
         collection_name: str = "my_collection",
         use_altLabels: bool = True,
         hnsw_index_params: dict | str | None = None,
-        model_args: dict | str | None = None,
+        embedding_model_name: str | None = None,
+        embedding_model_type: str = "internal",
+        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,
@@ -94,11 +99,14 @@ class EbmModel:
 
         # Parameters for embedding generator
         self.generator = None
+        self.embedding_model_type = embedding_model_type
         self.embedding_model_name = embedding_model_name
         self.embedding_dimensions = int(embedding_dimensions)
-        if isinstance(model_args, str) or not model_args:
-            model_args = ast.literal_eval(model_args) if model_args else {}
-        self.model_args = model_args
+        if isinstance(embedding_model_args, str) or not embedding_model_args:
+            embedding_model_args = (
+                ast.literal_eval(embedding_model_args) if embedding_model_args else {}
+            )
+        self.embedding_model_args = embedding_model_args
         if isinstance(encode_args_vocab, str) or not encode_args_vocab:
             encode_args_vocab = (
                 ast.literal_eval(encode_args_vocab) if encode_args_vocab else {}
@@ -171,12 +179,25 @@ class EbmModel:
             None
         """
         if self.generator is None:
-            self.logger.info("initializing embedding generator")
-            self.generator = EmbeddingGenerator(
-                model_name=self.embedding_model_name,
-                embedding_dimensions=self.embedding_dimensions,
-                **self.model_args,
-            )
+            if self.embedding_model_type == "internal":
+                self.logger.info("initializing internal embedding generator")
+                self.generator = EmbeddingGeneratorInternal(
+                    model_name=self.embedding_model_name,
+                    embedding_dimensions=self.embedding_dimensions,
+                    **self.embedding_model_args,
+                )
+            elif self.embedding_model_type == "mock":
+                self.logger.info("initializing mock embedding generator")
+                self.generator = EmbeddingGeneratorMock(self.embedding_dimensions)
+            elif self.embedding_model_type == "HuggingFaceTEI":
+                self.logger.info("initializing API embedding generator")
+                self.generator = EmbeddingGeneratorHuggingFaceTEI(
+                    embedding_dimensions=self.embedding_dimensions,
+                    **self.embedding_model_args,
+                )
+            else:
+                self.logger.error("unsupportet API for embedding generator")
+                raise NotImplementedError
 
     def init_logger(
         self, log_path: str | None = None, logger: logging.Logger | None = None

ebm4subjects/embedding_generator.py

@@ -1,30 +1,124 @@
 import os
 
 import numpy as np
+import requests
 from sentence_transformers import SentenceTransformer
 
 
 class EmbeddingGenerator:
+    """
+    A base class for embedding generators.
+    """
+
+    def __init__(self) -> None:
+        """
+        Base method fot the initialization of an EmbeddingGenerator.
+        """
+        pass
+
+    def generate_embeddings(self, texts: list[str], **kwargs) -> np.ndarray:
+        """
+        Base method fot the creating embeddings with an EmbeddingGenerator.
+
+        Args:
+            texts (list[str]): A list of input texts.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            np.ndarray: A numpy array of shape (len(texts), embedding_dimensions)
+                containing the generated embeddings.
+        """
+        pass
+
+
+class EmbeddingGeneratorAPI(EmbeddingGenerator):
+    """
+    A base class for API embedding generators.
+
+    Attributes:
+        embedding_dimensions (int): The dimensionality of the generated embeddings.
+    """
+
+    def __init__(
+        self,
+        embedding_dimensions: int,
+        **kwargs,
+    ) -> None:
+        """
+        Initializes the API EmbeddingGenerator.
+
+        Sets the embedding dimensions, and initiliazes and
+        prepares a session with the API.
+        """
+
+        self.embedding_dimensions = embedding_dimensions
+
+        self.session = requests.Session()
+        self.api_address = kwargs.get("api_address")
+        self.headers = kwargs.get("headers", {"Content-Type": "application/json"})
+
+
+class EmbeddingGeneratorHuggingFaceTEI(EmbeddingGeneratorAPI):
+    """
+    A class for generating embeddings using the HuggingFaceTEI API.
+    """
+
+    def generate_embeddings(self, texts: list[str], **kwargs) -> np.ndarray:
+        """
+        Generates embeddings for a list of input texts using a model
+        via the HuggingFaceTEI 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 each text
+        for text in texts:
+            # send a request to the HuggingFaceTEI API
+            data = {"inputs": text}
+            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])
+            else:
+                embeddings.append([0 for _ in range(self.embedding_dimensions)])
+
+        return np.array(embeddings)
+
+
+class EmbeddingGeneratorInternal(EmbeddingGenerator):
     """
     A class for generating embeddings using a given SentenceTransformer model.
 
     Args:
-        model_name (str, SentenceTransformer): The name of the SentenceTransformer
-            model or an SentenceTransformer model to use.
+        model_name (str): The name of the SentenceTransformer model.
         embedding_dimensions (int): The dimensionality of the generated embeddings.
         **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.
-        model (SentenceTransformer): The SentenceTransformer model instance.
     """
 
-    def __init__(
-        self, model_name: str | SentenceTransformer, embedding_dimensions: int, **kwargs
-    ) -> None:
+    def __init__(self, model_name: str, embedding_dimensions: int, **kwargs) -> None:
         """
-        Initializes the EmbeddingGenerator.
+        Initializes the internal EmbeddingGenerator.
 
         Sets the model name, embedding dimensions, and creates a
         SentenceTransformer model instance.
@@ -34,13 +128,9 @@ class EmbeddingGenerator:
 
         # Create a SentenceTransformer model instance with the given
         # model name and embedding dimensions
-        # or set model to the given SentenceTransformer
-        if type(model_name) is str:
-            self.model = SentenceTransformer(
-                model_name, truncate_dim=embedding_dimensions, **kwargs
-            )
-        else:
-            self.model = model_name
+        self.model = SentenceTransformer(
+            model_name, truncate_dim=embedding_dimensions, **kwargs
+        )
 
         # Disabel parallelism for tokenizer
         # Needed because process might be already parallelized
@@ -68,3 +158,44 @@ class EmbeddingGenerator:
 
         # Generate embeddings using the SentenceTransformer model and return them
         return self.model.encode(texts, **kwargs)
+
+
+class EmbeddingGeneratorMock(EmbeddingGenerator):
+    """
+    A mock class for generating fake embeddings. Used for testing.
+
+    Args:
+        embedding_dimensions (int): The dimensionality of the generated embeddings.
+        **kwargs: Additional keyword arguments to pass to the model.
+
+    Attributes:
+        embedding_dimensions (int): The dimensionality of the generated embeddings.
+    """
+
+    def __init__(self, embedding_dimensions: int, **kwargs) -> None:
+        """
+        Initializes the mock EmbeddingGenerator.
+
+        Sets the embedding dimensions.
+        """
+        self.embedding_dimensions = embedding_dimensions
+
+    def generate_embeddings(self, texts: list[str], **kwargs) -> np.ndarray:
+        """
+        Generates embeddings for a list of input texts.
+
+        Args:
+            texts (list[str]): A list of input texts.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            np.ndarray: A numpy array of shape (len(texts), embedding_dimensions)
+                containing the generated 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))
+
+        # Generate mock embeddings return them
+        return np.ones((len(texts), 1024))

{ebm4subjects-0.5.1.dist-info → ebm4subjects-0.5.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ebm4subjects
-Version: 0.5.1
+Version: 0.5.2
 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.1.dist-info → ebm4subjects-0.5.2.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=PVFtljF3oZK8u0lA6df82lsTdAD8H1Y9CHvWq1jWF2M,29125
-ebm4subjects/embedding_generator.py,sha256=DZhZxkjcsy_4NA62_2V-4UPbIUkg5qMPat_cIgsoIAA,2609
+ebm4subjects/ebm_model.py,sha256=5LpZtYM3vZEyyJTTaQCsHzpsDLlJDdkvP8G61fPWrjs,30227
+ebm4subjects/embedding_generator.py,sha256=tPnb8XsNNNU4bYtjtc408eWI8KrZMyY8TI9d43AWHKo,6634
 ebm4subjects/prepare_data.py,sha256=vQ-BdXkIP3iZJdPXol0WDlY8cRFMHkjzzL7oC7EbouE,3084
-ebm4subjects-0.5.1.dist-info/METADATA,sha256=QkOBvOAI49_AUipc3yAH6RVG9OVUs_8jO64Bjfy561U,8274
-ebm4subjects-0.5.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-ebm4subjects-0.5.1.dist-info/licenses/LICENSE,sha256=RpvAZSjULHvoTR_esTlucJ08-zdQydnoqQLbqOh9Ub8,13826
-ebm4subjects-0.5.1.dist-info/RECORD,,
+ebm4subjects-0.5.2.dist-info/METADATA,sha256=pzoU6iJv3Cxh8q5656uUWac1t_QXuwuQUiGVfC7dklM,8274
+ebm4subjects-0.5.2.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+ebm4subjects-0.5.2.dist-info/licenses/LICENSE,sha256=RpvAZSjULHvoTR_esTlucJ08-zdQydnoqQLbqOh9Ub8,13826
+ebm4subjects-0.5.2.dist-info/RECORD,,

{ebm4subjects-0.5.1.dist-info → ebm4subjects-0.5.2.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: hatchling 1.27.0
+Generator: hatchling 1.28.0
 Root-Is-Purelib: true
 Tag: py3-none-any