fix: return final model answer after reasoning

app.py

@@ -98,23 +98,29 @@ def generate(
     inputs = tokenizer.apply_chat_template(
         messages,
         add_generation_prompt=True,
+        tokenize=True,
+        return_dict=True,
         return_tensors="pt",
     ).to(model.device)
 
     with torch.inference_mode():
         generated = model.generate(
-            inputs,
-            max_new_tokens=512,
+            **inputs,
+            max_new_tokens=4096,
             do_sample=True,
             temperature=0.6,
             top_p=0.95,
             top_k=20,
         )
 
+    output_ids = generated[0, inputs["input_ids"].shape[-1] :].tolist()
+    think_end_token = tokenizer.convert_tokens_to_ids("</think>")
+    answer_start = len(output_ids) - output_ids[::-1].index(think_end_token)
+
     return tokenizer.decode(
-        generated[0, inputs.shape[-1] :],
+        output_ids[answer_start:],
         skip_special_tokens=True,
-    )
+    ).strip()
 
 
 CSS = """