Update app.py

app.py

@@ -10,6 +10,10 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
 model.eval()
 
+# 确定设备（CPU 或 GPU）
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)  # 将模型加载到正确的设备
+
 # 标签列表
 labels_list = ['上供', '中人', '中央亞', '中央行政', '中央軍', '主副食', '交通', '人事', '人文敎育',
                '人物', '任免', '住生活', '佛敎', '保健', '倉庫', '倫理', '倭', '儀式', '儒學', '元',
@@ -33,10 +37,20 @@ labels_list = ['上供', '中人', '中央亞', '中央行政', '中央軍', '
 # 推理函数
 @torch.no_grad()
 def classify(text: str, threshold: float = 0.5):
+    # 将输入文本转化为模型需要的格式
     enc = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
+    
+    # 将输入数据移动到正确的设备（GPU/CPU）
+    enc = {key: value.to(device) for key, value in enc.items()}
+    
+    # 获取模型的预测值
     logits = model(**enc).logits.sigmoid()  # 获取预测的概率
+    
+    # 获取预测结果
     scores = logits[0].tolist()
     preds = [labels_list[i] for i, score in enumerate(scores) if score >= threshold]
+    
+    # 返回预测标签及得分
     return preds, {labels_list[i]: score for i, score in enumerate(scores)}
 
 # 创建 Gradio 接口