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",  # Hindi
    "ben_Beng",  # Bengali
    "pan_Guru",  # Punjabi
    "guj_Gujr",  # Gujarati
    "tam_Taml",  # Tamil
    "tel_Telu",  # Telugu
    "mal_Mlym",  # Malayalam
    "mar_Deva",  # Marathi
    "kan_Knda",  # Kannada
    "asm_Beng",  # Assamese
    "kas_Arab",  # Kashmiri (Arabic)
    "kas_Deva",  # Kashmiri (Devanagari)
    "san_Deva",  # Sanskrit
    "brx_Deva",  # Bodo
    "mai_Deva",  # Maithili
    "sat_Olck",  # Santali
    "eng_Latn",  # English
    "urd_Arab"   # Urdu

]

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()