feat: Use Partial Prefix to Prompt Time Series Generation

docs/developer-guide/example-generation.md

@@ -259,12 +259,11 @@ these conditions are met:
 - Token budget exceeded -- the accumulated tokens reach the randomized budget
 - `max_sequences_per_example` reached
 
-### Initial prefill
+### Time-series group registry
 
-For each group, the first 3 training records are stored as `initial_prefill` --
-a dictionary mapping group IDs to seed text. During generation, the
-time-series backend uses these prefill strings to prime the model's context for
-each group.
+The artifact stores each typed group value in `timeseries_group_values`.
+Generation uses this registry to initialize one stream per training group.
+Each stream starts with a generated prefix as described below.
 
 ### Concrete example
 
@@ -446,23 +445,33 @@ When no BOS/EOS delimiters are found, behavior depends on
 !!! warning "Experimental"
     Time-series generation is experimental and its API may change.
 
-The `TimeseriesBackend` generates one group at a time using a sliding-window
-strategy:
+The `TimeseriesBackend` generates each group using a sliding history:
 
-1. Generation is seeded with an `initial_prefill` -- the first 3 records from
-   training data for that group.
-2. The model generates a continuation. Valid records are appended to the
-   group's context window.
-3. The updated context (most recent records) becomes the prefill for the
-   next prompt.
-4. This repeats until the target time range is covered or retries are
+1. The model receives up to the three most recently accepted records as
+   history and may generate one or more new records.
+2. The backend parses and validates each candidate, retains one response, and
+   appends its exact accepted record text to the history.
+3. This repeats until the target time range is covered or retries are
    exhausted.
 
+The first iteration is a special case because no generated history exists yet.
+The backend uses a prefix: an incomplete JSON record containing the group
+ID (for grouped data), configured start timestamp, and opening quote of the
+next field name. The model may complete that record and generate additional
+records. Before parsing, the backend prepends the prefix to each completion.
+After the first record is accepted, subsequent iterations use history only.
+
 Because each prompt sees the model's own prior output, sequential mode
 preserves temporal continuity -- timestamps, intervals, and trends carry
 forward naturally. The trade-off is that generation is inherently serial per
 group (though multiple groups are processed in parallel batches).
 
+Preprocessing places the group and timestamp columns first in the persisted
+schema and training JSONL. This lets the partial record match the field order
+seen during training. Older artifacts whose saved schema uses another order
+must be retrained; generation fails explicitly instead of silently changing
+prompting behavior.
+
 ---
 
 ## Grouped Generation: Validation Knobs

docs/user-guide/configuration.md

@@ -364,7 +364,6 @@ for the full field list.
 | `time_series.timestamp_format` | `null` | strftime format or `"elapsed_seconds"` | Required when `is_timeseries: true` |
 | `time_series.start_timestamp` | `null` | Override start timestamp for all groups (inferred from data if `null`) | Leave `null` to infer from data |
 | `time_series.stop_timestamp` | `null` | Override stop timestamp for all groups (inferred from data if `null`) | Leave `null` to infer from data |
