import gradio as gr
import time
from config import model_repo_id, src_lang, tgt_lang
from indictrans2 import initialize_model_and_tokenizer, batch_translate
from examples import example_sentences


def load_models():
    model_dict = {}

    print("\tLoading model: %s" % model_repo_id)

    # build model and tokenizer
    en_indic_tokenizer, en_indic_model, en_indic_lora_model = (
        initialize_model_and_tokenizer()
    )

    model_dict["_tokenizer"] = en_indic_tokenizer
    model_dict["_model"] = en_indic_model
    model_dict["_lora_model"] = en_indic_lora_model

    return model_dict


def translation(text):

    start_time = time.time()

    tokenizer = model_dict["_tokenizer"]
    model = model_dict["_model"]
    lora_model = model_dict["_lora_model"]

    # org translation
    org_translation = batch_translate(
        [text],
        model=model,
        tokenizer=tokenizer,
    )
    org_output = org_translation[0]
    end_time = time.time()

    # lora translation
    lora_translation = batch_translate(
        [text],
        model=lora_model,
        tokenizer=tokenizer,
    )
    lora_output = lora_translation[0]
    end_time2 = time.time()

    result = {
        "source": src_lang,
        "target": tgt_lang,
        "input": text,
        "it2_result": org_output,
        "it2_conv_result": lora_output,
        "it2_inference_time": end_time - start_time,
        "it2_conv_inference_time": end_time2 - end_time,
    }

    return result


print("\tinit models")

global model_dict

model_dict = load_models()

inputs = gr.Textbox(lines=5, label="Input text")
outputs = gr.JSON(container=True)
submit_btn = gr.Button("Translate", variant="primary")

title = "IndicTrans2 fine-tuned on conversation"
description = f"Note: LoRA is trained only on En-Hi pair.\nDetails: https://github.com/AI4Bharat/IndicTrans2.\nLoRA Model: https://huggingface.co/sam749/IndicTrans2-Conv"

gr.Interface(
    fn=translation,
    inputs=inputs,
    outputs=outputs,
    title=title,
    description=description,
    submit_btn=submit_btn,
    examples=example_sentences,
    examples_per_page=10,
    cache_examples=False,
).launch(share=True)