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

# Load model và tokenizer
MODEL_ID = "chi-vi/hirashiba-mt-tiny-zh-vi"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

TOKENIZER = AutoTokenizer.from_pretrained(MODEL_ID)
MODEL = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID).to(DEVICE)

# punctuation
ZH_PUNC = str.maketrans("", "", "，、：；．。！？…⋯～／‧•●『』「」（）《》〈〉［］【】〖〗〔〕｛｝")


@torch.inference_mode()
def dịch(input_text: str) -> str:
	# Tokenize input
	inputs = TOKENIZER(input_text, return_tensors="pt", padding=True, truncation=True).to(DEVICE)
	# Dịch với mô hình (không cần tính gradient)
	output_tokens = MODEL.generate(**inputs, max_length=512)
	# Giải mã kết quả
	return TOKENIZER.decode(output_tokens[0], skip_special_tokens=True)


def translate_text(input_text: str) -> str:
	lines = input_text.split("\n")  # Tách từng dòng
	translated_lines = []

	for line in lines:
		raw_text = line.strip()
		if not raw_text:
			translated_lines.append("")  # Giữ dòng trống
		elif not raw_text.translate(ZH_PUNC):
			translated_lines.append(raw_text)  # Giữ dòng chỉ có dấu câu
		else:
			translated_lines.append(dịch(raw_text))
		yield "\n".join(translated_lines)  # streaming outputs


app = gr.Interface(
	title="công cụ dịch truyện bằng AI",
	description="sử dụng mô hình đào tạo bởi Chivi",
	theme="soft",
	fn=translate_text,
	inputs=[gr.Textbox(label="text gốc", lines=10)],
	outputs=[gr.Textbox(label="text dịch", lines=10, interactive=False, show_copy_button=True)],
	submit_btn="Dịch",
	api_name="translate",
	deep_link=False,
)

app.launch()