Update main.py

main.py

@@ -10,19 +10,23 @@ from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import HTMLResponse
 import uvicorn
 
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel
 from pymongo import MongoClient
 import jwt
 from jwt import encode as jwt_encode
-from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from bson import ObjectId
 
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 import ctranslate2
 import sentencepiece as spm
 import fasttext
+import torch
 
-import pytz
 from datetime import datetime
+import gradio as gr
+import pytz
+import time
 import os
 
 app = FastAPI()
@@ -47,9 +51,9 @@ templates_folder = os.path.join(os.path.dirname(__file__), "templates")
 
 # Authentication
 class User(BaseModel):
- username: str = None # Make the username field optional
- email: str
- password: str
+    username: str = None # Make the username field optional
+    email: str
+    password: str
 
 # Connect to the MongoDB database
 client = MongoClient("mongodb://localhost:27017")
@@ -64,63 +68,63 @@ security = HTTPBearer()
 
 @app.post("/login")
 def login(user: User):
- # Check if user exists in the database
- user_data = users_collection.find_one(
- {"email": user.email, "password": user.password}
- )
- if user_data:
-     # Generate a token
-     token = generate_token(user.email)
-     # Convert ObjectId to string
-     user_data["_id"] = str(user_data["_id"])
-     # Store user details and token in local storage
-     user_data["token"] = token
-     return user_data
- return {"message": "Invalid email or password"}
+     # Check if user exists in the database
+     user_data = users_collection.find_one(
+     {"email": user.email, "password": user.password}
+     )
+     if user_data:
+         # Generate a token
+         token = generate_token(user.email)
+         # Convert ObjectId to string
+         user_data["_id"] = str(user_data["_id"])
+         # Store user details and token in local storage
+         user_data["token"] = token
+         return user_data
+     return {"message": "Invalid email or password"}
 
 #Implement the registration route:
 @app.post("/register")
 def register(user: User):
- # Check if user already exists in the database
- existing_user = users_collection.find_one({"email": user.email})
- if existing_user:
+     # Check if user already exists in the database
+     existing_user = users_collection.find_one({"email": user.email})
+     if existing_user:
      return {"message": "User already exists"}
- #Insert the new user into the database
- user_dict = user.dict()
- users_collection.insert_one(user_dict)
- # Generate a token
- token = generate_token(user.email)
- # Convert ObjectId to string
- user_dict["_id"] = str(user_dict["_id"])
- # Store user details and token in local storage
- user_dict["token"] = token
- return user_dict
+     #Insert the new user into the database
+     user_dict = user.dict()
+     users_collection.insert_one(user_dict)
+     # Generate a token
+     token = generate_token(user.email)
+     # Convert ObjectId to string
+     user_dict["_id"] = str(user_dict["_id"])
+     # Store user details and token in local storage
+     user_dict["token"] = token
+     return user_dict
 
 
 #Implement the `/api/user` route to fetch user data based on the JWT token
 @app.get("/api/user")
 def get_user(credentials: HTTPAuthorizationCredentials = Depends(security)):
- # Extract the token from the Authorization header
- token = credentials.credentials
- # Authenticate and retrieve the user data from the database based on the token
- # Here, you would implement the authentication logic and fetch user details
- # based on the token from the database or any other authentication mechanism
- # For demonstration purposes, assuming the user data is stored in local storage
- # Note: Local storage is not accessible from server-side code
- # This is just a placeholder to demonstrate the concept
- user_data = {
- "username": "John Doe",
- "email": "johndoe@example.com"
- }
-if user_data["username"] and user_data["email"]:
- return user_data
-raise HTTPException(status_code=401, detail="Invalid token")
+     # Extract the token from the Authorization header
+     token = credentials.credentials
+     # Authenticate and retrieve the user data from the database based on the token
+     # Here, you would implement the authentication logic and fetch user details
+     # based on the token from the database or any other authentication mechanism
+     # For demonstration purposes, assuming the user data is stored in local storage
+     # Note: Local storage is not accessible from server-side code
+     # This is just a placeholder to demonstrate the concept
+     user_data = {
+     "username": "John Doe",
+     "email": "johndoe@example.com"
+     }
+     if user_data["username"] and user_data["email"]:
+        return user_data
+     raise HTTPException(status_code=401, detail="Invalid token")
 
 #Define a helper function to generate a JWT token
 def generate_token(email: str) -> str:
- payload = {"email": email}
- token = jwt_encode(payload, SECRET_KEY, algorithm="HS256")
- return token
+     payload = {"email": email}
+     token = jwt_encode(payload, SECRET_KEY, algorithm="HS256")
+     return token
 
     
 # Get time of request
@@ -135,6 +139,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
@@ -154,13 +181,17 @@ sp_model_full_path = os.path.join(os.path.dirname(__file__), sp_model_file)
 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')
+print('\nDone importing models\n') 
 
     
 def translate_detect(userinput: str, target_lang: str):
@@ -213,6 +244,25 @@ 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):
@@ -258,5 +308,23 @@ async def translate_enter_endpoint(request: Request):
         "translated_text": translated_text_e,
     }
 
+@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.")
 
-print("\nAPI starting .......\n")
+    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")