hjxdl 0.0.12__py3-none-any.whl → 0.0.14__py3-none-any.whl

hdl/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.0.12'
-__version_tuple__ = version_tuple = (0, 0, 12)
+__version__ = version = '0.0.14'
+__version_tuple__ = version_tuple = (0, 0, 14)

hdl/utils/llm/embs.py

@@ -0,0 +1,71 @@
+class FlagEmbedder():
+    def __init__(
+        self,
+        emb_name: str = "bge",
+        emb_dir: str = None
+    ) -> None:
+        """Initializes the object with the specified embedding name and directory.
+        
+        Args:
+            emb_name (str): The name of the embedding. Defaults to "bge".
+            emb_dir (str): The directory path for the embedding model.
+        
+        Returns:
+            None
+        """
+        self.emb_name = emb_name
+        self.emb_dir = emb_dir
+
+        if "bge" in emb_name.lower():
+            from FlagEmbedding import BGEM3FlagModel
+            self.model = BGEM3FlagModel(
+                emb_dir,  
+                use_fp16=True
+            )
+        elif "bce" in emb_name.lower():
+            from BCEmbedding import EmbeddingModel
+            self.model = EmbeddingModel(
+                model_name_or_path=emb_dir,
+                use_fp16=True
+            )
+    
+    def encode(
+        self,
+        sentences
+    ):
+        """Encode the input sentences using the model.
+        
+            Args:
+                sentences (list): List of sentences to encode.
+        
+            Returns:
+                numpy.ndarray: Encoded representation of the input sentences.
+        """
+        output = self.model.encode(
+            sentences,
+            return_dense=True,
+            return_sparse=True,
+            return_colbert_vecs=False
+        )
+        if "bge" in self.emb_name.lower():
+            return output["dense_vecs"]
+        return output
+    
+    def sim(
+        self,
+        sentences_1,
+        sentences_2
+    ):
+        """Calculate the similarity between two sets of sentences.
+        
+            Args:
+                sentences_1 (list): List of sentences for the first set.
+                sentences_2 (list): List of sentences for the second set.
+        
+            Returns:
+                float: Similarity score between the two sets of sentences.
+        """
+        output_1 = self.encode(sentences_1)
+        output_2 = self.encode(sentences_2)
+        similarity = output_1 @ output_2.T
+        return similarity.item()

{hjxdl-0.0.12.dist-info → hjxdl-0.0.14.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: hjxdl
-Version: 0.0.12
+Version: 0.0.14
 Summary: A collection of functions for Jupyter notebooks
 Home-page: https://github.com/huluxiaohuowa/hdl
 Author: Jianxing Hu

{hjxdl-0.0.12.dist-info → hjxdl-0.0.14.dist-info}/RECORD

@@ -1,5 +1,5 @@
 hdl/__init__.py,sha256=5sZZNySv08wwfzJcSDssGTqUn9wlmDsR6R4XB8J8mFM,70
-hdl/_version.py,sha256=J9AUd4fAuSymrNadeGqOirSMYNZ9WQmqXwLZX1sdsUk,413
+hdl/_version.py,sha256=Z5KVh4ecKbPlzksuEdQrX-Mm8ihZCLuudKr2xk8fqgc,413
 hdl/args/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 hdl/args/loss_args.py,sha256=s7YzSdd7IjD24rZvvOrxLLFqMZQb9YylxKeyelSdrTk,70
 hdl/controllers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -85,9 +85,10 @@ hdl/utils/general/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU
 hdl/utils/general/glob.py,sha256=8-RCnt6L297wMIfn34ZAMCsGCZUjHG3MGglGZI1cX0g,491
 hdl/utils/llm/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 hdl/utils/llm/chat.py,sha256=NEbWdrPpQqS5ArU7ghA6xuam8wwvWXk3LL_CiNEEuIg,7055
+hdl/utils/llm/embs.py,sha256=rkMDaQIVX7cmTEtOhTBzqZYwouOzO3ox0DpJ493BXrs,2124
 hdl/utils/schedulers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 hdl/utils/schedulers/norm_lr.py,sha256=bDwCmdEK-WkgxQMFBiMuchv8Mm7C0-GZJ6usm-PQk14,4461
-hjxdl-0.0.12.dist-info/METADATA,sha256=a06qubt7_L4HGCpF0svsd1lGd4i6I0jurdc7fWV1wJg,543
-hjxdl-0.0.12.dist-info/WHEEL,sha256=Z4pYXqR_rTB7OWNDYFOm1qRk0RX6GFP2o8LgvP453Hk,91
-hjxdl-0.0.12.dist-info/top_level.txt,sha256=-kxwTM5JPhylp06z3zAVO3w6_h7wtBfBo2zgM6YZoTk,4
-hjxdl-0.0.12.dist-info/RECORD,,
+hjxdl-0.0.14.dist-info/METADATA,sha256=0mn1FFKnA5m5jqrIqlh-XioiqkonZ3DXD6I770Q62Xc,543
+hjxdl-0.0.14.dist-info/WHEEL,sha256=Z4pYXqR_rTB7OWNDYFOm1qRk0RX6GFP2o8LgvP453Hk,91
+hjxdl-0.0.14.dist-info/top_level.txt,sha256=-kxwTM5JPhylp06z3zAVO3w6_h7wtBfBo2zgM6YZoTk,4
+hjxdl-0.0.14.dist-info/RECORD,,