Update app.py

app.py

@@ -3,18 +3,18 @@ import torch.nn as nn
 import joblib
 import hashlib
 from collections import Counter
-import numpy as np
 import gradio as gr
 
-# --- utils ---
+# ========== utils ==========
 def ngrams(sentence, n=1, lc=True):
-    sentence = sentence.lower()
+    if lc:
+        sentence = sentence.lower()
     return [sentence[i:i+n] for i in range(len(sentence) - n + 1)]
 
 def all_ngrams(sentence, max_ngram=3, lc=True):
     result = []
     for i in range(1, max_ngram + 1):
-        result += [ngrams(sentence, n=i, lc=lc)]
+        result.append(ngrams(sentence, n=i, lc=lc))
     return result
 
 MAX_CHARS = 521
@@ -54,11 +54,11 @@ def build_freq_dict(sentence):
     freqs = list(map(calc_rel_freq, hngrams))
     return shift_keys(freqs, MAX_SHIFT)
 
-# --- load artifacts ---
+# ========== load artifacts ==========
 vectorizer = joblib.load("nld_vectorizer.joblib")
 idx2lang = joblib.load("nld_lang_codes.joblib")
 
-input_dim = len(vectorizer.vocabulary_)
+input_dim = len(vectorizer.feature_names_)   # 确保和训练时一致
 num_classes = len(idx2lang)
 
 model = nn.Sequential(
@@ -66,31 +66,28 @@ model = nn.Sequential(
     nn.ReLU(),
     nn.Linear(50, num_classes)
 )
-model.load_state_dict(torch.load("nld (1).pth", map_location="cpu"))
+
+state_dict = torch.load("nld.pth", map_location="cpu")
+model.load_state_dict(state_dict)
 model.eval()
 
-# --- prediction ---
+# ========== prediction ==========
 def detect_lang(text: str):
     feat_dict = build_freq_dict(text)
     X = vectorizer.transform([feat_dict])
-    if hasattr(X, "toarray"):
-        X = X.toarray()
-    X = torch.from_numpy(X.astype("float32"))
-
+    X_tensor = torch.from_numpy(X.toarray().astype("float32"))
     with torch.no_grad():
-        logits = model(X)
-        pred_idx = torch.argmax(logits, dim=-1).item()
+        logits = model(X_tensor)
+        pred_idx = torch.argmax(logits, dim=1).item()
     return idx2lang[pred_idx]
 
-# --- UI ---
+# ========== Gradio UI ==========
 with gr.Blocks(title="Language Detector") as demo:
-    gr.Markdown("# Language Detector")
+    gr.Markdown("## Language Detector")
     with gr.Row():
-        with gr.Column():
-            src_text = gr.Textbox(label="Enter text", placeholder="Type here...")
-            btn = gr.Button("Detect Language")
-        with gr.Column():
-            out_lang = gr.Textbox(label="Predicted language", interactive=False)
-    btn.click(fn=detect_lang, inputs=src_text, outputs=out_lang)
+        text_in = gr.Textbox(label="Input text", placeholder="Type something...")
+        text_out = gr.Textbox(label="Predicted language", interactive=False)
+    btn = gr.Button("Detect")
+    btn.click(fn=detect_lang, inputs=text_in, outputs=text_out)
 
 demo.launch()