Update modules/app.py

modules/app.py

@@ -10,12 +10,15 @@ from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import HTMLResponse
 import uvicorn
 
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 import ctranslate2
 import sentencepiece as spm
 import fasttext
+import torch
 
-import pytz
 from datetime import datetime
+import pytz
+import time
 import os
 
 app = FastAPI()
@@ -30,12 +33,14 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+# set this key as an environment variable
+os.environ["HUGGINGFACEHUB_API_TOKEN"] = st.secrets['huggingface_token']
+
 fasttext.FastText.eprint = lambda x: None
 
 # User interface
 templates_folder = os.path.join(os.path.dirname(__file__), "templates")
-
-    
+  
 # Get time of request
 
 def get_time():
@@ -48,6 +53,29 @@ def get_time():
     
     full_date = f"{curr_day} | {curr_date} | {curr_time}"
     return full_date, curr_time
+
+
+def load_models():
+    # build model and tokenizer
+    model_name_dict = {
+            #'nllb-distilled-600M': 'facebook/nllb-200-distilled-600M',
+            #'nllb-1.3B': 'facebook/nllb-200-1.3B',
+            #'nllb-distilled-1.3B': 'facebook/nllb-200-distilled-1.3B',
+            #'nllb-3.3B': 'facebook/nllb-200-3.3B',
+            'nllb-moe-54b': 'facebook/nllb-moe-54b',
+            }
+
+    model_dict = {}
+
+    for call_name, real_name in model_name_dict.items():
+        print('\tLoading model: %s' % call_name)
+        model = AutoModelForSeq2SeqLM.from_pretrained(real_name)
+        tokenizer = AutoTokenizer.from_pretrained(real_name)
+        model_dict[call_name+'_model'] = model
+        model_dict[call_name+'_tokenizer'] = tokenizer
+
+    return model_dict 
+
     
 # Load the model and tokenizer ..... only once!
 beam_size = 1  # change to a smaller value for faster inference
@@ -68,10 +96,14 @@ sp = spm.SentencePieceProcessor()
 sp.load(sp_model_full_path)
 
 # Import The Translator model
+'''
 print("\nimporting Translator model")
 ct_model_file = "sematrans-3.3B"
 ct_model_full_path = os.path.join(os.path.dirname(__file__), ct_model_file)
 translator = ctranslate2.Translator(ct_model_full_path, device)
+'''
+print("\nimporting Translator model")
+model_dict = load_models()
 
 print('\nDone importing models\n')
 
@@ -126,6 +158,26 @@ def translate_enter(userinput: str, source_lang: str, target_lang: str):
   # Return the source language and the translated text
   return translations_desubword[0]
 
+
+def translate_faster(userinput3: str, source_lang3: str, target_lang3: str):
+    if len(model_dict) == 2:
+        model_name = 'nllb-moe-54b'
+        
+    start_time = time.time()
+    
+    model = model_dict[model_name + '_model']
+    tokenizer = model_dict[model_name + '_tokenizer']
+
+    translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=source_lang3, tgt_lang=target_lang3)
+    output = translator(userinput3, max_length=400)
+    end_time = time.time()
+
+    output = output[0]['translation_text']
+    result = {'inference_time': end_time - start_time,
+              'source': source,
+              'target': target,
+              'result': output}
+    return result
     
 @app.get("/", response_class=HTMLResponse)
 async def read_root(request: Request):
@@ -172,4 +224,23 @@ async def translate_enter_endpoint(request: Request):
     }
 
 
-print("\nAPI starting .......\n")
+@app.post("/translate_faster/")
+async def translate_faster_endpoint(request: Request):
+    dataf = await request.json()
+    userinputf = datae.get("userinput")
+    source_langf = datae.get("source_lang")
+    target_langf = datae.get("target_lang")
+    ffull_date = get_time()[0]
+    print(f"\nrequest: {ffull_date}\nSource_language; {source_langf}, Target Language; {target_langf}, User Input: {userinputf}\n")
+
+    if not userinputf or not target_langf:
+        raise HTTPException(status_code=422, detail="'userinput' 'sourc_lang'and 'target_lang' are required.")
+
+    translated_text_f = translate_faster(userinputf, source_langf, target_langf)
+    fcurrent_time = get_time()[1]
+    print(f"\nresponse: {fcurrent_time}; ... Translated Text: {translated_text_f}\n\n")
+    return {
+        "translated_text": translated_text_f,
+    }
+    
+print("\nAPI started successfully 😁\n")