import fasttext
import gradio as gr

# 训练数据（直接在代码中写入）
train_data = [
    "__label__查余额 查询余额",
    "__label__查余额 账户里有多少钱",
    "__label__信用卡申请 申请信用卡",
    "__label__信用卡申请 我想办一张信用卡",
    "__label__贷款利率 贷款利率是多少",
    "__label__贷款利率 现在房贷利息多少"
]

# 将训练数据保存到一个临时文件，以便 FastText 使用
train_file = "train_data.txt"
with open(train_file, "w") as f:
    for line in train_data:
        f.write(line + "\n")

# 训练 FastText 分类模型
model = fasttext.train_supervised(input=train_file, epoch=25, lr=0.5, wordNgrams=2, dim=100)

# 保存训练好的模型
model.save_model("intent_classifier.bin")

# 定义一个用于预测的函数
def classify_intent(text):
    label, confidence = model.predict(text)
    return label[0].replace("__label__", ""), confidence[0]

# 创建 Gradio 界面
iface = gr.Interface(fn=classify_intent,
                     inputs=gr.Textbox(label="请输入查询内容"),
                     outputs=[gr.Textbox(label="预测意图"), gr.Textbox(label="置信度")])

# 启动 Gradio 应用
iface.launch()