import gradio as gr
import torch
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
import os
import sys

# Add local IndicTransToolkit path
sys.path.append(os.path.abspath("libs/IndicTransToolkit"))
from IndicTransToolkit.processor import IndicProcessor

# Load processor and model
ip = IndicProcessor(inference=True)
tokenizer = AutoTokenizer.from_pretrained("ai4bharat/indictrans2-en-indic-dist-200M", trust_remote_code=True)
model = AutoModelForSeq2SeqLM.from_pretrained("ai4bharat/indictrans2-en-indic-dist-200M", trust_remote_code=True)

LANG_OPTIONS = [
    "hin_Deva", "ben_Beng", "pan_Guru", "guj_Gujr", 
    "tam_Taml", "tel_Telu", "mal_Mlym", 
    "mar_Deva", "kan_Knda", "asm_Beng"
]

def translate(text, target_lang):
    if not text.strip():
        return "Please enter some text."
    
    try:
        batch = ip.preprocess_batch([text], src_lang="eng_Latn", tgt_lang=target_lang)
        batch = tokenizer(batch, padding="longest", truncation=True, max_length=256, return_tensors="pt")

        with torch.inference_mode():
            outputs = model.generate(**batch, num_beams=5, max_length=256)

        with tokenizer.as_target_tokenizer():
            decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True, clean_up_tokenization_spaces=True)

        translated = ip.postprocess_batch(decoded, lang=target_lang)[0]
        return translated
    except Exception as e:
        return f"Error: {e}"

demo = gr.Interface(
    fn=translate,
    inputs=[
        gr.Textbox(label="Enter text in English", lines=5),
        gr.Dropdown(choices=LANG_OPTIONS, label="Select Target Language")
    ],
    outputs="text",
    title="IndicTrans Translator",
    description="Translate English text into Indian languages using IndicTrans2."
)

if __name__ == "__main__":
    demo.launch()