speculators.data_generation.preprocessing

def build_speculator_training_dataset(
    dataset: HFDataset,
    processor: ProcessorLike,
    max_length: int = 2048,
    num_proc: int = 8,
    *,
    render_endpoint: str | None = None,
    minimum_valid_tokens: int | None = None,
) -> HFDataset:
    """Build a speculator training dataset with render-boundary loss masks.

    Both accepted representations contain responses produced by the target
    model. Natural-language conversations are tokenized by the vLLM ``/render``
    endpoint and masked at each assistant-turn boundary, fanning out to one row
    per assistant turn. Rendering only converts representation; it does not
    generate responses or make arbitrary data on-policy. Speculator-format rows
    already carry ``input_ids`` and ``loss_mask`` and pass straight through.

    Args:
        dataset: On-policy natural-language conversations, or speculator-format
            rows containing ``input_ids`` and ``loss_mask``.
        processor: Processor, used to detect multimodal inputs and to decode.
        max_length: Maximum sequence length.
        num_proc: Number of worker processes; each renders concurrently.
        render_endpoint: Base URL of a vLLM server. Required unless the dataset
            is already in speculator format.
        minimum_valid_tokens: Minimum supervised tokens for a row to be kept.
    """
    original_cols = dataset.column_names
    # These rows carry their supervision mask, so _preprocess_batch passes them
    # through without rendering or boundary derivation.
    pretokenized = {"input_ids", "loss_mask"} <= set(original_cols)
    # Multimodal rows keep their `messages` so the images survive to hidden-state
    # extraction. Compute once here rather than pickling the heavyweight processor
    # into every map worker just to recheck it.
    is_multimodal = isinstance(processor, ProcessorMixin)

    if pretokenized:
        log.info("Speculator-format rows: using their loss mask, skipping render")
    elif render_endpoint is None:
        raise ValueError(
            "render_endpoint is required to convert natural-language "
            "conversations to speculator training rows. Pass --render-endpoint "
            "pointing at the target model's vLLM server."
        )
    else:
        log.info("Deriving loss masks from vLLM render boundaries")

    # Avoid CPU contention for MM processing:
    # https://github.com/vllm-project/vllm/pull/31879
    with set_default_torch_num_threads() if is_multimodal else nullcontext():
        dataset = dataset.map(
            lambda examples: _preprocess_batch(
                examples,
                is_multimodal,
                render_endpoint,
                max_length,
                minimum_valid_tokens,
            ),
            batched=True,
            num_proc=num_proc,
            batch_size=1000,
            remove_columns=original_cols,
            keep_in_memory=True,  # skip caching
        )

    dataset.set_format(type="torch")
    return dataset

def default_preprocessing_workers(cpus: int | None = None) -> int:
    """Choose preprocessing workers within the shared render CPU budget."""
    if cpus is None:
        cpus = usable_cpu_count()
    return max(
        1,
        min(
            MAX_PREPROCESSING_WORKERS,
            int(cpus * CPU_BUDGET_FRACTION) // EFFECTIVE_CPUS_PER_PREPROCESSING_WORKER,
        ),
    )

