Aggressive fix for CoT truncation: increase min_new_tokens to 1500, suppress EOS token for CoT requests, cap max_new_tokens

gradio_app.py

@@ -106,8 +106,10 @@ def generate_response(prompt, temperature=0.8):
         
         # Set minimum tokens based on request type
         if is_cot_request:
-            min_tokens = 1000  # Much higher minimum for CoT to ensure complete responses
-            logger.info("Detected Chain of Thinking request - using min_new_tokens=1000")
+            min_tokens = 1500  # Even higher minimum for CoT to ensure complete responses
+            # Also reduce max_new_tokens to ensure we don't hit context limits
+            gen_max_new_tokens = min(gen_max_new_tokens, 2048)  # Cap at 2048 for CoT
+            logger.info(f"Detected Chain of Thinking request - using min_new_tokens={min_tokens}, max_new_tokens={gen_max_new_tokens}")
         else:
             min_tokens = 200  # Standard minimum
 
@@ -126,23 +128,37 @@ def generate_response(prompt, temperature=0.8):
         
         # Generate response
         with torch.no_grad():
-            outputs = model_manager.model.generate(
+            # For CoT requests, be more aggressive about preventing early stopping
+            if is_cot_request:
+                # Suppress EOS token for CoT to prevent early termination
+                eos_token_id = None
+                suppress_tokens = [model_manager.tokenizer.eos_token_id] if model_manager.tokenizer.eos_token_id is not None else None
+            else:
+                eos_token_id = model_manager.tokenizer.eos_token_id
+                suppress_tokens = None
+                
+            generation_kwargs = {
                 **inputs,
-                max_new_tokens=gen_max_new_tokens,
-                temperature=temperature,
-                top_p=0.95,
-                do_sample=True,
-                num_beams=1,
-                pad_token_id=model_manager.tokenizer.eos_token_id,
-                # Keep EOS but rely primarily on post-decode stop to capture full JSON
-                eos_token_id=model_manager.tokenizer.eos_token_id,
-                early_stopping=False,
-                repetition_penalty=1.05,
-                no_repeat_ngram_size=0,
-                length_penalty=1.0,
-                # Dynamic minimum based on request type
-                min_new_tokens=min_tokens
-            )
+                "max_new_tokens": gen_max_new_tokens,
+                "temperature": temperature,
+                "top_p": 0.95,
+                "do_sample": True,
+                "num_beams": 1,
+                "pad_token_id": model_manager.tokenizer.eos_token_id,
+                "early_stopping": False,
+                "repetition_penalty": 1.05,
+                "no_repeat_ngram_size": 0,
+                "length_penalty": 1.0,
+                "min_new_tokens": min_tokens
+            }
+            
+            # Add EOS suppression for CoT
+            if eos_token_id is not None:
+                generation_kwargs["eos_token_id"] = eos_token_id
+            if suppress_tokens is not None:
+                generation_kwargs["suppress_tokens"] = suppress_tokens
+                
+            outputs = model_manager.model.generate(**generation_kwargs)
         
         # Decode the response
         generated_text = model_manager.tokenizer.decode(outputs[0], skip_special_tokens=True)