import gradio as gr
import subprocess
subprocess.check_call(["pip", "install", "--upgrade", "huggingface-hub"])
subprocess.check_call(["pip", "install", "transformers"])
subprocess.check_call(["pip", "install", "torch"])
  

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

tokenizer = AutoTokenizer.from_pretrained("balaramas/mbart-sahitrans_new_data")
model = AutoModelForSeq2SeqLM.from_pretrained("balaramas/mbart-sahitrans_new_data")


def sanmt(txt):
    
    tokenizer.src_lang = "hi_IN"
    encoded_ar = tokenizer(txt, return_tensors="pt")   
    generated_tokens = model.generate(
        **encoded_ar,
        forced_bos_token_id=tokenizer.lang_code_to_id["hi_IN"]
    )
    output = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]   
    return output

iface = gr.Interface(
    fn=sanmt, 
    inputs=gr.Textbox(label="Enter text in Sanskrit", placeholder="Type here..."), 
    outputs=gr.Textbox(label="Translated Hindi Text"),
    title="Sanskrit to Hindi Translator"
)
iface.launch()