#!/usr/bin/env python3
# -*- coding: utf-8 -*-

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

CURRENT_MODEL: str = None
TRANSLATOR = None
TOKENIZER = None
SAVE_FILE = "ban_dich.txt"


def load_model(model_name: str) -> None:
	"""Downloads and loads the translator and tokenizer for the given model if not already active."""
	global CURRENT_MODEL, TRANSLATOR, TOKENIZER
	if CURRENT_MODEL == model_name:
		return TRANSLATOR, TOKENIZER

	if model_name == "ngocdang83/HachimiMT-60-zh-vi":
		subfolder = "/ct2-int8_float32"
	elif model_name == "DanVP/MoxhiMT-60":
		subfolder = "/ct2-int8"
	else:
		raise ValueError(f"Unknown model: {model_name}")

	model_path = snapshot_download(model_name)

	TRANSLATOR = Translator(model_path + subfolder, device="cpu")
	TOKENIZER = AutoTokenizer.from_pretrained(model_name)
	CURRENT_MODEL = model_name


def dịch(text: str, file_obj: str, batch_size: int, progress=gr.Progress()) -> str:
	if file_obj is not None:
		with open(file_obj, "r", encoding="utf-8", errors="ignore") as f:
			text = f.read()

	inputs = []
	for line in text.split("\n"):
		if (l := line.strip()) != "":
			inputs.append(TOKENIZER.convert_ids_to_tokens(TOKENIZER.encode(l, truncation=True)))

	if len(inputs) == 0:
		return ""

	results = []
	for i in progress.tqdm(range(0, len(inputs), batch_size)):
		outputs = TRANSLATOR.translate_batch(
			inputs[i : i + batch_size],
			max_decoding_length=TOKENIZER.model_max_length,
			max_batch_size=batch_size,
			beam_size=4,
			no_repeat_ngram_size=2,
			repetition_penalty=1.2
		)
		results.extend([
			TOKENIZER.decode(TOKENIZER.convert_tokens_to_ids(i.hypotheses[0]), skip_special_tokens=True)
			for i in outputs
		])
	return "\n".join(results)  # streaming outputs


def lưu(text: str) -> str:
	if text.strip() == "":
		return None
	with open(SAVE_FILE, "w", encoding="utf-8") as f:
		f.write(text)

DESCRIPTION = """# dịch máy nhanh truyện chữ tiếng Trung

sử dụng model:
- https://huggingface.co/ngocdang83/HachimiMT-60-zh-vi
- https://huggingface.co/DanVP/MoxhiMT-60

[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/phineas-pta/gg_colab_AI_playground/blob/main/trans_ZH_VI.ipynb)
"""

with gr.Blocks(title="dịch máy nhanh truyện chữ tiếng Trung") as APP:
	gr.Markdown(DESCRIPTION)

	with gr.Row():
		model_dropdown = gr.Dropdown(label="Chọn model dịch", choices=["DanVP/MoxhiMT-60", "ngocdang83/HachimiMT-60-zh-vi"], value="DanVP/MoxhiMT-60")
		batch_size_slider = gr.Number(label="batch size", minimum=1, maximum=2048, value=64, step=1, precision=0)
		translate_btn = gr.Button("Dịch", variant="primary")

	with gr.Row():
		with gr.Column():
			input_text = gr.Textbox(label="Nhập hoặc dán text tiếng Trung", lines=12)
			upload_file = gr.File(label="Tải lên tệp văn bản (.txt)", file_types=[".txt"], file_count="single")

		with gr.Column():
			output_text = gr.Textbox(label="Bản dịch tiếng Việt", lines=16, interactive=False)
			save_btn = gr.DownloadButton("Lưu bản dịch (.txt)", variant="secondary", value=SAVE_FILE)

	model_dropdown.change(load_model, inputs=[model_dropdown])
	translate_btn.click(dịch, inputs=[input_text, upload_file, batch_size_slider], outputs=[output_text])
	save_btn.click(lưu, inputs=[output_text])

	APP.load(load_model, inputs=[model_dropdown]) # pre-load

if __name__ == "__main__":
	APP.launch(theme="citrus", ssr_mode=False) # to display badge (see gradio-app/gradio#10140)