fix: add a mechanism for allowing more max tokens

Signed-off-by: Matt Kornfield <mkornfield@nvidia.com>

Author: mckornfield

src/nemo_safe_synthesizer/config/generate.py

@@ -297,6 +297,28 @@ class GenerateParameters(Parameters, BaseModel):
         ),
     ] = 0
 
+    max_tokens_multiplier: Annotated[
+        float,
+        ValueValidator(value_func=lambda v: v > 0),
+        Field(
+            title="max_tokens_multiplier",
+            description=(
+                "Margin applied to the longest tokenized training example "
+                "(``max_tokens_per_example``) when sizing the per-sample generation "
+                "budget (``SamplingParams.max_tokens``). The effective budget is "
+                "``int(max_tokens_per_example * max_tokens_multiplier)``, still clamped "
+                "to the remaining context window. The default adds only a small jitter "
+                "margin, which suffices for most tables but is too tight for long, "
+                "unbounded free-text columns: a model that over-generates slightly past "
+                "the longest training example truncates mid-JSON and produces no "
+                "parseable record (``length``-dominated finish reasons, zero valid "
+                "records). Raise this (e.g. to 1.5-2.0) to give the model room to emit "
+                "the closing tokens on such datasets; there is no benefit to values that "
+                "push the budget past the context window. Must be > 0."
+            ),
+        ),
+    ] = 1.2  # mirrors llm.metadata.GENERATION_MAX_TOKENS_SAFETY_MULTIPLIER (kept a literal to avoid a config->llm import cycle)
+
     structured_generation: StructuredGenerationParameters = Field(
         description="Structured generation parameters controlling schema-constrained output.",
         default_factory=StructuredGenerationParameters,

src/nemo_safe_synthesizer/generation/timeseries_backend.py

@@ -921,7 +921,10 @@ def generate(
             top_p=self.config.generation.top_p,
             top_k=FIXED_RUNTIME_GENERATE_ARGS["top_k"],
             min_p=FIXED_RUNTIME_GENERATE_ARGS["min_p"],
-            max_tokens=self.model_metadata.generation_max_tokens_for(self._get_prompt_token_count()),
+            max_tokens=self.model_metadata.generation_max_tokens_for(
+                self._get_prompt_token_count(),
+                multiplier=self.config.generation.max_tokens_multiplier,
+            ),
             skip_special_tokens=True,
             include_stop_str_in_output=False,
             ignore_eos=False,

src/nemo_safe_synthesizer/generation/vllm_backend.py

@@ -746,7 +746,10 @@ def _run_generation(self, data_actions_fn: utils.DataActionsFn | None) -> None:
             top_p=self.config.generation.top_p,
             top_k=FIXED_RUNTIME_GENERATE_ARGS["top_k"],
             min_p=FIXED_RUNTIME_GENERATE_ARGS["min_p"],
-            max_tokens=self.model_metadata.generation_max_tokens_for(self._get_prompt_token_count()),
+            max_tokens=self.model_metadata.generation_max_tokens_for(
+                self._get_prompt_token_count(),
+                multiplier=self.config.generation.max_tokens_multiplier,
+            ),
             skip_special_tokens=not need_special_token_outputs,
             include_stop_str_in_output=need_special_token_outputs,
             ignore_eos=False,

src/nemo_safe_synthesizer/llm/metadata.py

@@ -462,13 +462,13 @@ def max_seq_length(self) -> int:
             rsf = self.rope_scaling.factor
         return int((self.base_max_seq_length or DEFAULT_MAX_SEQ_LENGTH) * rsf)
 
-    def generation_max_tokens_for(self, prompt_len: int) -> int:
+    def generation_max_tokens_for(self, prompt_len: int, multiplier: float | None = None) -> int:
         """Per-sample ``max_tokens`` ceiling, prompt-aware.
 
         Returns the smaller of:
 
-        1. ``int(max_tokens_per_example * GENERATION_MAX_TOKENS_SAFETY_MULTIPLIER)``
-           when the assembler stat is populated, else ``max_seq_length``.
+        1. ``int(max_tokens_per_example * multiplier)`` when the assembler stat
+           is populated, else ``max_seq_length``.
         2. ``max_seq_length - prompt_len`` -- vLLM raises when
            ``len(prompt) + max_tokens > max_model_len``
            (`vllm#33418 <https://github.com/vllm-project/vllm/issues/33418>`_).
@@ -480,15 +480,29 @@ def generation_max_tokens_for(self, prompt_len: int) -> int:
         clamp is a defensive belt for legacy adapters where the assembler
         stat is missing and for prompts longer than those seen in training.
 
+        The default ``multiplier`` (``GENERATION_MAX_TOKENS_SAFETY_MULTIPLIER``)
+        adds only a small jitter margin, which is enough for most tables but
+        too tight for long, unbounded free-text columns: a model that
+        over-generates slightly past the longest training example truncates
+        mid-JSON and yields no parseable record. Callers wire the user-facing
+        ``generation.max_tokens_multiplier`` knob through here to widen the
+        budget (bounded by the context window) for such datasets.
+
         Args:
             prompt_len: Tokenized length of the prompt this sample will
                 run against. Pass ``0`` to disable the prompt clamp.
+            multiplier: Margin applied to ``max_tokens_per_example``. Defaults
+                to ``GENERATION_MAX_TOKENS_SAFETY_MULTIPLIER`` when ``None`` so
+                non-generation callers (e.g. the training eval callback) keep
+                the legacy sizing.
 
         Returns:
             Non-negative ``max_tokens`` value safe to feed to ``SamplingParams``.
         """
+        if multiplier is None:
+            multiplier = GENERATION_MAX_TOKENS_SAFETY_MULTIPLIER
         if self.max_tokens_per_example and self.max_tokens_per_example > 0:
-            sized = int(self.max_tokens_per_example * GENERATION_MAX_TOKENS_SAFETY_MULTIPLIER)
+            sized = int(self.max_tokens_per_example * multiplier)
         else:
             sized = self.max_seq_length
         return max(0, min(sized, self.max_seq_length - prompt_len))

tests/config/test_generate.py

@@ -209,3 +209,22 @@ def test_legacy_keys_with_no_nested_dict_are_migrated(self) -> None:
         assert params.structured_generation.backend == "xgrammar"
         assert params.structured_generation.schema_method == "structural_tag"
         assert params.structured_generation.use_single_sequence is True
+
+
+@pytest.mark.unit
+class TestMaxTokensMultiplier:
+    def test_default_matches_metadata_constant(self) -> None:
+        """The config default mirrors the metadata safety-margin constant."""
+        from nemo_safe_synthesizer.llm.metadata import GENERATION_MAX_TOKENS_SAFETY_MULTIPLIER
+
+        assert GenerateParameters().max_tokens_multiplier == GENERATION_MAX_TOKENS_SAFETY_MULTIPLIER
+
+    def test_accepts_widened_value(self) -> None:
+        """Users can widen the budget for long free-text datasets."""
+        assert GenerateParameters(max_tokens_multiplier=1.8).max_tokens_multiplier == 1.8
+
+    @pytest.mark.parametrize("value", [0, -0.5])
+    def test_rejects_non_positive(self, value: float) -> None:
+        """Non-positive multipliers are rejected by the validator."""
+        with pytest.raises(ValidationError):
+            GenerateParameters(max_tokens_multiplier=value)

tests/generation/test_vllm_backend.py

@@ -1048,8 +1048,10 @@ def capture_and_stop(**kwargs):
         assert captured["max_tokens"] == 4_200
         # Engine is not initialized in this plumbing test, so the cached
         # prompt-token count falls back to 0; the helper is still called
-        # exactly once with that value.
-        mock_model_metadata.generation_max_tokens_for.assert_called_once_with(0)
+        # exactly once with that value plus the configured budget multiplier.
+        mock_model_metadata.generation_max_tokens_for.assert_called_once_with(
+            0, multiplier=base_params.generation.max_tokens_multiplier
+        )
 
     def test_passes_cached_prompt_token_count_when_engine_initialized(
         self, base_params, mock_model_metadata, mock_schema, mock_workdir
@@ -1079,7 +1081,9 @@ def capture_and_stop(**kwargs):
             backend.generate()
 
         assert captured["max_tokens"] == 4_096
-        mock_model_metadata.generation_max_tokens_for.assert_called_once_with(37)
+        mock_model_metadata.generation_max_tokens_for.assert_called_once_with(
+            37, multiplier=base_params.generation.max_tokens_multiplier
+        )
         # Cached: a second access does not retokenize.
         assert backend._get_prompt_token_count() == 37
         fake_tokenizer.encode.assert_called_once_with(backend.prompt)

tests/llm/test_metadata.py

@@ -689,6 +689,32 @@ def test_metadata_generation_max_tokens_for_never_returns_negative(self, sample_
         # Beyond-context prompts still produce a safe value vLLM can accept.
         assert sample_model_metadata.generation_max_tokens_for(sample_model_metadata.max_seq_length + 64) == 0
 
+    def test_metadata_generation_max_tokens_for_none_multiplier_uses_default(self, sample_model_metadata):
+        """``multiplier=None`` reproduces the legacy safety-margin sizing."""
+        sample_model_metadata.max_tokens_per_example = 1000
+        expected = int(1000 * GENERATION_MAX_TOKENS_SAFETY_MULTIPLIER)
+        assert sample_model_metadata.generation_max_tokens_for(10, multiplier=None) == expected
+        # Explicitly passing the default constant matches the None path.
+        assert (
+            sample_model_metadata.generation_max_tokens_for(10, multiplier=GENERATION_MAX_TOKENS_SAFETY_MULTIPLIER)
+            == expected
+        )
+
+    def test_metadata_generation_max_tokens_for_custom_multiplier_widens_budget(self, sample_model_metadata):
+        """A larger multiplier widens the stat-derived budget (the long-text fix)."""
+        sample_model_metadata.max_tokens_per_example = 1000
+        # 1000 * 1.8 = 1800, still within the 2048 window given a small prompt.
+        assert sample_model_metadata.generation_max_tokens_for(10, multiplier=1.8) == 1800
+
+    def test_metadata_generation_max_tokens_for_custom_multiplier_still_clamped_to_window(self, sample_model_metadata):
+        """The prompt/window clamp still binds even with an aggressive multiplier."""
+        assert sample_model_metadata.max_seq_length == 2048
+        sample_model_metadata.max_tokens_per_example = 1500  # * 4.0 = 6000, far past the window
+        prompt_len = 48
+        assert sample_model_metadata.generation_max_tokens_for(prompt_len, multiplier=4.0) == (
+            sample_model_metadata.max_seq_length - prompt_len
+        )
+
     def test_metadata_max_tokens_per_example_round_trips_through_metadata_json(self, sample_model_metadata):
         """``max_tokens_per_example`` persists through save → load."""
         sample_model_metadata.max_tokens_per_example = 1500