"""
Write a gradio ui to demonstrate the classifier
"""

import transformers
import gradio as gr

# load cls pipeline
cls_type = transformers.pipeline("text-classification", model="./data/type/output/")
cls_location = transformers.pipeline(
    "text-classification", model="./data/location/output/"
)
cls_department = transformers.pipeline(
    "text-classification", model="./data/department/output/"
)


def classify_texts(text1, text2):
    text = text1 + text2
    # 示例分类逻辑
    type_ = cls_type(text)[0]["label"]
    location = cls_location(text)[0]["label"]
    department = cls_department(text)[0]["label"]
    return type_, location, department

# 创建 Gradio 界面，自定义输入和输出标签
iface = gr.Interface(
    fn=classify_texts,
    inputs=[gr.Textbox(label="具体情况描述"), gr.Textbox(label="市民诉求")],
    outputs=[
        gr.Label(label="地区分类结果"),
        gr.Label(label="部门分类结果"),
        gr.Label(label="工单类型分类结果"),
    ],
    title="12345文本分类器",
    description="输入文本，输出分类结果。",
)

# 启动界面
iface.launch()