Update app.py

app.py

@@ -136,24 +136,24 @@ def answer(system: str, context: str, question: str, user_id="demo", history="No
         # 3.  Generate and strip everything before the assistant tag
         load_chat()
         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()}
+            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()}
+    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
+    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
     except Exception as e:
         return f"Error in app.py: {e}"
     finally:
@@ -241,17 +241,17 @@ def rag(req:QueryReq):
     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()}
+    )
+    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()}
+    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,
-)
+    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()