hjxdl 0.0.13__py3-none-any.whl → 0.0.15__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.13'
-__version_tuple__ = version_tuple = (0, 0, 13)
+__version__ = version = '0.0.15'
+__version_tuple__ = version_tuple = (0, 0, 15)

hdl/utils/llm/embs.py

@@ -1,23 +1,33 @@
-from FlagEmbedding import BGEM3FlagModel
-
-
-class FlagEmbedder():
+class HFEmbedder():
     def __init__(
         self,
+        emb_name: str = "bge",
         emb_dir: str = None
     ) -> None:
-        """Initialize the class with the specified embedding directory.
+        """Initializes the object with the specified embedding name and directory.
         
         Args:
-            emb_dir (str): The directory containing the embeddings.
-            
+            emb_name (str): The name of the embedding. Defaults to "bge".
+            emb_dir (str): The directory path for the embedding model.
+        
         Returns:
             None
         """
-        self.model = BGEM3FlagModel(
-            emb_dir,  
-            use_fp16=True
-        )
+        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,
@@ -31,12 +41,16 @@ class FlagEmbedder():
             Returns:
                 numpy.ndarray: Encoded representation of the input sentences.
         """
+        if isinstance(sentences, str):
+            sentences = [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(
@@ -55,5 +69,5 @@ class FlagEmbedder():
         """
         output_1 = self.encode(sentences_1)
         output_2 = self.encode(sentences_2)
-        similarity = output_1['dense_vecs'] @ output_2['dense_vecs'].T
-        return similarity.item()
+        similarity = output_1 @ output_2.T
+        return similarity

{hjxdl-0.0.13.dist-info → hjxdl-0.0.15.dist-info}/METADATA

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

{hjxdl-0.0.13.dist-info → hjxdl-0.0.15.dist-info}/RECORD

@@ -1,5 +1,5 @@
 hdl/__init__.py,sha256=5sZZNySv08wwfzJcSDssGTqUn9wlmDsR6R4XB8J8mFM,70
-hdl/_version.py,sha256=FMLMzW4kZnBuc3z5uydtuxnwzYOv6BBwKN68kOrxpEo,413
+hdl/_version.py,sha256=pzShTmESM-5FMm5QD9BQxtLcs8dS3tIY17WJZlF9dSQ,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,10 +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=TB3SnD-DZ_d8uV-UT63Q_8SECFDAOzsWev4SbPj3-ug,1605
+hdl/utils/llm/embs.py,sha256=yCFtc25gUFas6kwgOGBFydeaHNyQMq5y1Chxl8TNEUQ,2190
 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.13.dist-info/METADATA,sha256=vg17dj7a7mLuDjaMgh1lSNe8wIWqInI3Xjvnw16lNiI,543
-hjxdl-0.0.13.dist-info/WHEEL,sha256=Z4pYXqR_rTB7OWNDYFOm1qRk0RX6GFP2o8LgvP453Hk,91
-hjxdl-0.0.13.dist-info/top_level.txt,sha256=-kxwTM5JPhylp06z3zAVO3w6_h7wtBfBo2zgM6YZoTk,4
-hjxdl-0.0.13.dist-info/RECORD,,
+hjxdl-0.0.15.dist-info/METADATA,sha256=uVY_O0levGtYTHXWtCIstsHzulDmM4JLwQ7Xu6uZ5sg,543
+hjxdl-0.0.15.dist-info/WHEEL,sha256=Z4pYXqR_rTB7OWNDYFOm1qRk0RX6GFP2o8LgvP453Hk,91
+hjxdl-0.0.15.dist-info/top_level.txt,sha256=-kxwTM5JPhylp06z3zAVO3w6_h7wtBfBo2zgM6YZoTk,4
+hjxdl-0.0.15.dist-info/RECORD,,