push max token

app.py

@@ -391,21 +391,27 @@ Produce ONLY valid JSON with these exact fields:
         prompt_tokens = tokenizer.encode(prompt, return_tensors="pt")
         prompt_token_count = prompt_tokens.shape[1]
         
+        max_context = 4096
         max_input_tokens = 3800
-        max_output_tokens = 4096 - max_input_tokens
         
         if prompt_token_count > max_input_tokens:
             logger.warning(f"Prompt is {prompt_token_count} tokens, truncating to {max_input_tokens}")
             prompt_tokens = prompt_tokens[:, :max_input_tokens]
             prompt = tokenizer.decode(prompt_tokens[0], skip_special_tokens=True)
+            prompt_token_count = max_input_tokens
         
-        logger.info(f"Input tokens: ~{prompt_token_count}, Max output tokens: {max_output_tokens}")
+        max_output_tokens = max_context - prompt_token_count - 50
         
-        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=max_input_tokens).to(model.device)
+        logger.info(f"Input tokens: {prompt_token_count}, Available output tokens: {max_output_tokens}")
+        
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=prompt_token_count).to(model.device)
+        
+        output_limit = min(1500, max_output_tokens)
+        logger.info(f"Setting max_new_tokens to {output_limit}")
         
         outputs = model.generate(
             **inputs,
-            max_new_tokens=min(1500, max_output_tokens),
+            max_new_tokens=output_limit,
             temperature=0.3,
             do_sample=True,
             top_p=0.95,
@@ -462,19 +468,23 @@ Produce a FINAL comprehensive review with the same JSON structure as before, con
         prompt_tokens = tokenizer.encode(prompt, return_tensors="pt")
         prompt_token_count = prompt_tokens.shape[1]
         
+        max_context = 4096
         max_input_tokens = 3800
-        max_output_tokens = 4096 - max_input_tokens
         
         if prompt_token_count > max_input_tokens:
             logger.warning(f"Combine prompt is {prompt_token_count} tokens, truncating to {max_input_tokens}")
             prompt_tokens = prompt_tokens[:, :max_input_tokens]
             prompt = tokenizer.decode(prompt_tokens[0], skip_special_tokens=True)
+            prompt_token_count = max_input_tokens
         
-        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=max_input_tokens).to(model.device)
+        max_output_tokens = max_context - prompt_token_count - 50
         
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=prompt_token_count).to(model.device)
+        
+        output_limit = min(1500, max_output_tokens)
         outputs = model.generate(
             **inputs,
-            max_new_tokens=min(1500, max_output_tokens),
+            max_new_tokens=output_limit,
             temperature=0.3,
             do_sample=True,
             top_p=0.95,
@@ -583,19 +593,23 @@ Return ONLY valid JSON:
         prompt_tokens = tokenizer.encode(prompt, return_tensors="pt")
         prompt_token_count = prompt_tokens.shape[1]
         
+        max_context = 4096
         max_input_tokens = 3800
-        max_output_tokens = 4096 - max_input_tokens
         
         if prompt_token_count > max_input_tokens:
             logger.warning(f"Improvement prompt is {prompt_token_count} tokens, truncating to {max_input_tokens}")
             prompt_tokens = prompt_tokens[:, :max_input_tokens]
             prompt = tokenizer.decode(prompt_tokens[0], skip_special_tokens=True)
+            prompt_token_count = max_input_tokens
+        
+        max_output_tokens = max_context - prompt_token_count - 50
         
-        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=max_input_tokens).to(model.device)
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=prompt_token_count).to(model.device)
         
+        output_limit = min(1000, max_output_tokens)
         outputs = model.generate(
             **inputs,
-            max_new_tokens=min(1000, max_output_tokens),
+            max_new_tokens=output_limit,
             temperature=0.4,
             do_sample=True,
             top_p=0.95,