Update app.py

app.py

@@ -77,14 +77,18 @@ model.load_state_dict(torch.load("nld.pth", map_location="cpu"))
 model.eval()
 
 # --- prediction function ---
-# 修改预测函数为sklearn版本
 def detect_lang(src_sentence):
-    # 特征提取逻辑不变
-    X_test = vectorizer.transform([build_freq_dict(src_sentence)])
-    # 使用sklearn模型预测
-    pred_idx = clf.predict(X_test)[0]
-    return idx2lang[pred_idx]
-    
+    # 直接对单条文本提取特征，用列表推导式（与notebook一致）
+    test_feat_dicts = [build_freq_dict(src_sentence)]
+    # 转换为模型输入
+    X_test = vectorizer.transform(test_feat_dicts)
+    # 后续处理不变
+    if hasattr(X_test, "toarray"):
+        X_test = X_test.toarray()
+    Y_logits = model(torch.Tensor(X_test))
+    pred_languages = torch.argmax(Y_logits, dim=-1).tolist()
+    return list(map(idx2lang.get, pred_languages))[0]
+
 # --- Gradio UI ---
 with gr.Blocks(title="Antons language detector") as demo:
     gr.Markdown("# Antons language detector")