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

gr.Interface.load("huggingface/kurianbenoy/kde_en_ml_translation_model")

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

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

def translate_text(text):
    model = 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",
    examples=["Add All Found Feeds to Akregator.", "Hello world",]
)

iface.launch(
    share=False,
    cache_examples=True,
    enable_queue=True,
    )