import streamlit as st
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
st.title("IndicTrans Translator")
st.write("Translate English text into Indian languages using IndicTrans2.")

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)

text = st.text_area("Enter text in English:", height=150)
target_lang = st.selectbox("Select Target Language", [
    "hin_Deva", "ben_Beng", "pan_Guru", "guj_Gujr", "tam_Taml", "tel_Telu", "mal_Mlym", "mar_Deva", "kan_Knda", "asm_Beng"
])

if st.button("Translate"):
    if not text.strip():
        st.warning("Please enter some text.")
    else:
        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]
            st.success(f"Translation: {translated}")
        except Exception as e:
            st.error(f"Error: {e}")