Update app.py

app.py

@@ -41,12 +41,12 @@ def load_embedder():
         emb_model.eval()
 
 @torch.no_grad()
-def embed(text: str) -> torch.Tensor:
+def embed(text:str)->torch.Tensor:
+    """Return L2-normalised embedding vector."""
     load_embedder()
-    with torch.no_grad():
-        inputs = emb_tokenizer(text, return_tensors="pt", truncation=True).to(emb_model.device)
-        vec = emb_model(**inputs).last_hidden_state[:, 0]
-        return F.normalize(vec, dim=-1).cpu()  
+    inputs = emb_tokenizer(text, return_tensors="pt", truncation=True).to(emb_model.device)
+    vec = emb_model(**inputs).last_hidden_state[:, 0]  # CLS pooling
+    return F.normalize(vec, dim=-1).squeeze(0)
 
 # ---------- 2. tiny in-memory KB shared by Gradio & API ----------------------
 # ---------- 2. Tiny in-memory knowledge-base -------------------------------
@@ -67,7 +67,7 @@ def add_docs(user_id: str, docs: list[str]) -> int:
         return 0
 
     load_embedder()                                # lazy-load once
-    new_vecs = torch.stack([embed(t) for t in docs]).cpu() 
+    new_vecs = torch.stack([embed(t) for t in docs])
     store = kb[user_id]                            # auto-creates via defaultdict
     store["texts"].extend(docs)
     store["vecs"] = (
@@ -119,7 +119,7 @@ def answer(system: str, context: str, question: str, user_id="demo", history="No
         context_list = [context]
         # 1.  Retrieve top-k similar passages
         if history == "Some":
-            q_vec  = embed(question).cpu()
+            q_vec  = embed(question)
             store  = kb[user_id]
             sims   = torch.matmul(store["vecs"], q_vec)          # [N]
             k      = min(4, sims.numel())
@@ -130,23 +130,7 @@ def answer(system: str, context: str, question: str, user_id="demo", history="No
             context_list += store["texts"]
 
         # 2.  Build a Qwen-chat prompt (helper defined earlier)
-        MAX_PROMPT_TOKENS = 8192          # 8 k is ~4 GB KV-cache
-
-    prompt = build_qwen_prompt(system, context_list, question)
-    tokens = tokenizer(prompt, return_tensors="pt", add_special_tokens=False)
-
-    if tokens.input_ids.size(1) > MAX_PROMPT_TOKENS:
-    # keep the last MAX_PROMPT_TOKENS tokens (most recent content)
-        tokens = {k: v[:, -MAX_PROMPT_TOKENS:] for k, v in tokens.items()}
-
-    tokens = {k: v.to(chat_model.device) for k, v in tokens.items()}
-
-    output = chat_model.generate(
-        **tokens,
-        max_new_tokens=512,
-        max_length=MAX_PROMPT_TOKENS + 512,
-    )
-
+        prompt = build_qwen_prompt(system, context_list, question)
 
         # 3.  Generate and strip everything before the assistant tag
         load_chat()
@@ -247,4 +231,4 @@ if __name__ == "__main__":
     # launch Gradio on a background thread
     demo.queue().launch(share=False, prevent_thread_lock=True)
     # then start FastAPI (uvicorn blocks main thread)
-    uvicorn.run(api, host="0.0.0.0", port=8000)
+    uvicorn.run(api, host="0.0.0.0", port=8000)