import streamlit as st
from fairseq.models.transformer import TransformerModel
import time

@st.cache_resource
def load_model():
    # Load English-to-Indic model
    en2indic = TransformerModel.from_pretrained(
        'ai4bharat/indictrans2-en-indic-1B',
        checkpoint_file='model.pt',
        bpe='sentencepiece',
        sentencepiece_model='ai4bharat/indictrans2-en-indic-1B/en-indic.model',
        device='cpu'
    )
    
    # Load Indic-to-English model
    indic2en = TransformerModel.from_pretrained(
        'ai4bharat/indictrans2-indic-en-1B',
        checkpoint_file='model.pt',
        bpe='sentencepiece',
        sentencepiece_model='ai4bharat/indictrans2-indic-en-1B/indic-en.model',
        device='cpu'
    )
    
    return en2indic, indic2en

en2indic, indic2en = load_model()

st.title("IndicTrans2 Translator")
src_lang = st.selectbox("Source Language", ["English", "Hindi", "Tamil", "Telugu", "Malayalam"])
tgt_lang = st.selectbox("Target Language", ["English", "Hindi", "Tamil", "Telugu", "Malayalam"])
input_text = st.text_area("Enter text to translate")

if st.button("Translate"):
    start_time = time.time()
    
    # Language code mapping
    lang_codes = {
        "English": "eng_Latn",
        "Hindi": "hin_Deva",
        "Tamil": "tam_Taml",
        "Telugu": "tel_Telu",
        "Malayalam": "mal_Mlym"
    }
    
    if src_lang == "English":
        translation = en2indic.translate(input_text, lang_pair=(lang_codes[src_lang], lang_codes[tgt_lang]))
    else:
        translation = indic2en.translate(input_text, lang_pair=(lang_codes[src_lang], lang_codes[tgt_lang]))
    
    st.text_area("Translation", translation[0], height=200)
    st.write(f"Translation time: {time.time() - start_time:.2f} seconds")