Merge pull request #62 from charleschile/embedding-generator

bge-m3 embedding generator

Author: peng3307165

modelcache/embedding/__init__.py

@@ -7,6 +7,7 @@
 paddlenlp = LazyImport("paddlenlp", globals(), "modelcache.embedding.paddlenlp")
 timm = LazyImport("timm", globals(), "modelcache.embedding.timm")
 huggingface_tei = LazyImport("huggingface_tei", globals(), "modelcache.embedding.huggingface_tei")
+bge_m3 = LazyImport("bge_m3", globals(), "modelcache.embedding.bge_m3")
 
 
 def Huggingface(model="sentence-transformers/all-mpnet-base-v2"):
@@ -33,4 +34,7 @@ def Timm(model="resnet50", device="default"):
     return timm.Timm(model, device)
 
 def HuggingfaceTEI(base_url, model):
-    return huggingface_tei.HuggingfaceTEI(base_url, model)
+    return huggingface_tei.HuggingfaceTEI(base_url, model)
+
+def BgeM3Embedding(model_path="model/bge-m3"):
+    return bge_m3.BgeM3Embedding(model_path)

modelcache/embedding/bge_m3.py

@@ -0,0 +1,28 @@
+# -*- coding: utf-8 -*-
+import numpy as np
+from modelcache.embedding.base import BaseEmbedding
+from transformers import AutoTokenizer, AutoModel
+from FlagEmbedding import BGEM3FlagModel
+
+class BgeM3Embedding(BaseEmbedding):
+    def __init__(self, model_path: str = "model/bge-m3"):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.model = AutoModel.from_pretrained(model_path)
+        
+        self.bge_model = BGEM3FlagModel(model_name_or_path=model_path, 
+                                        model=self.model, 
+                                        tokenizer=self.tokenizer, 
+                                        use_fp16=False)
+        
+        self.__dimension = 768 
+
+    def to_embeddings(self, data, **_):
+        if not isinstance(data, list):
+            data = [data]
+        
+        embeddings = self.bge_model.encode(data, batch_size=12, max_length=8192)['dense_vecs']
+        return np.array(embeddings).astype("float32")
+
+    @property
+    def dimension(self):
+        return self.__dimension