import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline

tokenizervin = AutoTokenizer.from_pretrained("vinai/vinai-translate-en2vi-v2", src_lang="en_XX")
modelvin = AutoModelForSeq2SeqLM.from_pretrained("vinai/vinai-translate-en2vi-v2")

def translate(en_text: str) -> str:
    input_ids = tokenizervin(en_text, return_tensors="pt").input_ids
    output_ids = modelvin.generate(
        input_ids,
        decoder_start_token_id=tokenizervin.lang_code_to_id["vi_VN"],
        num_return_sequences=1,
        max_length=4096,
        num_beams=5,
        early_stopping=True
    )
    vi_text = tokenizervin.batch_decode(output_ids, skip_special_tokens=True)
    vi_text = '\n'.join(vi_text)
    return vi_text

tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-1.3B", src_lang="eng_Latn")
model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-1.3B")

def translation(text):

    translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang="eng_Latn", tgt_lang="vie_Latn")
    output = translator(text, max_length=4096)

    output = output[0]['translation_text']
    return output

description_string = """
    <div align="center">
    <img src="https://i.spread.name/83a58f6a-475f-4b57-9c0f-b937a7919126_BIMBIP.png" width=200px>
    <span style="font-size: 24px; font-weight: bold;">BIPSOUP DRIVE-THRU</span><br>
    <span style="font-size: 18px; font-weight: bold;">You are almost ready to be entrusted with the secret ingredient of my Biptopia Soup!</span>
    </div>
    """

iface = gr.Interface(
    fn=translate,
    description=description_string,
    inputs=[gr.Textbox(lines=7, max_lines=100, placeholder="From BimBip with love", label="Ingredient"),],
    outputs=[gr.Textbox(lines=7, max_lines=100, placeholder="Nhanh, Gọn, Mướt", label="Fastsoup", show_label=True, show_copy_button=True),])

iface2 = gr.Interface(
    fn=translation,
    description=description_string,
    inputs=[gr.Textbox(lines=7, max_lines=100, placeholder="From BimBip with love", label="Ingredient"),],
    outputs=[gr.Textbox(lines=7, max_lines=100, placeholder="Nhanh, Gọn, Mướt", label="Fastsoup", show_label=True, show_copy_button=True),])

demo = gr.TabbedInterface([iface, iface2], ["VinAI", "NLLB"], theme=gr.themes.Soft())

if __name__ == "__main__":
    demo.launch()