ebm4subjects 0.5.1__tar.gz → 0.5.2__tar.gz

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

@@ -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 → ebm4subjects-0.5.2}/pyproject.toml

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

{ebm4subjects-0.5.1 → ebm4subjects-0.5.2}/src/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-0.5.2/src/ebm4subjects/embedding_generator.py

@@ -0,0 +1,201 @@
+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): 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.
+    """
+
+    def __init__(self, model_name: str, embedding_dimensions: int, **kwargs) -> None:
+        """
+        Initializes the internal EmbeddingGenerator.
+
+        Sets the model name, embedding dimensions, and creates a
+        SentenceTransformer model instance.
+        """
+        self.model_name = model_name
+        self.embedding_dimensions = embedding_dimensions
+
+        # Create a SentenceTransformer model instance with the given
+        # model name and embedding dimensions
+        self.model = SentenceTransformer(
+            model_name, truncate_dim=embedding_dimensions, **kwargs
+        )
+
+        # Disabel parallelism for tokenizer
+        # Needed because process might be already parallelized
+        # before embedding creation
+        os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+    def generate_embeddings(self, texts: list[str], **kwargs) -> np.ndarray:
+        """
+        Generates embeddings for a list of input texts using the
+        SentenceTransformer model.
+
+        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.
+        """
+        # 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 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/src/ebm4subjects/embedding_generator.py

@@ -1,70 +0,0 @@
-import os
-
-import numpy as np
-from sentence_transformers import SentenceTransformer
-
-
-class 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.
-        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:
-        """
-        Initializes the EmbeddingGenerator.
-
-        Sets the model name, embedding dimensions, and creates a
-        SentenceTransformer model instance.
-        """
-        self.model_name = model_name
-        self.embedding_dimensions = embedding_dimensions
-
-        # 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
-
-        # Disabel parallelism for tokenizer
-        # Needed because process might be already parallelized
-        # before embedding creation
-        os.environ["TOKENIZERS_PARALLELISM"] = "false"
-
-    def generate_embeddings(self, texts: list[str], **kwargs) -> np.ndarray:
-        """
-        Generates embeddings for a list of input texts using the
-        SentenceTransformer model.
-
-        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.
-        """
-        # 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 embeddings using the SentenceTransformer model and return them
-        return self.model.encode(texts, **kwargs)