Update app.py

app.py

@@ -1,32 +1,33 @@
 import gradio as gr
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
+import torch
 
-# 加载模型和tokenizer
-model = AutoModelForSequenceClassification.from_pretrained("SecurityXuanwuLab/HaS-820m")
-tokenizer = AutoTokenizer.from_pretrained("SecurityXuanwuLab/HaS-820m")
+# 加载模型和分词器
+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)
-    outputs = model(**inputs)
-    return outputs.logits.argmax(dim=-1).item()
+    with torch.no_grad():
+        outputs = model(**inputs)
+    logits = outputs.logits
+    predicted_class = torch.argmax(logits, dim=-1).item()
+    return f"预测类别：{predicted_class}"
 
-# 使用 gr.Blocks 创建界面
+# 使用 Gradio Blocks 构建界面
 with gr.Blocks() as demo:
-    # 设置页面标题
     gr.Markdown("<h1 style='text-align: center;'>欢迎使用 HaS-820m 模型预测</h1>")
     
-    # 创建一个简单的文本框和按钮
-    with gr.Row():  # Row 用于水平布局
-        with gr.Column():  # Column 用于垂直布局
+    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)
     
     # 可选：自定义样式