-
 See [`TimeSeriesParameters`][nemo_safe_synthesizer.config.time_series.TimeSeriesParameters]
 for the full schema. For detailed descriptions and constraints, see the
 [Time Series README](https://github.com/NVIDIA-NeMo/Safe-Synthesizer/blob/main/src/nemo_safe_synthesizer/TIMESERIES_README.md).

src/nemo_safe_synthesizer/TIMESERIES_README.md

@@ -123,6 +123,11 @@ Time series preprocessing occurs during training data preparation in `src/nemo_s
    - If timestamps differ across groups, raises a `DataError`.
    - Sets `start_timestamp` and `stop_timestamp` in config based on validated values.
 
+8. Identity Column Ordering
+   - Places the group and timestamp columns before all generated value columns.
+   - This order is persisted in the schema and training JSONL so generation can begin with those known fields.
+   - The pseudo-group remains internal and is excluded from the persisted schema.
+
 ### Code Flow
 
 ```
@@ -135,6 +140,7 @@ HuggingFaceBackend._process_timeseries()
             │       ├── group length validation
             │       ├── interval consistency validation
             │       └── start/stop consistency validation
+            ├── order group and timestamp columns first
             └── Return (processed_df, updated_config)
 ```
 
@@ -157,10 +163,10 @@ The `SequentialExampleAssembler` in `src/nemo_safe_synthesizer/data_processing/a
 
 4. Pseudo-Group Handling: When no group column is specified, preprocessing adds a `__pseudo_group__` column so ungrouped time series is treated as a single group. This unifies the grouped and ungrouped code paths.
 
-5. Initial Prefill Extraction
-   - Dictionary mapping each group to its first 3 decoded samples (including pseudo-group for single sequences).
-   - Stored in `model_metadata.initial_prefill` for use during generation.
-   - Used by `TimeseriesBackend` to seed each group's context.
+5. Group Registry Extraction
+   - Ordered list of typed group values, including the pseudo-group for single sequences.
+   - Stored in `model_metadata.timeseries_group_values` so generation can recover the trained groups.
+   - Training records are not copied into production generation prompts.
 
 6. Train/Test Split
    - Split is done by group boundaries using `grouped_train_test_split`.
@@ -228,41 +234,48 @@ TimeseriesBackend(VllmBackend)
 ### Key Concepts
 
 - Time-Range Based Generation: The number of records generated is determined by `(stop_timestamp - start_timestamp) / interval_seconds`, not by a target count. The `config.generation.num_records` parameter is used only for progress tracking.
-- Sliding Window: Maintains a window of recent records (controlled by `_prefill_context_size`) included in each prompt for context continuity.
-- Groups from Training: Groups are the same as those seen during training (from `model_metadata.initial_prefill`).
+- Partial-Record Initialization: Every group starts with an incomplete JSON record containing its known group ID and start timestamp, plus the opening quote of the next field name. Including the complete training `,"` token makes the prefix tokenization identical to a full training record while leaving the field name and value for the model.
+- Training-Dialect Serialization: Constructed prefixes and rolling records use the same compact JSON representation as training, including escaped slashes and schema field order.
+- Training-Compatible Token Boundary: Generation explicitly reproduces the prompt BOS/EOS settings and sequence BOS token used by training. The first JSON byte follows the sequence BOS directly, without added whitespace.
+- Sliding Window: Maintains the three most recent generated records for context continuity.
+- Rolling Context Budget: Each generation batch clamps `max_tokens` against its longest current rolling prompt so prompt plus completion cannot exceed the
+model context.
+- Groups from Training: Groups are the same as those seen during training (from `model_metadata.timeseries_group_values`).
 
 ### Sliding Window Approach
 
-1. Prefill Initialization: Start with initial prefill from training data (first 3 records per group).
-2. Batch Generation: Generate multiple samples (default 5) per prompt for each active group.
-3. Response Selection: Keep the response with the most valid records per group.
-4. Context Update: Update sliding window with new valid records.
-5. Repeat: Continue until stop timestamp is reached for each group.
+1. Partial Prefix Initialization: Start an incomplete first record with the group and start timestamp fields.
+2. Token-Prompt Assembly: Reproduce the training prompt and sequence special-token boundary, then append the prefix or history bytes.
+3. Batch Generation: Clamp completion length to the remaining context and generate multiple candidate suffixes (default 5) per prompt for each active group.
+4. Record Reconstruction: During the first iteration, prepend the JSON-only prefix before parsing each candidate.
+5. Response Selection: Keep the response with the most valid records per group.
+6. History Update: Switch from the prefix to a sliding history containing exact accepted record text.
+7. Repeat: Continue until stop timestamp is reached for each group.
 
 ### Key Parameters
 
 | Parameter | Value | Description |
 |-----------|-------|-------------|
 | `_samples_per_prompt` | 5 | Number of samples generated per prompt |
 | `_max_prompts_per_batch` | 100 | Max prompts per batch in parallel generation |
-| `_prefill_context_size` | 3 | Number of recent records in sliding window |
+| `_history_window_size` | 3 | Internal number of recent records in the history window |
 
 ### Parallel Group Generation Flow
 
 All time series use parallel group generation (single-sequence is just 1 group):
 
 ```
-1. Initialize GroupState for each group with prefill from training
+1. Initialize GroupState for each group with a partial first record
 2. Compute expected records per group: (stop - start) / interval + 1
 3. While groups remain pending or active:
    a. Fill active slots with pending groups (up to max_groups_per_batch)
-   b. Build prompts for all active groups using current prefill
+   b. Build prompts for all active groups using the initial prefix or generated history
    c. Generate completions for all prompts in single LLM batch call
    d. Process LLM outputs into per-group Batch objects
    e. For each group:
       - Validate chronological order against group's last timestamp
       - Retain response with most valid records (discard others)
-      - Update group state (prefill, last_timestamp)
+      - Update group state (history, last_timestamp)
       - Check if stop timestamp reached (marks group complete)
       - Track low valid fraction; fail group after max retries
    f. Remove completed/failed groups from active list
@@ -277,10 +290,8 @@ Each group maintains independent state:
 ```python
 @dataclass
 class GroupState:
-    group_id: str
-    initial_prefill: str           # Original prefill (first few records)
-    current_prefill: str           # Updated as generation progresses
-    recent_records: list[dict]     # Sliding window context
+    group_id: TimeSeriesGroupValue
+    prompt_state: RecordPromptState  # Prefix/history prompt state
     expected_records: int          # Based on (stop - start) / interval
     last_timestamp_seconds: int | None
     low_valid_fraction_count: int  # Counter for consecutive bad batches
@@ -298,8 +309,11 @@ Per-Group Stopping:
 
 Global Stopping:
 - Natural completion: All groups processed (pending and active lists empty).
-- No records: Too many consecutive batches with no valid records globally.
-- Target reached: Target number of records reached (for progress tracking).
+- Per-group retry state is authoritative. A zero-valid result for one group does
+  not trigger the generic global no-record stop while other parallel groups can
+  still progress.
+- `num_records` remains a progress target and does not stop time-range-based
+  generation.
 
 ### Progress Checkpoints
 
@@ -559,5 +573,5 @@ The `timestamp_validation_mode` parameter controls how timestamps are handled du
 | `generation/processors.py` | `TimeSeriesDataProcessor` class |
 | `generation/timeseries_backend.py` | `TimeseriesBackend` generation class |
 | `evaluation/components/autocorrelation_similarity.py` | ACF evaluation metric |
-| `llm/metadata.py` | `initial_prefill` field for time series |
+| `llm/metadata.py` | Typed time-series group registry and source column order |
 | `defaults.py` | `PSEUDO_GROUP_COLUMN` constant |

src/nemo_safe_synthesizer/data_processing/assembler.py

@@ -41,9 +41,10 @@
     ParameterError,
 )
 from ..holdout.holdout import grouped_train_test_split, naive_train_test_split
