Update app.py

app.py

@@ -72,8 +72,12 @@ def get_prediction(inputs):
         outputs = model(**inputs)
         logits = outputs.last_hidden_state[:, 0, :]  # 取CLS标记的输出进行分类
         pred_prob = torch.softmax(logits, dim=1)
-        pred = torch.argmax(pred_prob, dim=1)
-        return class_names[pred.item()]
+        pred = torch.argmax(pred_prob, dim=1).item()
+        if pred in class_names:
+            return class_names[pred]
+        else:
+            print(f"Warning: Prediction index {pred} not found in class_names.")
+            return "Unknown"  # 或者其他默认的响应
 
 # vectorizer= nltk_u.vectorizer()
 # vectorizer.fit(train_data.text)