Update apis/chat_api.py

apis/chat_api.py

@@ -1,7 +1,14 @@
 import argparse
 import uvicorn
 import sys
+import os
+import io
+from transformers import M2M100Tokenizer, M2M100ForConditionalGeneration
+import time
 import json
+from typing import List
+import torch
+import logging
 import string
 import random
 import base64
@@ -31,12 +38,12 @@ class ChatAPIApp:
         )
         self.setup_routes()
 
-    def get_available_models(self):
+    def get_available_langs(self):
         f = open('apis/lang_name.json', "r")
         self.available_models = json.loads(f.read())
         return self.available_models
 
-    class ChatCompletionsPostItem(BaseModel):
+    class TranslateCompletionsPostItem(BaseModel):
         from_language: str = Field(
             default="auto",
             description="(str) `Detect`",
@@ -51,7 +58,7 @@ class ChatAPIApp:
         )
    
 
-    def chat_completions(self, item: ChatCompletionsPostItem):
+    def translate_completions(self, item: TranslateCompletionsPostItem):
         translator = Translator()
         f = open('apis/lang_name.json', "r")
         available_langs = json.loads(f.read())
@@ -73,6 +80,60 @@ class ChatAPIApp:
         json_compatible_item_data = jsonable_encoder(item_response)
         return JSONResponse(content=json_compatible_item_data)
 
+    def translate_ai_completions(self, item: TranslateCompletionsPostItem):
+        translator = Translator()
+        f = open('apis/lang_name.json', "r")
+        available_langs = json.loads(f.read())
+        from_lang = 'en'
+        to_lang = 'en'
+        for lang_item in available_langs:
+          if item.to_language == lang_item['code']:
+              to_lang = item.to_language
+          if item.from_language == lang_item['code']:
+              from_lang = item.from_language
+
+        if to_lang == 'auto':
+            to_lang = 'en'
+
+        if from_lang == 'auto':
+            from_lang = translator.detect(item.input_text).lang
+            
+        if torch.cuda.is_available():
+            device = torch.device("cuda:0")
+        else:
+            device = torch.device("cpu")
+            logging.warning("GPU not found, using CPU, translation will be very slow.")
+
+        time_start = time.time()
+        
+        tokenizer = M2M100Tokenizer.from_pretrained(pretrained_model, cache_dir=cache_dir)
+        model = M2M100ForConditionalGeneration.from_pretrained(
+            "facebook/m2m100_1.2B", cache_dir="models/"
+        ).to(device)
+        model.eval()
+
+        tokenizer.src_lang = from_lang
+        with torch.no_grad():
+            encoded_input = tokenizer(item.input_text, return_tensors="pt").to(device)
+            generated_tokens = model.generate(
+               **encoded_input, forced_bos_token_id=tokenizer.get_lang_id(to_lang)
+            )
+            translated_text = tokenizer.batch_decode(
+            generated_tokens, skip_special_tokens=True
+            )[0]
+
+        time_end = time.time()
+        translated = translated_text
+        item_response = {
+            "from_language": from_lang,
+            "to_language": to_lang,
+            "text": item.input_text,
+            "translate": translated,
+            "start": str(time_start),
+            "end": str(time_end)
+        }
+        json_compatible_item_data = jsonable_encoder(item_response)
+        return JSONResponse(content=json_compatible_item_data)
         
     class DetectLanguagePostItem(BaseModel):
         input_text: str = Field(
@@ -125,15 +186,21 @@ class ChatAPIApp:
     def setup_routes(self):
         for prefix in ["", "/v1"]:
             self.app.get(
-                prefix + "/models",
+                prefix + "/langs",
                 summary="Get available languages",
-            )(self.get_available_models)
+            )(self.get_available_langs)
 
             self.app.post(
                 prefix + "/translate",
                 summary="translate text",
-            )(self.chat_completions)
+            )(self.translate_completions)
+
+            self.app.post(
+                prefix + "/translate/ai",
+                summary="translate text with ai",
+            )(self.translate_ai_completions)
 
+            
             self.app.post(
                 prefix + "/detect",
                 summary="detect language",