-from ..llm.metadata import ModelMetadata
+from ..llm.metadata import ModelMetadata, TimeSeriesGroupValue
 from ..observability import get_logger
 from .budget import NUM_SPECIAL_TOKENS, compute_max_new_tokens
+from .prompt_tokens import encode_prompt_token_ids, wrap_sequence_token_ids
 
 logger = get_logger(__name__)
 
@@ -170,12 +171,11 @@ def __init__(
         self.tokenizer = tokenizer
         self.metadata = metadata
 
-        self.input_ids = self.tokenizer.encode(prompt, add_special_tokens=False)
-
-        if self.metadata.prompt_config.add_bos_token_to_prompt:
-            self.input_ids = [self.metadata.prompt_config.bos_token_id] + self.input_ids
-        if self.metadata.prompt_config.add_eos_token_to_prompt:
-            self.input_ids = self.input_ids + [self.metadata.prompt_config.eos_token_id]
+        self.input_ids = encode_prompt_token_ids(
+            prompt,
+            tokenizer=self.tokenizer,
+            prompt_config=self.metadata.prompt_config,
+        )
 
         # We use -100 to ignore the prompt tokens when calculating the loss.
         self.labels = [-100] * len(self.input_ids)
@@ -199,7 +199,11 @@ def add_sequence(self, seq: dict[str, list[int]], add_special_tokens: bool = Tru
             GenerationError: If the number of tokens in the example exceeds the context length.
         """
         input_ids = (
-            [self.metadata.prompt_config.bos_token_id] + seq["input_ids"] + [self.metadata.prompt_config.eos_token_id]
+            wrap_sequence_token_ids(
+                seq["input_ids"],
+                prompt_config=self.metadata.prompt_config,
+                include_eos=True,
+            )
             if add_special_tokens
             else seq["input_ids"]
         )
@@ -742,10 +746,9 @@ class SequentialExampleAssembler(TabularDataExampleAssembler):
         - Pseudo-Group Handling: When no group column is specified, preprocessing
           adds a PSEUDO_GROUP_COLUMN so ungrouped time series is treated as a single
           group. This unifies the grouped and ungrouped code paths.
-        - Initial Prefill: For each group, the first 3 records are stored in
-          `model_metadata.initial_prefill` as a dict mapping group_id -> prefill string.
-          This is used by TimeseriesBackend during generation to seed each group's
-          context.
+        - Group Registry: Typed group values are stored in model metadata so
+          generation can initialize one stream per training group without
+          copying training records into prompts.
 
     Processing Flow:
         1. Initialization:
@@ -957,40 +960,29 @@ def num_groups_validation(self) -> int:
         """Number of unique groups in the validation split."""
         return self._count_groups(self.validation_dataset)
 
-    def _get_initial_prefill(self) -> dict[str, str]:
-        """Return sample records from the training dataset for each group.
-
-        Returns a dictionary mapping each group_id to its prefill string (first 3 samples per group).
-        For pseudo-grouped single sequences, returns a dict with one key.
-
-        Note:
-            This method assumes the dataset is already ordered by the timestamp/order column
-            within each group (done by _prepare_dataset_for_training).
-
-        Returns:
-            Dict mapping group values to prefill strings (first 3 samples per group).
-        """
+    def _get_timeseries_group_values(self) -> list[TimeSeriesGroupValue]:
+        """Return distinct typed group values in training-dataset order."""
         if self.training_dataset is None or len(self.training_dataset) == 0:
-            return {}
+            return []
 
         if self.group_by_column not in self.training_dataset.column_names:
-            return {}
-
-        # Get first 3 samples from each group, return as dict
-        # Use the preserved 'text' column directly to avoid encode/decode roundtrip issues
-        seen_groups: dict[str, list[str]] = {}
-        for record in self.training_dataset:
-            group_value = record[self.group_by_column]
-            if group_value not in seen_groups:
-                seen_groups[group_value] = []
-            if len(seen_groups[group_value]) < 3:
-                seen_groups[group_value].append(record["text"])
-
-        # Each sample line is already newline-terminated (see
-        # _convert_records_to_jsonl), so concatenate directly: joining with
-        # "\n" would insert blank lines between records, a shape that never
-        # occurs in training examples.
-        return {group: " " + "".join(samples) for group, samples in seen_groups.items()}
+            return []
+
+        group_values: list[TimeSeriesGroupValue] = []
+        seen_values: set[TimeSeriesGroupValue] = set()
+        for raw_value in self.training_dataset[self.group_by_column]:
+            value = raw_value.item() if isinstance(raw_value, np.generic) else raw_value
+            if not isinstance(value, (str, int, float, bool)) or (
+                isinstance(value, float) and not math.isfinite(value)
+            ):
+                raise ParameterError(
+                    f"Time-series group value {value!r} cannot be represented as a finite JSON scalar."
+                )
+            if value in seen_values:
+                continue
+            seen_values.add(value)
+            group_values.append(value)
+        return group_values
 
     def _apply_train_test_split(self, dataset: Dataset) -> None:
         """Override split logic to preserve record order and split along group boundaries."""

src/nemo_safe_synthesizer/data_processing/prompt_tokens.py

@@ -0,0 +1,53 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Shared prompt and sequence token-boundary assembly."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Protocol
+
+if TYPE_CHECKING:
+    from ..llm.metadata import LLMPromptConfig
+
+__all__ = [
+    "EncodeOnlyTokenizer",
+    "encode_prompt_token_ids",
+    "wrap_sequence_token_ids",
+]
+
+
+class EncodeOnlyTokenizer(Protocol):
+    """Tokenizer interface required for prompt construction."""
+
+    def encode(self, text: str, *, add_special_tokens: bool) -> list[int]:
+        """Encode text while explicitly controlling tokenizer special tokens."""
+
+
+def encode_prompt_token_ids(
+    prompt: str,
+    *,
+    tokenizer: EncodeOnlyTokenizer,
+    prompt_config: LLMPromptConfig,
+) -> list[int]:
+    """Encode a prompt with its configured BOS and EOS boundaries."""
+    prompt_ids = list(tokenizer.encode(prompt, add_special_tokens=False))
+    if prompt_config.add_bos_token_to_prompt:
+        prompt_ids.insert(0, prompt_config.bos_token_id)
+    if prompt_config.add_eos_token_to_prompt:
+        prompt_ids.append(prompt_config.eos_token_id)
+    return prompt_ids
+
+
+def wrap_sequence_token_ids(
+    token_ids: Sequence[int],
+    *,
+    prompt_config: LLMPromptConfig,
+    include_eos: bool,
+) -> list[int]:
+    """Return sequence IDs with a leading BOS and an optional trailing EOS."""
+    sequence_ids = [prompt_config.bos_token_id, *token_ids]
+    if include_eos:
+        sequence_ids.append(prompt_config.eos_token_id)
+    return sequence_ids