source/tests/unit/train/test_infinity_parser2_prepare_data.py

import argparse
import json
from pathlib import Path

import pytest
from datasets import Dataset

from scripts import infinity_parser2_prepare_data as prepare_module


class _FakeProcessor:
    chat_template = "{{ messages }}"

    def apply_chat_template(self, *args, **kwargs):
        return ""

    def decode(self, _tokens):
        return "token"


def _raw_dataset(candidate_ranks: list[int]) -> Dataset:
    return Dataset.from_list(
        [
            {
                "id": f"v1.12:{source_index:010d}",
                "conversations": [
                    {"role": "user", "content": "question"},
                    {"role": "assistant", "content": "answer"},
                ],
                "provenance": {
                    "source_line_index": source_index,
                    "candidate_rank": candidate_rank,
                    "sample_manifest_sha256": "a" * 64,
                    "generation_config_sha256": "b" * 64,
                },
            }
            for source_index, candidate_rank in enumerate(candidate_ranks)
        ]
    )


def _patch_pipeline(
    monkeypatch: pytest.MonkeyPatch,
    raw: Dataset,
    *,
    dropped_ranks: set[int] | None = None,
    fanout_by_rank: dict[int, int] | None = None,
) -> None:
    dropped_ranks = dropped_ranks or set()
    fanout_by_rank = fanout_by_rank or {}

    def fake_build(dataset, *, preserve_columns, **_kwargs):
        rows = [
            dataset[index]
            for index in range(len(dataset))
            if dataset[index]["_candidate_rank"] not in dropped_ranks
            for _ in range(fanout_by_rank.get(dataset[index]["_candidate_rank"], 1))
        ]
        result = Dataset.from_dict(
            {
                "input_ids": [[10, 11] for _ in rows],
                "loss_mask": [[0, 1] for _ in rows],
                "seq_len": [2 for _ in rows],
                **{
                    column: [row[column] for row in rows] for column in preserve_columns
                },
            }
        )
        result.set_format(type="torch")
        return result

    monkeypatch.setattr(
        prepare_module,
        "load_processor",
        lambda *args, **kwargs: _FakeProcessor(),
    )
    monkeypatch.setattr(
        prepare_module,
        "load_raw_dataset",
        lambda _path: (raw, None),
    )
    monkeypatch.setattr(
        prepare_module,
        "build_speculator_training_dataset",
        fake_build,
    )


def _args(tmp_path: Path, target_samples: int) -> argparse.Namespace:
    return argparse.Namespace(
        model="target-model",
        data=tmp_path / "success-pool.jsonl",
        output=tmp_path,
        ranked_target_samples=target_samples,
        seq_length=8,
        token_freq_train_ratio=0.67,
        render_endpoint="http://127.0.0.1:8000",
        # Both are store_true flags, so False is their parser default. Rendering
        # goes through the endpoint set above, which --local-render is the
        # alternative to.
        local_render=False,
        drop_clipped_rows=False,
        minimum_valid_tokens=None,
        num_preprocessing_workers=1,
        preprocessing_batch_size=2,
        trust_remote_code=False,
    )


def test_ranked_preprocessing_backfills_and_uses_train_prefix(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
):
    _patch_pipeline(
        monkeypatch,
        _raw_dataset([4, 0, 3, 1, 2]),
        dropped_ranks={1},
        fanout_by_rank={0: 2, 2: 3},
    )
    captured = {}

    def capture_frequency(*, dataset, output_path):
        captured["dataset"] = dataset.with_format(None)
        captured["output_path"] = output_path

    monkeypatch.setattr(
        prepare_module,
        "save_token_frequency_distribution",
        capture_frequency,
    )
    args = _args(tmp_path, 3)
    dataset = prepare_module.prepare_ranked_dataset(args)

    plain = dataset.with_format(None)
    assert plain["candidate_rank"] == [0, 0, 2, 2, 2, 3]
    assert captured["dataset"]["_candidate_rank"] == [0, 0, 2, 2, 2]
    assert captured["output_path"] == tmp_path / "token_freq.pt"
    metadata = json.loads(dataset.info.description)["ranked_preprocessing"]
    assert metadata == {
        "dataset_order": "candidate_rank_ascending_assistant_turn",
        "eligible_records": 4,
        "eligible_training_rows": 7,
        "generation_config_sha256": "b" * 64,
        "sample_manifest_sha256": "a" * 64,
        "selected_records": 3,
        "selected_training_rows": 6,
    }


def test_ranked_preprocessing_checks_the_complete_reserve(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
):
    _patch_pipeline(monkeypatch, _raw_dataset([0, 1, 2, 2]))
    monkeypatch.setattr(
        prepare_module,
        "save_token_frequency_distribution",
        lambda **kwargs: None,
    )

    with pytest.raises(ValueError, match="not unique"):
        prepare_module.prepare_ranked_dataset(_args(tmp_path, 2))


def test_ranked_preprocessing_requires_enough_survivors(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
):
    _patch_pipeline(monkeypatch, _raw_dataset([0, 1]))
    monkeypatch.setattr(
        prepare_module,
        "save_token_frequency_distribution",
        lambda **kwargs: None,
    )

    with pytest.raises(ValueError, match="below target"):
        prepare_module.prepare_ranked_dataset(_args(tmp_path, 3))


def test_selection_manifest_is_minimal_and_published_last(tmp_path: Path):
    output = tmp_path / "prepared"
    output.mkdir()
    dataset = Dataset.from_dict(
        {
            "input_ids": [[10, 11], [10, 12], [12, 13]],
            "loss_mask": [[0, 1], [0, 1], [0, 1]],
            "seq_len": [2, 2, 2],
            "id": [
                "v1.12:0000000003",
                "v1.12:0000000003",
                "v1.12:0000000001",
            ],
            "source_line_index": [3, 3, 1],
            "candidate_rank": [0, 0, 2],
        }
    )
    dataset.info.description = json.dumps(
        {
            "ranked_preprocessing": {
                "sample_manifest_sha256": "a" * 64,
                "generation_config_sha256": "b" * 64,
                "selected_records": 2,
                "selected_training_rows": 3,
            }
        }
    )

    staged, manifest = prepare_module.stage_selection(
        dataset,
        output=output,
        target_samples=2,
    )
    assert manifest == {
        "complete": True,
        "target_records": 2,
        "training_rows": 3,
        "sample_manifest_sha256": "a" * 64,
        "generation_config_sha256": "b" * 64,
        "selection": {
            "path": prepare_module.SELECTION_NAME,
            "records": 2,
        },
    }
    assert not (output / prepare_module.MANIFEST_NAME).exists()

    dataset.save_to_disk(output)
    prepare_module.publish_selection(
        output,
        staged_selection=staged,
        manifest=manifest,
    )
    rows = [
        json.loads(line)
        for line in (output / prepare_module.SELECTION_NAME).read_text().splitlines()
    ]
    assert [row["source_line_index"] for row in rows] == [1, 3]
    assert not list(output.glob("*.partial"))

    args = argparse.Namespace(output=output, ranked_target_samples=2)
    prepare_module.validate_existing_output(args)
    args.ranked_target_samples = 3
    with pytest.raises(ValueError, match="different settings"):
        prepare_module.validate_existing_output(args)