diff --git a/src/nemo_safe_synthesizer/data_processing/assembler.py b/src/nemo_safe_synthesizer/data_processing/assembler.py index f0242e699..b9b9cb4f4 100644 --- a/src/nemo_safe_synthesizer/data_processing/assembler.py +++ b/src/nemo_safe_synthesizer/data_processing/assembler.py @@ -986,8 +986,11 @@ def _get_initial_prefill(self) -> dict[str, str]: if len(seen_groups[group_value]) < 3: seen_groups[group_value].append(record["text"]) - # Convert lists to joined strings - return {group: " " + "\n".join(samples) for group, samples in seen_groups.items()} + # 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()} def _apply_train_test_split(self, dataset: Dataset) -> None: """Override split logic to preserve record order and split along group boundaries.""" diff --git a/src/nemo_safe_synthesizer/generation/timeseries_backend.py b/src/nemo_safe_synthesizer/generation/timeseries_backend.py index a140fad63..7c96f5998 100644 --- a/src/nemo_safe_synthesizer/generation/timeseries_backend.py +++ b/src/nemo_safe_synthesizer/generation/timeseries_backend.py @@ -67,8 +67,14 @@ class GroupState: current_prefill: str """Current prefill string, updated as generation progresses to include recently generated records.""" - recent_records: list[dict] = field(default_factory=list) - """Sliding window of recently generated records used to build the next prompt context.""" + recent_records: list[ParsedRecord] = field(default_factory=list) + """Sliding window of recently generated records used to build the next prompt context. + + Holds ``ParsedRecord`` so the next prompt context can reuse the exact + bytes the model emitted (``.text``): re-serializing the parsed dicts + would drift the prompt into a JSON dialect the model was not trained + on (spacing, slash escaping, float formatting). + """ expected_records: int = 0 """Target record count, calculated from ``(stop_timestamp - start_timestamp) / interval_seconds``.""" @@ -513,12 +519,18 @@ def _check_chronological_for_group(self, batch: Batch, group_state: GroupState) if record.is_valid: record.invalidate(error) - def _update_group_state(self, group_state: GroupState, records: list[dict]) -> None: - """Update a group's state with new records. + def _update_group_state(self, group_state: GroupState, records: list[ParsedRecord]) -> None: + """Update a group's state with new valid records. + + The rebuilt prefill reuses each record's original ``text`` (the + bytes the model emitted, which match the training serialization) + and mirrors the initial prefill's shape from + ``SequentialExampleAssembler._get_initial_prefill``: a leading + space, then newline-terminated records. Args: group_state: The group state to update. - records: The new valid records. + records: The new valid records (``parsed`` is set on each). """ if not records: return @@ -529,12 +541,12 @@ def _update_group_state(self, group_state: GroupState, records: list[dict]) -> N # Update prefill tail = group_state.recent_records[-self._prefill_context_size :] - lines = [json.dumps(record, ensure_ascii=False) for record in tail] - group_state.current_prefill = "\n".join(lines) + "\n" + group_state.current_prefill = " " + "".join(f"{record.text}\n" for record in tail) # Update last timestamp - 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 is not None else None + timestamp_seconds = self._parse_timestamp_seconds(timestamp_value) if timestamp_seconds is not None: group_state.last_timestamp_seconds = timestamp_seconds @@ -621,9 +633,9 @@ def _process_group_result( if self.config.time_series.timestamp_interval_seconds is not None: self._check_chronological_for_group(batch, state) - batch_records = self._retain_single_valid_response(batch) - reached_stop = self._has_reached_stop_time(batch_records) - self._update_group_state(state, batch_records) + retained_records = self._retain_single_valid_response(batch) + reached_stop = self._has_reached_stop_time([r.parsed for r in retained_records if r.parsed is not None]) + self._update_group_state(state, retained_records) # Check if batch has high invalid fraction invalid_fraction = 1.0 - batch.valid_record_fraction @@ -852,7 +864,7 @@ def _generate_parallel_groups( return all_groups_succeeded - def _retain_single_valid_response(self, batch: Batch) -> list[dict]: + def _retain_single_valid_response(self, batch: Batch) -> list[ParsedRecord]: """Retain the response with the most valid records, discarding all others. For time-series sliding window generation, only one response can be used @@ -865,9 +877,11 @@ def _retain_single_valid_response(self, batch: Batch) -> list[dict]: batch: The batch to retain the response from. Returns: - List of valid records from the retained response. + The retained response's valid ``ParsedRecord`` objects, keeping + both the parsed dicts and the original emitted text (the latter + feeds the next prompt context). """ - final_records: list[dict] = [] + final_records: list[ParsedRecord] = [] # Find the index of the response with the most valid records. max_valid_idx = None @@ -886,7 +900,7 @@ def _retain_single_valid_response(self, batch: Batch) -> list[dict]: # error statistics without carrying stale text/token counts. response.records = [ParsedRecord(text="", error=trim_error)] else: - final_records.extend(response.valid_records) + final_records.extend(r for r in response.records if r.is_valid and r.parsed is not None) return final_records diff --git a/tests/data_processing/test_assembler.py b/tests/data_processing/test_assembler.py index 6648a3b1f..96b7fc608 100644 --- a/tests/data_processing/test_assembler.py +++ b/tests/data_processing/test_assembler.py @@ -809,21 +809,15 @@ def test_sequential_assembler_initial_prefill( assert "A" in prefill assert "B" in prefill - # Each prefill should contain up to 3 records as JSONL (newline-separated) - # Group A has 3 records, Group B has 2 records - # Filter out empty lines that may appear between records - prefill_a_lines = [line for line in prefill["A"].strip().split("\n") if line] - prefill_b_lines = [line for line in prefill["B"].strip().split("\n") if line] - - assert len(prefill_a_lines) == 3 # All 3 records from group A - assert len(prefill_b_lines) == 2 # Both records from group B - - # Verify the records contain expected values (order should be by time) - assert '"value": 10' in prefill_a_lines[0] or '"value":10' in prefill_a_lines[0] - assert '"value": 20' in prefill_a_lines[1] or '"value":20' in prefill_a_lines[1] - assert '"value": 30' in prefill_a_lines[2] or '"value":30' in prefill_a_lines[2] - assert '"value": 100' in prefill_b_lines[0] or '"value":100' in prefill_b_lines[0] - assert '"value": 200' in prefill_b_lines[1] or '"value":200' in prefill_b_lines[1] + # Pin the exact byte shape: a leading space, then newline-terminated + # training-dialect (pandas to_json) records with single newlines between + # them -- the same shape training examples use. Blank lines or Python + # json.dumps spacing here would put the generation prompt in a dialect + # the model never saw in training. + assert prefill["A"] == ( + ' {"group":"A","time":1,"value":10}\n{"group":"A","time":2,"value":20}\n{"group":"A","time":3,"value":30}\n' + ) + assert prefill["B"] == ' {"group":"B","time":1,"value":100}\n{"group":"B","time":2,"value":200}\n' def test_should_flush_example_boundary_conditions(): diff --git a/tests/generation/test_timeseries_backend.py b/tests/generation/test_timeseries_backend.py index f9a65415b..912c101ae 100644 --- a/tests/generation/test_timeseries_backend.py +++ b/tests/generation/test_timeseries_backend.py @@ -22,6 +22,7 @@ TimeSeriesParameters, TrainingHyperparams, ) +from nemo_safe_synthesizer.data_processing.record_utils import ParsedRecord from nemo_safe_synthesizer.defaults import DEFAULT_MAX_SEQ_LENGTH, PSEUDO_GROUP_COLUMN from nemo_safe_synthesizer.generation.processors import TimeSeriesDataProcessor from nemo_safe_synthesizer.generation.results import GenerationBatches @@ -387,16 +388,47 @@ def test_appends_records_and_updates_prefill(self, timeseries_base_params, times ) records = [ - {"timestamp": "2024-01-01 00:00:00", "value": 1}, - {"timestamp": "2024-01-01 01:00:00", "value": 2}, + ParsedRecord( + text='{"timestamp":"2024-01-01 00:00:00","value":1}', + parsed={"timestamp": "2024-01-01 00:00:00", "value": 1}, + ), + ParsedRecord( + text='{"timestamp":"2024-01-01 01:00:00","value":2}', + parsed={"timestamp": "2024-01-01 01:00:00", "value": 2}, + ), ] backend._update_group_state(state, records) assert len(state.recent_records) == 2 - assert '"timestamp"' in state.current_prefill + # The rebuilt prefill must reuse the records' emitted text verbatim + # (training dialect) and mirror the initial prefill's shape: leading + # space, newline-terminated records, single newlines between them. + assert state.current_prefill == ( + ' {"timestamp":"2024-01-01 00:00:00","value":1}\n{"timestamp":"2024-01-01 01:00:00","value":2}\n' + ) assert state.last_timestamp_seconds is not None + def test_updates_timestamp_for_empty_string_column_name( + self, timeseries_base_params, timeseries_model_metadata, mock_workdir + ): + """Test that an empty-string timestamp column name remains valid.""" + backend = create_timeseries_backend(timeseries_base_params, timeseries_model_metadata, mock_workdir) + backend._time_column = "" + state = GroupState( + group_id="test", + initial_prefill="", + current_prefill="", + expected_records=10, + last_timestamp_seconds=0, + ) + timestamp = "2024-01-01 01:00:00" + records = [ParsedRecord(text='{"":"2024-01-01 01:00:00"}', parsed={"": timestamp})] + + backend._update_group_state(state, records) + + assert state.last_timestamp_seconds == backend._parse_timestamp_seconds(timestamp) + class TestGetTimestampFromPrefill: """Tests for the _get_timestamp_from_prefill method."""