fix

app.py

@@ -123,10 +123,10 @@ def answer(system: str, context: str, question: str, user_id="demo", history="No
             q_vec  = embed(question).view(-1).cpu()
             store  = kb[user_id]
             sims   = torch.matmul(store["vecs"], q_vec)          # [N]
-            if sims.numel() == 0:
-                return "Knowledge base is empty or corrupted."
+            if sims.shape[1] == 1:
+                sims = sims.squeeze(1)
             k      = min(4, store["vecs"].size(0))
-            idxs   = torch.topk(sims, k=k).indices.tolist()
+            idxs   = torch.topk(sims, k=k, dim=0).indices.tolist()
             context_list += [store["texts"][i] for i in idxs]
         elif history == "All":
             store  = kb[user_id]