fix(generation): keep timeseries prompt context in the training serialization - #684
Conversation
WalkthroughThe change preserves emitted record formatting when rebuilding timeseries prefills. It also updates sliding-window state to store ChangesPrefill record handling
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
| last_record = records[-1] | ||
| timestamp_seconds = self._parse_timestamp_seconds(last_record.get(self._time_column)) | ||
| last_parsed = records[-1].parsed or {} | ||
| timestamp_value = last_parsed.get(self._time_column) if self._time_column else None |
There was a problem hiding this comment.
[P2] Preserve empty-string timestamp column names
timestamp_column is typed as str | None, and the config/preflight code treats only None as absent, so "" remains a valid column name when the input DataFrame has such a column. The previous last_record.get(self._time_column) updated this case, but this truthiness check skips it and leaves last_timestamp_seconds stale. On the next batch, chronological validation compares against the old timestamp and rejects otherwise contiguous records. Please check self._time_column is not None (or access the key directly) instead.
There was a problem hiding this comment.
addressed! as an aside, should we disallow "" as a col name? that seems confusing in cases
There was a problem hiding this comment.
yes, we probably should!
|
Thank you for catching this! |
…lization The timeseries sliding-window prompt context diverged from the serialization the model is trained on, in three ways: - Dialect: _update_group_state rebuilt the per-group prefill with json.dumps (spaces after ':' and ',', unescaped '/', full-repr floats), while training examples serialize via pandas to_json (no spaces, escaped slashes, double_precision=10). From the second batch of every group onward, the prompt continued in a dialect the model never saw in training. - Blank lines: _get_initial_prefill joined already newline-terminated record lines with "\n", inserting blank lines between records; a shape absent from training examples. - Shape drift: the initial prefill (leading space, "\n\n" separators) and the rebuilt prefill (no leading space, "\n" separators) did not match each other, so batch-1 and batch-2+ prompts differed beyond record content. Fix: rebuild the prefill from each ParsedRecord's original emitted text (which is already in the training dialect) instead of re-serializing parsed dicts; re-serializing via records_to_jsonl would upcast ints to floats when a null shares the 3-record window and re-round float precision. GroupState.recent_records now carries ParsedRecord so the emitted bytes survive the sliding window, and both prefill producers emit the same shape: a leading space followed by newline-terminated records. Tests pin the exact prefill bytes on both producers; the previous assertions filtered out the blank lines rather than pinning them. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
0319c47 to
4de4816
Compare
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
Greptile SummaryAligns time-series prompt context with training serialization.
Confidence Score: 5/5The PR appears safe to merge, with the changed training and generation serialization paths remaining aligned. The assembler guarantees newline-terminated record text, generated ParsedRecord text is framed correctly when rebuilding the window, and all consumers of the changed retained-record type explicitly receive the expected ParsedRecord or parsed-dictionary representation. Important Files Changed
|
3 participants
Summary
The timeseries sliding-window prompt context diverged from the serialization the model is trained on, in three ways (all reproduced at runtime against
main):Serialization dialect.
TimeseriesBackend._update_group_staterebuilt each group's prefill withjson.dumps(record, ensure_ascii=False)— spaces after:and,, unescaped/, full-repr floats. Training examples serialize via pandasto_json(orient="records", lines=True)— no spaces,\/escaping,double_precision=10. From the second batch of every group onward, the prompt continued in a JSON dialect the model never saw in training:Blank lines in the initial prefill.
SequentialExampleAssembler._get_initial_prefilljoined already newline-terminated record lines with"\n", producing\n\nbetween records — training examples have single newlines.Shape drift between batches. The initial prefill (leading space,
\n\nseparators) and the rebuilt prefill (no leading space,\nseparators) did not match each other, so batch-1 and batch-2+ prompts differed beyond record content.Fix
ParsedRecord.text— the exact bytes the model emitted, which are already in the training dialect. Re-serializing the parsed dicts throughrecords_to_jsonlwas rejected: pandas serializes column-wise, so a null sharing the 3-record window upcasts another record's1to1.0, anddouble_precision=10re-rounds accepted model output.GroupState.recent_recordsnow carriesParsedRecord(wasdict) so the emitted bytes survive the sliding window;_retain_single_valid_responsereturns the retainedParsedRecords." {r1}\n{r2}\n{r3}\n"). The leading space is kept because the initial-prefill shape is persisted in trained-model metadata; the rebuilt prefill is pure runtime state.Besides output quality, byte-alignment of prompt context with the trained output format is a prerequisite for prompt-lookup speculative decoding (drafts only match if the prompt text token-matches what the model emits) and for stable per-group prefix caching.
Tests
tests/data_processing/test_assembler.py::test_sequential_assembler_initial_prefillnow pins the exact prefill bytes (previously it filtered out the blank lines rather than pinning them).tests/generation/test_timeseries_backend.py::TestUpdateGroupStatenow pins the exact rebuilt-prefill bytes and the verbatim reuse of emitted text.mise run format/typecheckclean; affected suites (52 tests) pass.Compatibility
The prefill string is consumed only by
_format_promptand_get_timestamp_from_prefill(regex-based, shape-agnostic) — both dialects parse identically, so no downstream consumer changes. Already-trained model artifacts keep their persisted initial prefill; the D2 fix affects newly trained models only.Summary by CodeRabbit
Bug Fixes
Tests