Update app.py

app.py

@@ -12,7 +12,7 @@ generator = pipeline(
     dtype=torch.bfloat16
 )
 
-def generate_with_chat_template(messages, max_new_tokens=2048, do_sample=True, temperature=0.7, num_return_sequences=1):
+def generate_with_chat_template(messages, max_new_tokens=1536, do_sample=True, temperature=0.7, num_return_sequences=1):
     """
     Applies the model's chat template to structure the prompts properly, 
     preventing prompt injection or completion confusion.
@@ -78,7 +78,7 @@ def tot_search(problem: str, branches: int = 3, max_depth: int = 3) -> str:
             
             outputs = generate_with_chat_template(
                 messages, 
-                max_new_tokens=2048, 
+                max_new_tokens=1536, 
                 do_sample=True, 
                 temperature=0.7, 
                 num_return_sequences=branches
@@ -104,7 +104,7 @@ def tot_search(problem: str, branches: int = 3, max_depth: int = 3) -> str:
                     }
                 ]
                 
-                eval_outs = generate_with_chat_template(eval_messages, max_new_tokens=2048, do_sample=False)
+                eval_outs = generate_with_chat_template(eval_messages, max_new_tokens=1536, do_sample=False)
                 eval_text = eval_outs[0].lower()
                 stats["nodes_evaluated"] += 1
                 
@@ -164,7 +164,7 @@ def tot_search(problem: str, branches: int = 3, max_depth: int = 3) -> str:
         }
     ]
     
-    final_outs = generate_with_chat_template(final_messages, max_new_tokens=2048, do_sample=False)
+    final_outs = generate_with_chat_template(final_messages, max_new_tokens=1536, do_sample=False)
     final_response = final_outs[0]
     
     # Strip everything down to get just the clean answer