from blurr.data.all import *
from blurr.modeling.all import *
from huggingface_hub import hf_hub_download
from fastai.learner import load_learner
from transformers import *
import gradio as gr

def download():
    pretrained_model_name = "Helsinki-NLP/opus-mt-en-ml"
    model_cls = AutoModelForSeq2SeqLM
    hf_arch, hf_config, hf_tokenizer, hf_model = BLURR.get_hf_objects(pretrained_model_name, model_cls=model_cls)

    model = load_learner(hf_hub_download("kurianbenoy/kde_en_ml_translation_model", "saved_model.pkl"))

def translate_text(text):
    download()
    outputs = model.blurr_translate(text)
    return outputs[0]["translation_texts"]
    

iface = gr.Interface(
    fn=translate_text,
    inputs=gr.inputs.Textbox(lines=2, label="Enter Text in English", placeholder="Type text in English Here..."),
    outputs="text",
    theme="huggingface",
)

iface.launch()