Update apis/chat_api.py

apis/chat_api.py

@@ -4,6 +4,7 @@ import sys
 import os
 import io
 from transformers import M2M100Tokenizer, M2M100ForConditionalGeneration
+from transformers import pipeline
 import time
 import json
 from typing import List
@@ -137,7 +138,75 @@ class ChatAPIApp:
         }
         json_compatible_item_data = jsonable_encoder(item_response)
         return JSONResponse(content=json_compatible_item_data)
-        
+    class TranslateAiPostItem(BaseModel):
+        model: str = Field(
+            default="t5-base",
+            description="(str) `Model Name`",
+        )
+        from_language: str = Field(
+            default="en",
+            description="(str) `Detect`",
+        )
+        to_language: str = Field(
+            default="fa",
+            description="(str) `en`",
+        )
+        input_text: str = Field(
+            default="Hello",
+            description="(str) `Text for translate`",
+        )    
+    def ai_translate(self, item:TranslateAiPostItem):
+        MODEL_MAP = {
+        "t5-base": "t5-base",
+        "t5-small": "t5-small",
+        "t5-large": "t5-3b",
+        "t5-3b": "t5-3b",
+        "mbart-large-50-many-to-many-mmt": "facebook/mbart-large-50-many-to-many-mmt",
+        "nllb-200-distilled-600M": "facebook/nllb-200-distilled-600M",
+        "madlad400-3b-mt": "jbochi/madlad400-3b-mt",    
+        "default": "t5-base",
+        }
+        if item.model in MODEL_MAP.keys():
+            target_model = item.model
+        else:
+            target_model = "default"
+            
+        model_fullname = MODEL_MAP[target_model]   
+        request_url = (
+            f"https://api-inference.huggingface.co/models/{model_fullname}"
+        )
+        request_headers = {
+            "Content-Type": "application/json",
+        }
+        request_body = {
+            "inputs": item.input_text,
+            "parameters": {
+                "src_lang": item.from_language,
+                "tgt_lang": item.to_language,
+                "return_full_text": True,
+            },
+            "stream": False,
+        }
+        enver.set_envs(proxies=True)
+        stream = requests.post(
+            request_url,
+            headers=request_headers,
+            json=request_body,
+            proxies=enver.requests_proxies,
+            stream=False,
+        )
+        status_code = stream.status_code
+        if status_code == 200:
+            result = stream.json()
+        else:
+            result = 'unknown error'
+            
+        item_response = {
+            "statue": status_code,
+            "result": result,
+            }
+        json_compatible_item_data = jsonable_encoder(item_response)
+        return JSONResponse(content=json_compatible_item_data)
     class DetectLanguagePostItem(BaseModel):
         input_text: str = Field(
             default="Hello, how are you?",
@@ -201,7 +270,7 @@ class ChatAPIApp:
             self.app.post(
                 prefix + "/translate/ai",
                 summary="translate text with ai",
-            )(self.translate_ai_completions)
+            )(self.ai_translate)
 
             
             self.app.post(