import gradio as gr
from gec_model import GecBERTModel

title = "Vietnamese Capitalization and Punctuation recovering"
description = "Auto capitalize and punctuations for plain, lowercase Vietnamese text."
article = "Coming soon"

examples = [
  ["viBERT", "theo đó thủ tướng dự kiến tiếp bộ trưởng nông nghiệp mỹ tom wilsack bộ trưởng thương mại mỹ gina raimondo bộ trưởng tài chính janet yellen gặp gỡ thượng nghị sĩ patrick leahy và một số nghị sĩ mỹ khác"],
  ["viBERT", "những gói cước năm g mobifone sẽ mang đến cho bạn những trải nghiệm mới lạ trên cả tuyệt vời so với mạng bốn g thì tốc độ truy cập mạng năm g mobifone được nhận định là siêu đỉnh với mức truy cập nhanh gấp 10 lần"],
]

model_kwargs = {
    "max_len": 64,
    "min_len": 3,
    "iterations": 3,
    "min_error_probability": 0.2,
    "lowercase_tokens": False,
    "log": False,
    "confidence": 0,
    "split_chunk": True,
    "chunk_size": 48,
    "overlap_size": 16,
    "min_words_cut": 8
}
model_dict = {
    "viBERT": GecBERTModel(
        vocab_path="vocabulary",
        model_paths='dragonSwing/vibert-capu',
        **model_kwargs
    ),
    "XLM-RoBERTa-base": GecBERTModel(
        vocab_path="vocabulary",
        model_paths='dragonSwing/xlm-roberta-capu',
        **model_kwargs
    ),
}

def fn(model_choice, input):
    if model_choice in model_dict:
        return model_dict[model_choice](input)[0]
    else:
        return "Unsupported model"

gr.Interface(fn, [gr.inputs.Dropdown(list(model_dict.keys())), gr.inputs.Textbox(lines=5)], "text", examples=examples, title=title, description=description, article=article).launch()