import spaces
import gradio as gr
import torch

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from optimum.bettertransformer import BetterTransformer

tokenizer = AutoTokenizer.from_pretrained(
    "google/madlad400-3b-mt",
    use_fast=True
)

model_hf = AutoModelForSeq2SeqLM.from_pretrained(
    "google/madlad400-3b-mt",
    torch_dtype=torch.bfloat16
)

model = BetterTransformer.transform(model_hf, keep_original=True)

@spaces.GPU
def translate(text):
    """
    Translates the input text from English to Hawaiian.
    """
    text = "<2haw> " + text
    
    inputs = tokenizer(
        text,
        return_tensors="pt"
    )
    
    outputs = model.generate(**inputs, max_new_tokens=1000)
    text_translated = tokenizer.batch_decode(outputs, skip_special_tokens=True)
    
    return text_translated[0]

demo = gr.Interface(
    fn=translate,
    inputs=[gr.Textbox(label="English")],
    outputs=[gr.Textbox(label="Hawaiian")],
    title="MADLAD-400-3B-MT English-to-Hawaiian Translation",
    description="[Code](https://github.com/darylalim/madlad-400-3b-mt-eng-to-haw-translation)")

demo.queue()

demo.launch()