REEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEE

app.py

@@ -17,6 +17,7 @@ from collections import defaultdict
 HF_TOKEN = os.getenv("HF_token")
 CHAT_MODEL_ID = "QWen/Qwen1.5-7B-Chat"
 EMB_MODEL_ID  = "mixedbread-ai/mxbai-embed-large-v1"
+MAX_PROMPT_TOKENS = 8192
 
 # --- lazy loaders (unchanged) -------------------------------------------------
 tokenizer, chat_model = None, None
@@ -42,11 +43,11 @@ def load_embedder():
 
 @torch.no_grad()
 def embed(text:str)->torch.Tensor:
-    """Return L2-normalised embedding vector."""
     load_embedder()
-    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)
+    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()
 
 # ---------- 2. tiny in-memory KB shared by Gradio & API ----------------------
 # ---------- 2. Tiny in-memory knowledge-base -------------------------------
@@ -67,7 +68,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])
+    new_vecs = torch.stack([embed(t) for t in docs]).cpu()
     store = kb[user_id]                            # auto-creates via defaultdict
     store["texts"].extend(docs)
     store["vecs"] = (
@@ -119,7 +120,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)
+            q_vec  = embed(question).cpu()
             store  = kb[user_id]
             sims   = torch.matmul(store["vecs"], q_vec)          # [N]
             k      = min(4, sims.numel())
@@ -134,8 +135,22 @@ def answer(system: str, context: str, question: str, user_id="demo", history="No
 
         # 3.  Generate and strip everything before the assistant tag
         load_chat()
-        inputs = tokenizer(prompt, return_tensors="pt").to(chat_model.device)
-        output = chat_model.generate(**inputs, max_new_tokens=512)
+        tokens = tokenizer(
+    prompt,
+    return_tensors="pt",
+    add_special_tokens=False,     # important – we already built chat template
+)
+if tokens.input_ids.size(1) > MAX_PROMPT_TOKENS:
+    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()}
+
+# --- generate ------------------------------------------------------
+output = chat_model.generate(
+    **tokens,
+    max_new_tokens=512,
+    max_length=MAX_PROMPT_TOKENS + 512,
+)
         full   = tokenizer.decode(output[0], skip_special_tokens=True)
         reply  = full.split("<|im_start|>assistant")[-1].strip()
         return reply
@@ -220,8 +235,22 @@ def rag(req:QueryReq):
     prompt = build_qwen_prompt(context, req.question)
 
     load_chat()
-    inputs = tokenizer(prompt, return_tensors="pt").to(chat_model.device)
-    out = chat_model.generate(**inputs, max_new_tokens=512)
+    tokens = tokenizer(
+    prompt,
+    return_tensors="pt",
+    add_special_tokens=False,
+)
+if tokens.input_ids.size(1) > MAX_PROMPT_TOKENS:
+    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()}
+
+out = chat_model.generate(
+    **tokens,
+    max_new_tokens=512,
+    max_length=MAX_PROMPT_TOKENS + 512,
+)
+
     full = tokenizer.decode(out[0], skip_special_tokens=True)
     ans  = full.split("<|im_start|>assistant")[-1].strip()
     return {"answer": ans}