import gradio as gr
import ctranslate2
import transformers
from huggingface_hub import snapshot_download

MODEL_REPO = "avans06/madlad400-7b-mt-bt-ct2-int8_float16"
TOKENIZER_NAME = "jbochi/madlad400-7b-mt-bt"

MODEL_DIR = snapshot_download(MODEL_REPO)
translator = ctranslate2.Translator(MODEL_DIR, compute_type="int8")
tokenizer = transformers.AutoTokenizer.from_pretrained(TOKENIZER_NAME)

LANG_CODES = {
    "Vietnamese (vi)": "<2vi>",
    "French (fr)": "<2fr>",
    "Chinese (zh)": "<2zh>",
    "Japanese (ja)": "<2ja>",
    "Korean (ko)": "<2ko>",
}

def translate(text, target_language):
    prefix = LANG_CODES[target_language]
    input_tokens = tokenizer.convert_ids_to_tokens(tokenizer.encode(prefix + " " + text))
    results = translator.translate_batch([input_tokens], beam_size=4)
    output_tokens = results[0].hypotheses[0]
    return tokenizer.decode(tokenizer.convert_tokens_to_ids(output_tokens))

with gr.Blocks(title="MADLAD-400-7B-BT Translate") as demo:
    gr.Markdown("## MADLAD-400-7B-BT — CTranslate2 int8 (CPU)")
    with gr.Row():
        input_text = gr.Textbox(lines=6, label="Input (English)")
        output_text = gr.Textbox(lines=6, label="Translation", interactive=False)
    target_lang = gr.Dropdown(choices=list(LANG_CODES.keys()), value="Vietnamese (vi)", label="Target language")
    btn = gr.Button("Translate", variant="primary")
    btn.click(fn=translate, inputs=[input_text, target_lang], outputs=output_text)

demo.launch()