Tin113/text_error_correction

1import os2import torch3 4# CRITICAL: Redirect cache to temporary storage5os.environ['TORCH_HOME'] = '/tmp/torch_cache'6os.environ['HUB_DIR'] = '/tmp/torch_hub'7os.environ['TMPDIR'] = '/tmp'8torch.hub.set_dir('/tmp/torch_hub')9 10 11import gradio as gr12import torch13from transformers import (14    AutoTokenizer,15    AutoModelForSeq2SeqLM,16    T5Tokenizer,17    T5ForConditionalGeneration18)19import re20 21HF_USERNAME = "Tin113"22# -----------------------------------------23 24BART_MODEL_REPO = f"{HF_USERNAME}/bart_model"25VIT5_MODEL_REPO = f"{HF_USERNAME}/vit5_model"26 27DEVICE = "cuda" if torch.cuda.is_available() else "cpu"28print(f"Thiết bị sử dụng: {DEVICE}")29 30# Tải các model31try:32    print(f"Đang tải model BART từ {BART_MODEL_REPO}...")33    tokenizer_bart = AutoTokenizer.from_pretrained(BART_MODEL_REPO)34    model_bart = AutoModelForSeq2SeqLM.from_pretrained(BART_MODEL_REPO).to(DEVICE)35    model_bart.eval()36    print("Tải model BART thành công.")37except Exception as e:38    print(f"Lỗi khi tải model BART: {e}")39    model_bart, tokenizer_bart = None, None40 41try:42    print(f"Đang tải model ViT5 từ {VIT5_MODEL_REPO}...")43    tokenizer_vit5 = T5Tokenizer.from_pretrained(VIT5_MODEL_REPO)44    model_vit5 = T5ForConditionalGeneration.from_pretrained(VIT5_MODEL_REPO).to(DEVICE)45    model_vit5.eval()46    print("Tải model ViT5 thành công.")47except Exception as e:48    print(f"Lỗi khi tải model ViT5: {e}")49    model_vit5, tokenizer_vit5 = None, None50 51def clean_text(text):52    if not isinstance(text, str): return ""53    return re.sub(r'\s+', ' ', text).strip()54 55def correct_grammar(sentence, model_choice):56    if not sentence.strip(): return "Vui lòng nhập một câu."57 58    model, tokenizer, prefix = None, None, ""59    if model_choice == "BARTpho-syllable":60        if model_bart:61            model, tokenizer, prefix = model_bart, tokenizer_bart, "Fix: "62        else:63            return "Lỗi: Model BART không khả dụng. Vui lòng kiểm tra lại Space."64    elif model_choice == "ViT5-base":65        if model_vit5:66            model, tokenizer, prefix = model_vit5, tokenizer_vit5, "sửa lỗi: "67        else:68            return "Lỗi: Model ViT5 không khả dụng. Vui lòng kiểm tra lại Space."69    70    input_text = prefix + sentence71    input_ids = tokenizer(input_text, return_tensors="pt", max_length=256, truncation=True, padding=True).input_ids.to(DEVICE)72 73    with torch.no_grad():74        outputs = model.generate(input_ids, max_length=276, num_beams=2, early_stopping=True, repetition_penalty=1.05, no_repeat_ngram_size=2)75    76    return clean_text(tokenizer.decode(outputs[0], skip_special_tokens=True))77 78description = """79Demo sửa lỗi chính tả tiếng Việt sử dụng hai model: BARTpho-syllable và ViT5-base.801. Nhập câu lỗi vào ô bên dưới.812. Chọn model bạn muốn dùng.823. Nhấn "Submit" để xem kết quả.83"""84 85demo = gr.Interface(86    fn=correct_grammar,87    inputs=[88        gr.Textbox(lines=5, label="Nhập câu tiếng Việt bị lỗi"),89        gr.Radio(choices=["BARTpho-syllable", "ViT5-base"], value="ViT5-base", label="Chọn Model")90    ],91    outputs=gr.Textbox(label="Câu đã được sửa"),92    title="Sửa lỗi chính tả Tiếng Việt",93    description=description,94)95 96if __name__ == "__main__":97    demo.launch()