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

# 載入 tokenizer 和模型
tokenizer = BertTokenizer.from_pretrained("ckiplab/bert-base-chinese")
model = BertForSequenceClassification.from_pretrained("ckiplab/bert-base-chinese", num_labels=2)
model.eval()

# 預測函式
def predict(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        prediction = torch.argmax(logits, dim=1).item()
        return "詐騙訊息" if prediction == 1 else "正常訊息"

# Gradio UI
iface = gr.Interface(fn=predict, inputs="text", outputs="text", title="Line詐騙訊息辨識器")
iface.launch()