Upload handler.py

handler.py

@@ -181,6 +181,7 @@ class _DecodingParams:
     temperature: float = 0.8
     top_p: float = 0.9
     max_new_tokens: int = 256
+    min_new_tokens: int = 16  # Minimum tokens before allowing EOS
     stop_quality: float = 0.6
 
     @classmethod
@@ -508,16 +509,30 @@ class EndpointHandler:
             if logits is None:
                 break
             last_index = min(len(token_ids) - 1, logits.shape[1] - 1)
-            next_logits = logits[0, last_index]
+            next_logits = logits[0, last_index].copy()
+
+            # Apply strong penalty to EOS token if we haven't reached min_new_tokens
+            # This reduces the probability of generating EOS prematurely
+            if steps < decoding.min_new_tokens:
+                next_logits[self._tokenizer.eos_token_id] -= 10.0
+
             next_token = self._sample_next_token(next_logits, decoding, rng)
             token_ids.append(int(next_token))
             steps += 1
 
+            # Check if we generated EOS prematurely and replace with space
+            if token_ids[-1] == self._tokenizer.eos_token_id and steps < decoding.min_new_tokens:
+                # Find space token ID (fallback to 'a' if space not found)
+                space_token_id = self._tokenizer._token_to_id.get(" ", self._tokenizer._token_to_id.get("a", self._tokenizer.unk_token_id))
+                token_ids[-1] = space_token_id
+                # Note: In production, add logging here to track how often this happens
+
             outputs = self._run_candidate(base_feed, token_ids)
             formatted_outputs = outputs
             quality = self._extract_q_hat(outputs)
 
-            if token_ids[-1] == self._tokenizer.eos_token_id:
+            # Only allow EOS break if we've generated at least min_new_tokens (excluding BOS)
+            if token_ids[-1] == self._tokenizer.eos_token_id and steps >= decoding.min_new_tokens:
                 break
             if self._token_sequence_length > 0 and len(token_ids) >= self._token_sequence_length:
                 break