import gradio as gr
from transformers import BertTokenizerFast, BertForSequenceClassification
import torch

model = BertForSequenceClassification.from_pretrained('./ch-sent-check-model')
tokenizer = BertTokenizerFast.from_pretrained('./ch-sent-check-model')

def judge(sentence):
    input_ids = tokenizer(sentence,return_tensors='pt')['input_ids']
    out = model(input_ids)
    logits = out.logits
    pred = torch.argmax(logits,dim=-1).item()
    pred_text = 'Incorrect' if pred == 0 else 'Correct'
    return pred_text

iface = gr.Interface(
    fn=judge, 
    inputs=gr.Textbox(
        label="Initial text",
        lines=1,
        value="請注意用字的鄭確性",
    ), 
    outputs="text"
)
iface.launch()