Preserve instructions in large code prompts

app.py

@@ -36,7 +36,7 @@ from transformers import (
     StoppingCriteria,
     StoppingCriteriaList,
 )
-from generation import merge_eos_token_ids
+from generation import head_tail_token_counts, merge_eos_token_ids
 from openai_compat import (
     indexed_tool_calls,
     normalize_messages,
@@ -62,12 +62,10 @@ MODEL = os.getenv(
 
 MAX_CONTEXT_TOKENS = int(os.getenv("MAX_CONTEXT_TOKENS", "16384"))
 MAX_NEW_TOKENS = int(os.getenv("MAX_NEW_TOKENS", "1024"))
+PRESERVED_PREFIX_TOKENS = int(os.getenv("PRESERVED_PREFIX_TOKENS", "4096"))
 TOOL_PROTOCOL_MARKER = "OPENAI TOOL CALL FORMAT — MANDATORY"
 
-# Keep the most recent tool result and user instructions if a client sends a
-# conversation longer than this Space can serve reliably.
 tokenizer = AutoTokenizer.from_pretrained(MODEL)
-tokenizer.truncation_side = "left"
 
 print(f"Loading {MODEL} on ZeroGPU during startup...", flush=True)
 model = AutoModelForCausalLM.from_pretrained(
@@ -171,9 +169,34 @@ def gerar(
         prompt,
         return_tensors="pt",
         add_special_tokens=False,
-        truncation=True,
-        max_length=max(1, MAX_CONTEXT_TOKENS - output_tokens),
-    ).to("cuda")
+        truncation=False,
+    )
+    input_budget = max(1, MAX_CONTEXT_TOKENS - output_tokens)
+    input_length = inputs["input_ids"].shape[1]
+    if input_length > input_budget:
+        head_tokens, tail_tokens = head_tail_token_counts(
+            input_length,
+            input_budget,
+            PRESERVED_PREFIX_TOKENS,
+        )
+        for key, value in inputs.items():
+            if (
+                isinstance(value, torch.Tensor)
+                and value.ndim == 2
+                and value.shape[1] == input_length
+            ):
+                parts = []
+                if head_tokens:
+                    parts.append(value[:, :head_tokens])
+                if tail_tokens:
+                    parts.append(value[:, -tail_tokens:])
+                inputs[key] = torch.cat(parts, dim=1)
+        print(
+            f"Prompt truncated head+tail: original={input_length} "
+            f"head={head_tokens} tail={tail_tokens}",
+            flush=True,
+        )
+    inputs = inputs.to("cuda")
 
     print(
         f"Generation started: input_tokens={inputs['input_ids'].shape[1]} "

generation.py

@@ -5,6 +5,28 @@ from __future__ import annotations
 from collections.abc import Iterable
 
 
+def head_tail_token_counts(
+    total_tokens: int,
+    token_budget: int,
+    preserved_prefix_tokens: int,
+) -> tuple[int, int]:
+    """Split an oversized prompt budget between its prefix and recent tail.
+
+    Keeping only the tail can erase system instructions, tool definitions, and
+    a task stated before a large code block. Keeping a bounded prefix plus the
+    largest possible tail retains both the operating contract and the newest
+    conversation state.
+    """
+    total = max(0, int(total_tokens))
+    budget = max(1, int(token_budget))
+    if total <= budget:
+        return total, 0
+
+    prefix = max(0, int(preserved_prefix_tokens))
+    head = min(prefix, budget - 1)
+    return head, budget - head
+
+
 def ensure_bos_token(prompt: str, bos_token: str | None) -> str:
     """Prefix the model's BOS token when the chat template omits it."""
     if not bos_token or prompt.startswith(bos_token):

test_generation.py

@@ -4,7 +4,11 @@ from __future__ import annotations
 
 import unittest
 
-from generation import ensure_bos_token, merge_eos_token_ids
+from generation import (
+    ensure_bos_token,
+    head_tail_token_counts,
+    merge_eos_token_ids,
+)
 
 
 class GenerationConfigTests(unittest.TestCase):
@@ -30,6 +34,18 @@ class GenerationConfigTests(unittest.TestCase):
     def test_returns_none_without_valid_ids(self) -> None:
         self.assertIsNone(merge_eos_token_ids(None, None))
 
+    def test_short_prompt_is_not_split(self) -> None:
+        self.assertEqual(head_tail_token_counts(100, 200, 50), (100, 0))
+
+    def test_large_prompt_preserves_prefix_and_maximizes_tail(self) -> None:
+        self.assertEqual(
+            head_tail_token_counts(30_000, 16_000, 4_096),
+            (4_096, 11_904),
+        )
+
+    def test_tiny_budget_still_preserves_latest_token(self) -> None:
+        self.assertEqual(head_tail_token_counts(100, 1, 4_096), (0, 1))
+
 
 if __name__ == "__main__":
     unittest.main()