Fix generate(): full 2048 canvas + EOSBOS anchor + trim_at_eos for coherent generation

generate() now denoises a full canvas_length (default 2048) canvas instead of prompt+max_new_tokens, anchors the <|endoftext|><|beginoftext|> delimiter at prompt+max_new_tokens (max_new_tokens is the content budget), and trims output at the first EOS (trim_at_eos). Full untrimmed canvas returned as out.canvas. New params canvas_length/anchor_eosbos/trim_at_eos all resolve via generation_config.

generation_sumi.py

@@ -50,6 +50,7 @@ class SumiGenerationOutput(ModelOutput):
     past_key_values: Cache | None = None
     logits: None = None
     hidden_states: None = None
+    canvas: torch.LongTensor | None = None  # full untrimmed denoised canvas
 
 
 class SumiGenerationConfig(GenerationConfig):
@@ -333,6 +334,32 @@ class SumiGenerationMixin(GenerationMixin):
                     return resolved
         return default
 
+    def _trim_at_eos(self, generated_ids, prompt_length):
+        """Cut each row at the first EOS in its generated region (dropping the EOS, the
+        anchored BOS delimiter, and the denoised tail), then right-pad the batch back into
+        a rectangular tensor. Prompt tokens are kept. Mirrors the eval harness's
+        ``_extract_text`` (`gen_ids[:gen_ids.index(eos)]`). The pad filler is a special
+        token, so ``decode(..., skip_special_tokens=True)`` yields clean text."""
+        eos_id = self.config.eos_token_id
+        if eos_id is None:
+            return generated_ids
+        pad_id = self.config.pad_token_id
+        pad_id = eos_id if pad_id is None else pad_id
+        batch_size = generated_ids.shape[0]
+        rows, max_len = [], prompt_length
+        for i in range(batch_size):
+            row = generated_ids[i]
+            hit = (row[prompt_length:] == eos_id).nonzero(as_tuple=True)[0]
+            cut = prompt_length + int(hit[0].item()) if hit.numel() > 0 else row.shape[0]
+            rows.append(row[:cut])
+            max_len = max(max_len, cut)
+        out = torch.full((batch_size, max(max_len, 1)), pad_id,
+                         dtype=generated_ids.dtype, device=generated_ids.device)
+        for i, row in enumerate(rows):
+            if row.shape[0] > 0:
+                out[i, : row.shape[0]] = row
+        return out
+
     @torch.no_grad()
     def generate(
         self,
@@ -351,6 +378,9 @@ class SumiGenerationMixin(GenerationMixin):
         tokens_per_step: Optional[int] = None,
         frozen: Optional[List] = None,
         denoise_end: Optional[List[int]] = None,
+        canvas_length: Optional[int] = None,
+        anchor_eosbos: Optional[bool] = None,
+        trim_at_eos: Optional[bool] = None,
         progress_callback: Optional[Callable] = None,
         **kwargs,
     ) -> SumiGenerationOutput:
@@ -384,6 +414,9 @@ class SumiGenerationMixin(GenerationMixin):
         tokens_per_step = self._resolve(
             tokens_per_step, "tokens_per_step", 1, generation_config, self_generation_config
         )
+        canvas_length = self._resolve(canvas_length, "canvas_length", 2048, generation_config, self_generation_config)
+        anchor_eosbos = self._resolve(anchor_eosbos, "anchor_eosbos", True, generation_config, self_generation_config)
+        trim_at_eos = self._resolve(trim_at_eos, "trim_at_eos", True, generation_config, self_generation_config)
 
         device = self.model.embed_tokens.weight.device
         if input_ids is None:
@@ -422,14 +455,29 @@ class SumiGenerationMixin(GenerationMixin):
         if max_new_tokens < 0:
             raise ValueError("max_new_tokens must be non-negative.")
         if max_new_tokens == 0:
-            return SumiGenerationOutput(sequences=input_ids)
-
-        total_length = prompt_length + max_new_tokens
-        if total_length > self.config.max_position_embeddings:
+            return SumiGenerationOutput(sequences=input_ids, canvas=input_ids)
+
+        # The model is trained on a packed, fixed-length canvas, so generation runs on a
+        # full `canvas_length` canvas (default 2048) rather than just prompt+max_new_tokens.
+        # `max_new_tokens` is the content budget: the EOS,BOS document delimiter is anchored
+        # at prompt_length+max_new_tokens, the rest of the canvas is denoised as context, and
+        # decoding is cut at the first EOS (trim_at_eos, default True).
+        ceiling = self.config.max_position_embeddings
+        canvas_length = min(int(canvas_length), ceiling)
+        reserve = 2 if anchor_eosbos else 0
+        if prompt_length + reserve >= canvas_length:
             raise ValueError(
-                f"Requested sequence length {total_length} exceeds max_position_embeddings "
-                f"({self.config.max_position_embeddings})."
+                f"prompt_length ({prompt_length}) leaves no room in canvas_length ({canvas_length})."
             )
+        budget = max(1, min(max_new_tokens, canvas_length - prompt_length - reserve))
+        total_length = canvas_length
+        completion_length = total_length - prompt_length
+
+        # Anchor the EOS,BOS delimiter at the end of the content budget. The caller can
+        # override by passing an explicit `frozen` set (then it takes precedence).
+        if anchor_eosbos and frozen is None:
+            eos_pos = prompt_length + budget
+            frozen = [(eos_pos, self.config.eos_token_id), (eos_pos + 1, self.config.bos_token_id)]
 
         if attention_mask is None:
             attention_mask = torch.ones_like(input_ids)
@@ -449,7 +497,7 @@ class SumiGenerationMixin(GenerationMixin):
                 generator = torch.Generator()
             generator.manual_seed(seed)
 
-        completion_shape = (batch_size, max_new_tokens)
+        completion_shape = (batch_size, completion_length)
         try:
             completion_ids = torch.randint(
                 low=0,
@@ -540,7 +588,8 @@ class SumiGenerationMixin(GenerationMixin):
             if was_training:
                 self.train()
 
-        return SumiGenerationOutput(sequences=generated_ids)
+        sequences = self._trim_at_eos(generated_ids, prompt_length) if trim_at_eos else generated_ids
+        return SumiGenerationOutput(sequences=sequences, canvas=generated_ids)
 
 
 __all__ = [