import os
import gradio as gr
import torch
from huggingface_hub import login
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Optional login for private models/spaces
hf_token = os.environ.get("HF_TOKEN")
if hf_token:
    login(token=hf_token)

MODEL_NAME = "facebook/nllb-200-distilled-600M"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)

# Use GPU if available, otherwise CPU
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)


def translate(text, src_lang, tgt_lang):
    if not text or not text.strip():
        return ""

    tokenizer.src_lang = src_lang
    inputs = tokenizer(text, return_tensors="pt").to(device)

    forced_bos_token_id = tokenizer.convert_tokens_to_ids(tgt_lang)

    outputs = model.generate(
        **inputs,
        forced_bos_token_id=forced_bos_token_id,
        max_length=512,
    )

    translated = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return translated.strip()


demo = gr.Interface(
    fn=translate,
    inputs=[
        gr.Textbox(label="Input Text"),
        gr.Textbox(label="Source Language (eng_Latn, spa_Latn, vie_Latn)"),
        gr.Textbox(label="Target Language"),
    ],
    outputs=gr.Textbox(label="Translation"),
    title="NLLB Translator",
    description="English -> Target Language Translation using NLLB",
)

demo.queue()
demo.launch()