import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# ✅ Use valid model
MODEL_NAME = "ai4bharat/indictrans2-indic-indic-1B"

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME, trust_remote_code=True)

def translate(text: str, src_lang: str, tgt_lang: str) -> str:
    """Translate text from src_lang to tgt_lang using IndicTrans2."""
    if not text.strip():
        return "⚠️ Please enter some text to translate."

    src_lang = src_lang.strip().lower()
    tgt_lang = tgt_lang.strip().lower()

    try:
        # Format input as required by IndicTrans2
        formatted_text = f"{src_lang}>>{tgt_lang} {text}"
        inputs = tokenizer(formatted_text, return_tensors="pt")

        # Generate translation
        output_tokens = model.generate(**inputs, max_length=512)
        translation = tokenizer.decode(output_tokens[0], skip_special_tokens=True)

        return translation
    except Exception as e:
        return f"❌ Error: {str(e)}"

# Gradio interface
demo = gr.Interface(
    fn=translate,
    inputs=[
        gr.Textbox(label="Text", placeholder="Enter your text here..."),
        gr.Textbox(label="Source Language Code (e.g., ta, hi, kn)", placeholder="ta"),
        gr.Textbox(label="Target Language Code (e.g., en, hi, kn)", placeholder="en")
    ],
    outputs=gr.Textbox(label="Translated Text"),
    title="IndicTrans2 Language Translator",
    description=(
        "🌐 Translate text between Indian languages using "
        "[ai4bharat/indictrans2-indic-indic-1B](https://huggingface.co/ai4bharat/indictrans2-indic-indic-1B)."
    )
)

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