from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
from dotenv import load_dotenv
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import os


load_dotenv()


os.environ["HF_TOKEN"] = os.getenv('HF_TOKEN')


app = FastAPI()


origins = ["*"]


app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


class LanguageTextModel(BaseModel):
    languageText: str
    sourceLanguageCode: str
    targetLanguageCode: str


@app.get('/')
def welcome():
    return {
        'success': True,
        'message': 'server of "nllb language translator" is up and running successfully '
    }


@app.post('/translate')
async def translate_text(input: LanguageTextModel):
    
    model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")

    tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")

    translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=input.sourceLanguageCode, tgt_lang=input.targetLanguageCode, max_length=400)

    response = translator('I am not feeling well')

    return {
        "success": True,
        "translated_text": response[0]['translation_text']
    }