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

Author: liwenshipro

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")
+text_embeddings_inference = LazyImport("text_embeddings_inference", globals(), "modelcache.embedding.text_embeddings_inference")
 
 
 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 TextEmbeddingsInference(base_url, model):
+    return text_embeddings_inference.TextEmbeddingsInference(base_url, model)

modelcache/embedding/text_embeddings_inference.py

@@ -0,0 +1,31 @@
+# -*- coding: utf-8 -*-
+import requests
+import numpy as np
+from modelcache.embedding.base import BaseEmbedding
+
+class TextEmbeddingsInference(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