Merge pull request #39 from liwenshipro/main

feat: support huggingface/text-embeddings-inference for faster embedding inference

Author: peng3307165

examples/embedding/huggingface_tei_example.py

@@ -0,0 +1,17 @@
+# -*- coding: utf-8 -*-
+import sys
+sys.path.append(".")
+from modelcache.embedding.huggingface_tei import HuggingfaceTEI
+
+'''
+run tei server:
+text-embeddings-router --model-id BAAI/bge-large-zh-v1.5 --port 8080
+'''
+
+def run():
+    tei_instance = HuggingfaceTEI('http://127.0.0.1:8080/v1/embeddings', 'BAAI/bge-large-zh-v1.5')
+    print('dimenson', tei_instance.dimension)
+    print('embedding', tei_instance.to_embeddings('hello'))
+
+if __name__ == '__main__':
+    run()

modelcache/embedding/__init__.py

@@ -6,6 +6,7 @@
 fasttext = LazyImport("fasttext", globals(), "modelcache.embedding.fasttext")
 paddlenlp = LazyImport("paddlenlp", globals(), "modelcache.embedding.paddlenlp")
 timm = LazyImport("timm", globals(), "modelcache.embedding.timm")
+huggingface_tei = LazyImport("huggingface_tei", globals(), "modelcache.embedding.huggingface_tei")
 
 
 def Huggingface(model="sentence-transformers/all-mpnet-base-v2"):
@@ -30,3 +31,6 @@ def PaddleNLP(model="ernie-3.0-medium-zh"):
 
 def Timm(model="resnet50", device="default"):
     return timm.Timm(model, device)
+
+def HuggingfaceTEI(base_url, model):
+    return huggingface_tei.HuggingfaceTEI(base_url, model)

modelcache/embedding/huggingface_tei.py

@@ -0,0 +1,32 @@
+# -*- coding: utf-8 -*-
+import requests
+import numpy as np
+from modelcache.embedding.base import BaseEmbedding
+
+class HuggingfaceTEI(BaseEmbedding):
+    def __init__(self, base_url: str, model: str):
+        self.base_url = base_url
+        self.model = model
+        self.headers = {
+            'accept': 'application/json',
+            'Content-Type': 'application/json',
+        }
+        self.__dimension = self.to_embeddings('test').shape[0]
+
+    def to_embeddings(self, data, **_):
+        json_data = {
+            'input': data,
+            'model': self.model,
+        }
+
+        response = requests.post(self.base_url, headers=self.headers, json=json_data)
+        embedding = response.json()['data'][0]['embedding']
+        return np.array(embedding)
+
+    @property
+    def dimension(self):
+        """Embedding dimension.
+
+        :return: embedding dimension
+        """
+        return self.__dimension