import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

# 加载模型和分词器
model_name = "SecurityXuanwuLab/HaS-820m"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 定义推理函数
def predict(text):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        outputs = model(**inputs)
    logits = outputs.logits
    predicted_class = torch.argmax(logits, dim=-1).item()
    return f"预测类别：{predicted_class}"

theme = gr.themes.Monochrome(
    primary_hue="stone",
    secondary_hue="stone",
    neutral_hue="stone",
).set(
    body_background_fill='*secondary_200',
    body_background_fill_dark='*secondary_200',
    body_text_color='*secondary_950',
    body_text_color_dark='*secondary_200',
    body_text_color_subdued='*secondary_500',
    body_text_color_subdued_dark='*secondary_500'
)

# 使用 Gradio 的内置主题构建界面
with gr.Blocks(theme=theme) as demo:
    gr.Markdown("<h1 style='text-align: center;'>模型试用：腾讯玄武</h1>")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(label="输入文本", placeholder="请输入文本处理...", lines=4)
            output_text = gr.Textbox(label="处理结果", interactive=False)
    
    with gr.Row():
        predict_btn = gr.Button("进行处理")
    
    predict_btn.click(predict, inputs=input_text, outputs=output_text)

# 启动界面
demo.launch()