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

ebm4subjects/ebm_model.py

@@ -19,7 +19,7 @@ from ebm4subjects.embedding_generator import EmbeddingGenerator
 class EbmModel:
     def __init__(
         self,
-        embedding_model_name: str,
+        embedding_model_name: str | Any,
         embedding_dimensions: int | str,
         chunk_tokenizer: str | Any,
         max_chunk_count: int | str,

ebm4subjects/embedding_generator.py

@@ -9,7 +9,8 @@ class EmbeddingGenerator:
     A class for generating embeddings using a given SentenceTransformer model.
 
     Args:
-        model_name (str): The name of the SentenceTransformer model to use.
+        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.
 
@@ -19,7 +20,9 @@ class EmbeddingGenerator:
         model (SentenceTransformer): The SentenceTransformer model instance.
     """
 
-    def __init__(self, model_name: str, embedding_dimensions: int, **kwargs) -> None:
+    def __init__(
+        self, model_name: str | SentenceTransformer, embedding_dimensions: int, **kwargs
+    ) -> None:
         """
         Initializes the EmbeddingGenerator.
 
@@ -31,9 +34,13 @@ class EmbeddingGenerator:
 
         # Create a SentenceTransformer model instance with the given
         # model name and embedding dimensions
-        self.model = SentenceTransformer(
-            model_name, truncate_dim=embedding_dimensions, **kwargs
-        )
+        # 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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ebm4subjects
-Version: 0.5.0
+Version: 0.5.1
 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.0.dist-info → ebm4subjects-0.5.1.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=fNO8Hy4EAY8gwu3ysSxQC_Kp2pUecbdr-6Ue7pFoAkk,29119
-ebm4subjects/embedding_generator.py,sha256=jC4rz4W50tKndxYezD7Kaoqysl8zhN-TbWirxA_WIQc,2354
+ebm4subjects/ebm_model.py,sha256=PVFtljF3oZK8u0lA6df82lsTdAD8H1Y9CHvWq1jWF2M,29125
+ebm4subjects/embedding_generator.py,sha256=DZhZxkjcsy_4NA62_2V-4UPbIUkg5qMPat_cIgsoIAA,2609
 ebm4subjects/prepare_data.py,sha256=vQ-BdXkIP3iZJdPXol0WDlY8cRFMHkjzzL7oC7EbouE,3084
-ebm4subjects-0.5.0.dist-info/METADATA,sha256=8JA0eDC-5-M7FW9nq3hqJDIpSGes4KoKd9uUWgafq8k,8274
-ebm4subjects-0.5.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-ebm4subjects-0.5.0.dist-info/licenses/LICENSE,sha256=RpvAZSjULHvoTR_esTlucJ08-zdQydnoqQLbqOh9Ub8,13826
-ebm4subjects-0.5.0.dist-info/RECORD,,
+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,,