def load_and_preprocess_dataset(
    target_model_path: str,
    train_data_paths: list[str],
    *,
    seq_length: int,
    build_dataset_num_proc: int = 8,
    seed: int = 0,
    max_samples: int | None = None,
    token_freq_path: Path | str = "./token_freq.pt",  # noqa: S107
    render_endpoint: str | None = None,
    minimum_valid_tokens: int | None = None,
    allow_empty_output: bool = False,
    trust_remote_code: bool = False,
) -> tuple[HFDataset, ProcessorLike]:
    """Load, tokenize, and preprocess a dataset for speculator training.

    Natural-language conversations containing target-model responses are
    tokenized by a vLLM ``/render`` endpoint and masked at each assistant-turn
    boundary. Speculator-format rows pass straight through. Rendering converts
    representation; it does not generate or validate response provenance.
    Caching is handled automatically by HuggingFace datasets.

    Args:
        target_model_path: HuggingFace model ID or local path
        train_data_path: Dataset name or path to JSON/JSONL file
        seq_length: Maximum sequence length
        build_dataset_num_proc: Number of processes for dataset building
        seed: Random seed for shuffling
        max_samples: Optional limit on number of samples
        token_freq_path: Path to save token frequency distribution
        cache_dir: Directory to cache HuggingFace datasets (optional)
        render_endpoint: Base URL of a running vLLM server (e.g.
            ``http://localhost:8000``) used to render conversations. Required
            unless every dataset is already in speculator format.
        minimum_valid_tokens: Number of tokens to consider for a valid sample
        allow_empty_output: If True, allow returning an empty dataset instead of
                          raising when no samples survive preprocessing.
        trust_remote_code: If True, allows executing code from HF Hub.

    Returns:
        Tuple of (preprocessed_dataset, processor)
    """
    if minimum_valid_tokens is not None and minimum_valid_tokens < 0:
        raise ValueError("minimum_valid_tokens must be >= 0")
    log.section("Starting dataset preprocessing")
    if minimum_valid_tokens is not None:
        log.info(
            f"Filtering samples with fewer than {minimum_valid_tokens} valid tokens"
        )

    log.subsection("Loading processor")
    processor = load_processor(target_model_path, trust_remote_code=trust_remote_code)

    processor_has_chat_template = (
        hasattr(processor, "apply_chat_template")
        and getattr(processor, "chat_template", None) is not None
    )

    if render_endpoint is not None:
        log.info(f"Rendering conversations via vLLM endpoint: {render_endpoint}")

    processed_datasets = []
    for train_data_path in train_data_paths:
        log.subsection(f"Processing {train_data_path}")
        raw_dataset, normalize_fn = load_raw_dataset(train_data_path)
        raw_dataset = raw_dataset.shuffle(seed=seed)

        if max_samples is not None and len(raw_dataset) > 3 * max_samples:
            # Reduce size to 3 * max_samples to reduce processing
            # This will then be reduced further to max_samples
            # after combining datasets and shuffling
            raw_dataset = raw_dataset.select(range(3 * max_samples))

        if normalize_fn is not None:
            raw_dataset = raw_dataset.map(
                normalize_fn,
                num_proc=build_dataset_num_proc,
                keep_in_memory=True,  # skip caching
            )

        pretokenized = {"input_ids", "loss_mask"} <= set(raw_dataset.column_names)
        if not pretokenized and not processor_has_chat_template:
            raise ValueError(
                f"Processor for {target_model_path} does not support chat templates. "
                "Please use a model with a pre-configured chat template or provide "
                "pre-tokenized input_ids and loss_mask columns."
            )

        log.info(f"Loaded {len(raw_dataset)} samples")

        preprocessed_dataset = build_speculator_training_dataset(
            dataset=raw_dataset,
            processor=processor,
            max_length=seq_length,
            num_proc=build_dataset_num_proc,
            render_endpoint=render_endpoint,
            minimum_valid_tokens=minimum_valid_tokens,
        )
        if minimum_valid_tokens is not None:
            log.info(f"Kept {len(preprocessed_dataset)} samples after filtering")
        processed_datasets.append(preprocessed_dataset)

    combined_dataset = concatenate_datasets(processed_datasets)
    combined_dataset = combined_dataset.shuffle(seed=seed)
    if max_samples is not None and len(combined_dataset) > max_samples:
        combined_dataset = combined_dataset.select(range(max_samples))

    if len(combined_dataset) == 0 and not allow_empty_output:
        raise ValueError(
            "No samples remain after preprocessing. Check the dataset schema, "
            "assistant masking, and --minimum-valid-tokens. Pass "
            "--allow-empty-output if an empty dataset is intentional."
        )

    log.subsection("Computing token frequency distribution")
    save_token_frequency_distribution(
        dataset=combined_dataset,
        output_path=token_freq_path,
    )

    if len(combined_dataset) == 0:
        log.warning("No samples remain after preprocessing; skipping visualization")
    else:
        log.subsection("Visualizing sample")
        _visualize_sample(combined_dataset, processor, idx=0)

    log.section("Dataset preprocessing complete")

    return combined_dataset, processor

def load_raw_dataset(
    train_data_path: str,
) -> tuple[HFDataset, Callable[[dict], dict] | None]:
    """Load a raw dataset from one of several source types.

    Resolution order:
        1. Local ``.json``/``.jsonl`` file.
        2. Local directory: recursively load all ``*.json``/``*.jsonl`` files
           as a single dataset.
        3. Named preset from ``DATASET_CONFIGS``.
        4. ``hf:<id>[:<subset>:<split>]`` for an arbitrary HuggingFace dataset.

    Args:
        train_data_path: File path, directory path, preset name, or ``hf:`` spec.

    Returns:
        Tuple of (raw_dataset, normalize_fn). normalize_fn is None for sources
        already in conversations format.

    Raises:
        ValueError: If the source cannot be resolved or a local directory
            contains no ``.json``/``.jsonl`` files.
    """
    # 1. Local file
    if train_data_path.endswith((".jsonl", ".json")):
        return load_dataset("json", data_files=train_data_path, split="train"), None

    # 2. Local directory
    path = Path(train_data_path)
    if path.is_dir():
        data_files = sorted(
            str(p) for p in (*path.rglob("*.json"), *path.rglob("*.jsonl"))
        )
        if not data_files:
            raise ValueError(
                f"No .json/.jsonl files found in directory: {train_data_path}"
            )
        return load_dataset("json", data_files=data_files, split="train"), None

    # 3. Named preset
    if train_data_path in DATASET_CONFIGS:
        config = DATASET_CONFIGS[train_data_path]
        raw_dataset = load_dataset(
            config.hf_path, name=config.subset, split=config.split
        )
        if config.filter_fn is not None:
            raw_dataset = raw_dataset.filter(config.filter_fn)
        return raw_dataset, config.normalize_fn

    # 4. Arbitrary HuggingFace dataset
    if train_data_path.startswith("hf:"):
        return _load_hf_dataset(train_data_path)

    raise ValueError(
        f"Unsupported dataset: {train_data_path}. Supported: local .json/.jsonl "
        f"file, local directory of .json/.jsonl files, hf:<id>[:<subset>:<split>], "
        f"or a preset {list(DATASET_CONFIGS.keys())}."
    )

def usable_cpu_count() -> int:
    """Return the CPUs available to this process, respecting affinity."""
    if hasattr(os, "process_cpu_count"):  # Python 3.13+
        return os.process_cpu_count() or 1
    if hasattr(os, "sched_getaffinity"):  # Linux
        return len(os.sched_getaffinity(0))
    return os.cpu_count() or 1