Update apis/chat_api.py

apis/chat_api.py

@@ -29,26 +29,30 @@ from networks.openai_streamer import OpenaiStreamer
 from sentence_transformers import SentenceTransformer
 import tiktoken
 
-# class EmbeddingsAPIInference:
-#     def __init__(self, model_name):
-#         self.model_name=model_name
-#     def encode(self, x, headers=None, api_key=None):
-#         if api_key:
-#             if headers:
-#                 headers.update({"Authorization": f"Bearer {api_key}"})
-#             else:
-#                 headers = {"Authorization": f"Bearer {api_key}"}
+class EmbeddingsAPIInference:
+    def __init__(self, model_name):
+        self.model_name=model_name
+    def encode(self, x:str, api_key=None):
+        if api_key:
+            headers = {"Authorization": f"Bearer {api_key}"}
+        else:
+            headers = None
                 
-#         API_URL = "https://api-inference.huggingface.co/models/"+self.model_name
+        API_URL = "https://api-inference.huggingface.co/models/"+self.model_name
         
-#         payload = {
-#             "inputs": {
-#                 "source_sentence": "That is a happy person",
-#                 "sentences": ["That is a happy dog", "That is a very happy person", "Today is a sunny day"],
-#             }
-#         }
+        payload = {
+            "inputs": x,
+            "options":{"wait_for_model":True}
+        }
+        
+        return requests.post(API_URL, headers=headers, json=payload).json()
+
+
+class SentenceTransformer(SentenceTransformer):
+    def encode(self, **kwargs):
+        kwargs.pop("api_key", None)
+        return super().encode(**kwargs)
         
-#         requests.post(API_URL, headers=headers, json=payload).json()
         
         
 
@@ -65,6 +69,8 @@ class ChatAPIApp:
             "mxbai-embed-large":SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1"),
             "nomic-embed-text": SentenceTransformer("nomic-ai/nomic-embed-text-v1.5", trust_remote_code=True),
             "multilingual-e5-large-instruct":SentenceTransformer("intfloat/multilingual-e5-large-instruct"),
+            "intfloat/multilingual-e5-large-instruct":EmbeddingsAPIInference("intfloat/multilingual-e5-large-instruct"),
+            "mixedbread-ai/mxbai-embed-large-v1":EmbeddingsAPIInference("mixedbread-ai/mxbai-embed-large-v1")
         }
 
     def get_available_models(self):
@@ -350,7 +356,7 @@ class ChatAPIApp:
             model = request.model
             model_kwargs = request.options
             encoding = tiktoken.get_encoding("cl100k_base")
-            embeddings = self.embeddings[model].encode([encoding.decode(inp) for inp in request.input])#, **model_kwargs)
+            embeddings = self.embeddings[model].encode([encoding.decode(inp) for inp in request.input], api_key=api_key)#, **model_kwargs)
             return {
                 "object":"list", 
                 "data":[
@@ -366,7 +372,7 @@ class ChatAPIApp:
         try:
             model = request.model
             model_kwargs = request.options
-            embeddings = self.embeddings[model].encode(request.prompt)#, **model_kwargs)
+            embeddings = self.embeddings[model].encode(request.prompt, api_key=api_key)#, **model_kwargs)
             return {"embedding": embeddings.tolist()}
         except ValueError as e:
             raise HTTPException(status_code=400, detail=str(e))