training_scripts/train_model.py

#!/usr/bin/env -S uv run python
"""Reproduce the Bekko Static Embedding Multilingual v1 model.

What this script does
---------------------
This is the standalone production trainer for v1. From a YAML recipe, it
validates and materializes selected rows from the pinned public Bekko
unsupervised dataset, builds deterministic locality-aware batches, and trains a
mean-pooled static embedding table with Multiple Negatives Ranking Loss (MNRL)
and Matryoshka Representation Learning. It runs sampled HAKARI reranking
evaluations during training, saves the final SentenceTransformers artifact, raises only the final
tokenizer ceiling to the configured 32,768 tokens, and can launch the full
HAKARI evaluation after training.

How to run it
-------------
Install the dependencies listed in the repository's v1 reproduction guide,
prepare the YAML and data-selection JSON, then run::

    python train_model.py --config path/to/v1-reproduction.yaml

Run ``python train_model.py --help`` for materialization-only and sharded
materialization options. The YAML must use portable paths chosen by the person
running the recipe. Full evaluation requires a checkout of the public
``hakari-bench`` repository referenced by ``evaluation.hakari_repo``.

Reproducing v1
--------------
First run ``train_tokenizer.py`` and point ``model.tokenizer_model`` at its
output. Use the checked-in v1 YAML and data-selection JSON without changing the
model, objective, sampling, or optimizer settings. Machine-specific storage and
HAKARI checkout paths may be changed. The public Bekko ``shared_id`` field is
retained during resumable materialization; no UniR checkout, private dataset,
or project-local Python package is required. Exact commands, dependency pins,
artifact hashes, and completion checks are in the v1 reproduction guide.
"""

from __future__ import annotations

# ==============================================================================
# Configuration
# ==============================================================================
import math
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Literal, cast

import yaml


@dataclass(frozen=True)
class RunSettings:
    mode: Literal["smoke", "search", "full"]
    name: str
    output_dir: Path


@dataclass(frozen=True)
class ModelSettings:
    tokenizer_model: str
    text_preprocessing: Literal["none", "porter", "snowball", "lemma"]
    embedding_dim: int
    max_seq_length: int
    final_max_seq_length: int
    matryoshka_dims: tuple[int, ...]
    matryoshka_weights: tuple[float, ...]
    architecture: Literal["static", "sse"]
    sse_preprocessing: Literal["official_tokenize", "sentence_transformers"]
    sse_alpha_init: float
    sse_train_alpha: bool
    sse_train_beta: bool
    sse_train_bias: bool


@dataclass(frozen=True)
class SupplementalDataSettings:
    """A separately sampled Hugging Face component with auditable provenance.

    ``provenance_note`` is mandatory because released rows can be sampled or
    expanded instances rather than independent, query-unique observations.
    """

    name: str
    dataset: str
    config: str
    revision: str
    split: str
    task: Literal["pair", "triplet"]
    pair_type: Literal["query_doc", "doc_doc"]
    fraction: float
    provenance_note: str
    query_column: str = "query"
    positive_column: str = "pos"
    negative_column: str = "neg"
    batch_deduplication: Literal["text", "model_input"] = "text"
    fingerprint_cache: Path | None = None


@dataclass(frozen=True)
class DataSettings:
    spec: Path
    fraction: float
    seed: int
    tasks: tuple[str, ...]
    require_cache: bool
    shared_id_column: str
    shared_id_block_size: int
    supplements: tuple[SupplementalDataSettings, ...]
    stage_cache_root: Path | None = None
    require_stage_cache: bool = False
    direct_source_cache: bool = False
    materialize_workers: int = 1
    materialize_max_attempts: int = 3
    materialize_retry_delay_seconds: float = 30.0


@dataclass(frozen=True)
class TrainingSettings:
    epochs: float
    batch_size: int
    eval_batch_size: int
    learning_rate: float
    sse_learning_rate: float | None
    adam_beta2: float
    adam_epsilon: float
    weight_decay: float
    warmup_ratio: float
    target_evaluations: int
    exact_target_evaluations: bool
    dataloader_num_workers: int
    dataloader_prefetch_factor: int
    dataloader_locality_block_size: int | None
    save_checkpoints: bool
    save_total_limit: int
    seed: int


@dataclass(frozen=True)
class LossSettings:
    name: Literal["mnrl", "qwen3_improved_infonce"]
    temperature: float
    margin: float | None
    hard_negative_margin: float | None
    pair_type_routing: Literal["shared", "split"]
    doc_doc_bidirectional: bool
    query_query_weight: float
    document_document_weight: float
    cross_document_weight: float
    hard_negative_weight: float
    qat_int8_weight: float
    qat_binary_weight: float
    qat_n_quantized_per_step: int


@dataclass(frozen=True)
class EvaluationTargetSettings:
    dataset: str | None
    collection: str | None
    splits: tuple[str, ...]
    query_limit: int


@dataclass(frozen=True)
class EvaluationSettings:
    hakari_repo: Path
    batch_size: int
    query_limit: int
    query_seed: int
    rerank_top_k: int
    candidate_ranking: str
    targets: tuple[EvaluationTargetSettings, ...]
    full_after_training: bool
    full_all: bool
    full_results_dir: Path | None
    full_model_alias: str | None
    full_dtype: Literal["bf16", "fp16", "fp32"]
    full_device: str
    full_retrieval_score_device: Literal["auto", "cpu", "cuda"]
    full_result_format: Literal["json.xz", "json"]
    full_overwrite: bool


@dataclass(frozen=True)
class TrackingSettings:
    enabled: bool
    project: str


@dataclass(frozen=True)
class TrainConfig:
    run: RunSettings
    model: ModelSettings
    data: DataSettings
    training: TrainingSettings
    loss: LossSettings
    evaluation: EvaluationSettings
    tracking: TrackingSettings
    source_path: Path


def _literal_choice(value: object, choices: set[str], name: str) -> Any:
    text = str(value)
    if text not in choices:
        expected = ", ".join(sorted(choices))
        raise ValueError(f"{name} must be one of {expected}; got {text!r}")
    return text


def load_config(path: Path | str) -> TrainConfig:
    source_path = Path(path).resolve()
    with source_path.open(encoding="utf-8") as handle:
        raw = yaml.safe_load(handle)
    if not isinstance(raw, dict):
        raise TypeError("Configuration root must be a mapping")

    run = _mapping(raw, "run")
    model = _mapping(raw, "model")
    data = _mapping(raw, "data")
    training = _mapping(raw, "training")
    evaluation = _mapping(raw, "evaluation")
    tracking = _mapping(raw, "tracking")
    loss = raw.get("loss", {})
    if not isinstance(loss, dict):
        raise TypeError("Configuration section 'loss' must be a mapping")

    mode = str(run["mode"])
    if mode not in {"smoke", "search", "full"}:
        raise ValueError(f"Unsupported run mode: {mode}")
    fraction = float(data["fraction"])
    if not 0 < fraction <= 1:
        raise ValueError(f"data.fraction must be in (0, 1], got {fraction}")
    batch_size = int(training["batch_size"])
    if batch_size <= 0:
        raise ValueError("training.batch_size must be positive")
    materialize_workers = int(data.get("materialize_workers", 1))
    if materialize_workers <= 0:
        raise ValueError("data.materialize_workers must be positive")
    materialize_max_attempts = int(data.get("materialize_max_attempts", 3))
    if materialize_max_attempts <= 0:
        raise ValueError("data.materialize_max_attempts must be positive")
    materialize_retry_delay_seconds = float(
        data.get("materialize_retry_delay_seconds", 30.0)
    )
    if materialize_retry_delay_seconds < 0:
        raise ValueError("data.materialize_retry_delay_seconds must be non-negative")

    tracking_enabled = bool(tracking.get("enabled", mode != "smoke"))
    if mode == "smoke":
        tracking_enabled = False

    return TrainConfig(
        run=RunSettings(
            mode=mode,  # type: ignore[arg-type]
            name=str(run["name"]),
            output_dir=Path(run["output_dir"]),
        ),
        model=_load_model_settings(model),
        data=DataSettings(
            spec=Path(data["spec"]),
            fraction=fraction,
            seed=int(data.get("seed", 42)),
            tasks=tuple(
                str(value)
                for value in data.get("tasks", ["pair", "triplet", "fw2ir-triplet"])
            ),
            require_cache=bool(data.get("require_cache", True)),
            shared_id_column=str(data.get("shared_id_column", "shared_id")),
            shared_id_block_size=int(data.get("shared_id_block_size", 8192)),
            supplements=_load_supplemental_data_settings(data),
            stage_cache_root=(
                Path(str(data["stage_cache_root"])).expanduser()
                if data.get("stage_cache_root") is not None
                else None
            ),
            require_stage_cache=bool(data.get("require_stage_cache", False)),
            direct_source_cache=bool(data.get("direct_source_cache", False)),
            materialize_workers=materialize_workers,
            materialize_max_attempts=materialize_max_attempts,
            materialize_retry_delay_seconds=materialize_retry_delay_seconds,
        ),
        training=_load_training_settings(training, batch_size=batch_size),
        loss=_load_loss_settings(loss),
        evaluation=EvaluationSettings(
            hakari_repo=Path(evaluation["hakari_repo"]),
            batch_size=int(evaluation.get("batch_size", 1024)),
            query_limit=int(evaluation.get("query_limit", 10)),
            query_seed=int(evaluation.get("query_seed", 13)),
            rerank_top_k=int(evaluation.get("rerank_top_k", 100)),
            candidate_ranking=str(
                evaluation.get("candidate_ranking", "reranking_hybrid")
            ),
            targets=_load_evaluation_targets(evaluation),
            full_after_training=bool(evaluation.get("full_after_training", False)),
            full_all=bool(evaluation.get("full_all", False)),
            full_results_dir=(
                Path(str(evaluation["full_results_dir"])).expanduser()
                if evaluation.get("full_results_dir") is not None
                else None
            ),
            full_model_alias=(
                str(evaluation["full_model_alias"])
                if evaluation.get("full_model_alias") is not None
                else None
            ),
            full_dtype=_literal_choice(
                evaluation.get("full_dtype", "bf16"),
                {"bf16", "fp16", "fp32"},
                "evaluation.full_dtype",
            ),
            full_device=str(evaluation.get("full_device", "cuda:0")),
            full_retrieval_score_device=_literal_choice(
                evaluation.get("full_retrieval_score_device", "cpu"),
                {"auto", "cpu", "cuda"},
                "evaluation.full_retrieval_score_device",
            ),
            full_result_format=_literal_choice(
                evaluation.get("full_result_format", "json.xz"),
                {"json.xz", "json"},
                "evaluation.full_result_format",
            ),
            full_overwrite=bool(evaluation.get("full_overwrite", False)),
        ),
        tracking=TrackingSettings(
            enabled=tracking_enabled,
            project=str(tracking.get("project", "bekko-static-embedding")),
        ),
        source_path=source_path,
    )


def _load_evaluation_targets(
    evaluation: dict[str, Any],
) -> tuple[EvaluationTargetSettings, ...]:
    raw_targets = evaluation.get("targets", [])
    if not isinstance(raw_targets, list):
        raise TypeError("evaluation.targets must be a list")
    targets: list[EvaluationTargetSettings] = []
    default_query_limit = int(evaluation.get("query_limit", 10))
    for index, raw in enumerate(raw_targets):
        if not isinstance(raw, dict):
            raise TypeError(f"evaluation.targets[{index}] must be a mapping")
        target = cast(dict[str, Any], raw)
        dataset = str(target["dataset"]) if target.get("dataset") is not None else None
        collection = (
            str(target["collection"]) if target.get("collection") is not None else None
        )
        if (dataset is None) == (collection is None):
            raise ValueError(
                f"evaluation.targets[{index}] requires exactly one of dataset or collection"
            )
        splits_raw = target.get("splits", [])
        if not isinstance(splits_raw, list):
            raise TypeError(f"evaluation.targets[{index}].splits must be a list")
        query_limit = int(target.get("query_limit", default_query_limit))
        if query_limit <= 0:
            raise ValueError(
                f"evaluation.targets[{index}].query_limit must be positive"
            )
        targets.append(
            EvaluationTargetSettings(
                dataset=dataset,
                collection=collection,
                splits=tuple(str(value) for value in splits_raw),
                query_limit=query_limit,
            )
        )
    return tuple(targets)


def _load_supplemental_data_settings(
    data: dict[str, Any],
) -> tuple[SupplementalDataSettings, ...]:
    raw_supplements = data.get("supplements", [])
    if not isinstance(raw_supplements, list):
        raise TypeError("data.supplements must be a list")
    supplements: list[SupplementalDataSettings] = []
    names: set[str] = set()
    for index, raw in enumerate(raw_supplements):
        if not isinstance(raw, dict):
            raise TypeError(f"data.supplements[{index}] must be a mapping")
        raw_mapping = cast(dict[str, Any], raw)
        name = str(raw_mapping.get("name", "")).strip()
        if not name:
            raise ValueError(f"data.supplements[{index}].name must not be empty")
        if name in names:
            raise ValueError(f"duplicate supplemental data name: {name}")
        names.add(name)
        task = str(raw_mapping.get("task", ""))
        if task not in {"pair", "triplet"}:
            raise ValueError(f"Unsupported supplemental task: {task}")
        pair_type = str(raw_mapping.get("pair_type", ""))
        if pair_type not in {"query_doc", "doc_doc"}:
            raise ValueError(f"Unsupported supplemental pair_type: {pair_type}")
        fraction = float(raw_mapping.get("fraction", 1.0))
        if not 0 < fraction <= 1:
            raise ValueError(
                f"data.supplements[{index}].fraction must be in (0, 1], got {fraction}"
            )
        provenance_note = str(raw_mapping.get("provenance_note", "")).strip()
        if not provenance_note:
            raise ValueError(
                f"data.supplements[{index}].provenance_note must not be empty"
            )
        revision = str(raw_mapping.get("revision", "")).strip()
        if not revision:
            raise ValueError(f"data.supplements[{index}].revision must not be empty")
        batch_deduplication = str(raw_mapping.get("batch_deduplication", "text"))
        if batch_deduplication not in {"text", "model_input"}:
            raise ValueError(
                "data.supplements"
                f"[{index}].batch_deduplication must be text or model_input"
            )
        fingerprint_cache_raw = raw_mapping.get("fingerprint_cache")
        supplements.append(
            SupplementalDataSettings(
                name=name,
                dataset=str(raw_mapping["dataset"]),
                config=str(raw_mapping["config"]),
                revision=revision,
                split=str(raw_mapping.get("split", "train")),
                task=task,  # type: ignore[arg-type]
                pair_type=pair_type,  # type: ignore[arg-type]
                fraction=fraction,
                provenance_note=provenance_note,
                query_column=str(raw_mapping.get("query_column", "query")),
                positive_column=str(raw_mapping.get("positive_column", "pos")),
                negative_column=str(raw_mapping.get("negative_column", "neg")),
                batch_deduplication=batch_deduplication,  # type: ignore[arg-type]
                fingerprint_cache=(
                    Path(str(fingerprint_cache_raw))
                    if fingerprint_cache_raw is not None
                    else None
                ),
            )
        )
    return tuple(supplements)


def _load_loss_settings(raw: dict[str, Any]) -> LossSettings:
    name = str(raw.get("name", "mnrl"))
    if name not in {"mnrl", "qwen3_improved_infonce"}:
        raise ValueError(f"Unsupported loss.name: {name}")
    temperature = float(raw.get("temperature", 0.01))
    if temperature <= 0:
        raise ValueError("loss.temperature must be positive")
    margin_raw = raw.get("margin", 0.1)
    hard_margin_raw = raw.get("hard_negative_margin")
    pair_type_routing = str(raw.get("pair_type_routing", "shared"))
    if pair_type_routing not in {"shared", "split"}:
        raise ValueError(f"Unsupported loss.pair_type_routing: {pair_type_routing}")
    weights = {
        name: float(raw.get(name, 1.0))
        for name in (
            "query_query_weight",
            "document_document_weight",
            "cross_document_weight",
            "hard_negative_weight",
        )
    }
    for weight_name, value in weights.items():
        if not math.isfinite(value) or value < 0:
            raise ValueError(f"loss.{weight_name} must be finite and non-negative")
    qat_weights = {
        name: float(raw.get(name, 0.0))
        for name in ("qat_int8_weight", "qat_binary_weight")
    }
    for weight_name, value in qat_weights.items():
        if not math.isfinite(value) or value < 0:
            raise ValueError(f"loss.{weight_name} must be finite and non-negative")
    qat_n_quantized_per_step = int(raw.get("qat_n_quantized_per_step", -1))
    if qat_n_quantized_per_step not in {-1, 1, 2}:
        raise ValueError("loss.qat_n_quantized_per_step must be -1, 1, or 2")
    return LossSettings(
        name=name,  # type: ignore[arg-type]
        temperature=temperature,
        margin=None if margin_raw is None else float(margin_raw),
        hard_negative_margin=None
        if hard_margin_raw is None
        else float(hard_margin_raw),
        pair_type_routing=pair_type_routing,  # type: ignore[arg-type]
        doc_doc_bidirectional=bool(raw.get("doc_doc_bidirectional", False)),
        **weights,
        **qat_weights,
        qat_n_quantized_per_step=qat_n_quantized_per_step,
    )


def _load_model_settings(raw: dict[str, Any]) -> ModelSettings:
    embedding_dim = int(raw.get("embedding_dim", 1024))
    max_seq_length = int(raw.get("max_seq_length", 512))
    final_max_seq_length = int(raw.get("final_max_seq_length", max_seq_length))
    architecture = str(raw.get("architecture", "static"))
    if architecture not in {"static", "sse"}:
        raise ValueError(f"Unsupported model.architecture: {architecture}")
    sse_preprocessing = str(raw.get("sse_preprocessing", "sentence_transformers"))
    if sse_preprocessing not in {"official_tokenize", "sentence_transformers"}:
        raise ValueError(f"Unsupported model.sse_preprocessing: {sse_preprocessing}")
    sse_alpha_init = float(raw.get("sse_alpha_init", 0.5))
    text_preprocessing = str(raw.get("text_preprocessing", "none"))
    if text_preprocessing not in {"none", "porter", "snowball", "lemma"}:
        raise ValueError(f"Unsupported model.text_preprocessing: {text_preprocessing}")
    if architecture == "sse" and text_preprocessing != "none":
        raise ValueError(
            "model.text_preprocessing is supported only for static architecture"
        )
    if not math.isfinite(sse_alpha_init) or sse_alpha_init <= 0:
        raise ValueError("model.sse_alpha_init must be finite and positive")
    dims = tuple(int(value) for value in raw.get("matryoshka_dims", [embedding_dim]))
    weights = tuple(
        float(value) for value in raw.get("matryoshka_weights", [1.0] * len(dims))
    )
    if embedding_dim <= 0:
        raise ValueError("model.embedding_dim must be positive")
    if max_seq_length <= 0:
        raise ValueError("model.max_seq_length must be positive")
    if final_max_seq_length < max_seq_length:
        raise ValueError(
            "model.final_max_seq_length must be at least model.max_seq_length"
        )
    if any(value <= 0 or value > embedding_dim for value in dims):
        raise ValueError(
            "model.matryoshka_dims must be positive and no larger than embedding_dim"
        )
    is_ascending = tuple(sorted(set(dims))) == dims
    is_descending = tuple(sorted(set(dims), reverse=True)) == dims
    if not (is_ascending or is_descending):
        raise ValueError("model.matryoshka_dims must be unique and strictly monotonic")
    if len(weights) != len(dims):
        raise ValueError("model.matryoshka_weights must match model.matryoshka_dims")
    if any(not math.isfinite(value) or value < 0 for value in weights):
        raise ValueError("model.matryoshka_weights must be finite and non-negative")
    if weights and weights[0] <= 0:
        raise ValueError("the largest Matryoshka dimension must have positive weight")
    return ModelSettings(
        tokenizer_model=str(raw["tokenizer_model"]),
        text_preprocessing=text_preprocessing,  # type: ignore[arg-type]
        embedding_dim=embedding_dim,
        max_seq_length=max_seq_length,
        final_max_seq_length=final_max_seq_length,
        matryoshka_dims=dims,
        matryoshka_weights=weights,
        architecture=architecture,  # type: ignore[arg-type]
        sse_preprocessing=sse_preprocessing,  # type: ignore[arg-type]
        sse_alpha_init=sse_alpha_init,
        sse_train_alpha=bool(raw.get("sse_train_alpha", True)),
        sse_train_beta=bool(raw.get("sse_train_beta", True)),
        sse_train_bias=bool(raw.get("sse_train_bias", True)),
    )


def _load_training_settings(
    raw: dict[str, Any], *, batch_size: int
) -> TrainingSettings:
    learning_rate = float(raw["learning_rate"])
    sse_learning_rate_raw = raw.get("sse_learning_rate")
    sse_learning_rate = (
        None if sse_learning_rate_raw is None else float(sse_learning_rate_raw)
    )
    adam_beta2 = float(raw.get("adam_beta2", 0.999))
    adam_epsilon = float(raw.get("adam_epsilon", 1e-8))
    weight_decay = float(raw.get("weight_decay", 0.0))
    locality_raw = raw.get("dataloader_locality_block_size")
    locality_block_size = None if locality_raw is None else int(locality_raw)
    if not math.isfinite(learning_rate) or learning_rate <= 0:
        raise ValueError("training.learning_rate must be finite and positive")
    if sse_learning_rate is not None and (
        not math.isfinite(sse_learning_rate) or sse_learning_rate < 0
    ):
        raise ValueError("training.sse_learning_rate must be finite and non-negative")
    if not math.isfinite(adam_beta2) or not 0 < adam_beta2 < 1:
        raise ValueError("training.adam_beta2 must be finite and in (0, 1)")
    if not math.isfinite(adam_epsilon) or adam_epsilon <= 0:
        raise ValueError("training.adam_epsilon must be finite and positive")
    if not math.isfinite(weight_decay) or weight_decay < 0:
        raise ValueError("training.weight_decay must be finite and non-negative")
    if locality_block_size is not None and locality_block_size < batch_size:
        raise ValueError(
            "training.dataloader_locality_block_size must be at least batch_size"
        )
    return TrainingSettings(
        epochs=float(raw.get("epochs", 1)),
        batch_size=batch_size,
        eval_batch_size=int(raw.get("eval_batch_size", 2048)),
        learning_rate=learning_rate,
        sse_learning_rate=sse_learning_rate,
        adam_beta2=adam_beta2,
        adam_epsilon=adam_epsilon,
        weight_decay=weight_decay,
        warmup_ratio=float(raw.get("warmup_ratio", 0.1)),
        target_evaluations=int(raw["target_evaluations"]),
        exact_target_evaluations=bool(raw.get("exact_target_evaluations", False)),
        dataloader_num_workers=int(raw.get("dataloader_num_workers", 15)),
        dataloader_prefetch_factor=int(raw.get("dataloader_prefetch_factor", 4)),
        dataloader_locality_block_size=locality_block_size,
        save_checkpoints=bool(raw.get("save_checkpoints", True)),
        save_total_limit=int(raw.get("save_total_limit", 3)),
        seed=int(raw.get("seed", 42)),
    )


def _mapping(raw: dict[str, Any], key: str) -> dict[str, Any]:
    value = raw.get(key)
    if not isinstance(value, dict):
        raise TypeError(f"Configuration section {key!r} must be a mapping")
    return value


# ==============================================================================
# Text preprocessing
# ==============================================================================

import json
import re
from functools import lru_cache
from typing import Self

import numpy as np
import simplemma
import torch
from nltk.stem import PorterStemmer
from nltk.stem.snowball import EnglishStemmer
from sentence_transformers.sentence_transformer.modules import StaticEmbedding
from tokenizers import Tokenizer
from transformers import PreTrainedTokenizerFast

TextPreprocessing = Literal["none", "porter", "snowball", "lemma"]

_WORD_RE = re.compile(r"[^\W\d_]+", flags=re.UNICODE)
_PORTER = PorterStemmer()
_SNOWBALL = EnglishStemmer()


def _restore_case(source: str, transformed: str) -> str:
    if source.isupper():
        return transformed.upper()
    if source[:1].isupper() and source[1:].islower():
        return transformed[:1].upper() + transformed[1:]
    return transformed


@lru_cache(maxsize=200_000)
def _porter_word(word: str) -> str:
    return _restore_case(word, _PORTER.stem(word.lower(), to_lowercase=False))


@lru_cache(maxsize=200_000)
def _snowball_word(word: str) -> str:
    return _restore_case(word, _SNOWBALL.stem(word.lower()))


@lru_cache(maxsize=200_000)
def _lemma_word(word: str) -> str:
    return _restore_case(word, simplemma.lemmatize(word.lower(), lang="en"))


def preprocess_text(text: str, mode: TextPreprocessing) -> str:
    if mode == "none":
        return text
    if mode == "porter":
        transform = _porter_word
    elif mode == "snowball":
        transform = _snowball_word
    elif mode == "lemma":
        transform = _lemma_word
    else:
        raise ValueError(f"Unsupported text preprocessing mode: {mode}")
    return _WORD_RE.sub(lambda match: transform(match.group(0)), text)


class PreprocessedStaticEmbedding(StaticEmbedding):
    """StaticEmbedding with a serialized lexical preprocessing step."""

    config_filename = "text_preprocessing_config.json"

    def __init__(
        self,
        tokenizer: Tokenizer | PreTrainedTokenizerFast,
        embedding_weights: np.ndarray | torch.Tensor | None = None,
        embedding_dim: int | None = None,
        *,
        text_preprocessing: TextPreprocessing,
        **kwargs: Any,
    ) -> None:
        if text_preprocessing not in {"porter", "snowball", "lemma"}:
            raise ValueError(
                "PreprocessedStaticEmbedding requires porter, snowball, or lemma preprocessing"
            )
        super().__init__(
            tokenizer,
            embedding_weights=embedding_weights,
            embedding_dim=embedding_dim,
            **kwargs,
        )
        self.text_preprocessing = text_preprocessing

    def preprocess(
        self,
        inputs: list[str],
        prompt: str | None = None,
        **kwargs: Any,
    ) -> dict[str, torch.Tensor]:
        transformed = [
            preprocess_text(text, self.text_preprocessing) for text in inputs
        ]
        return super().preprocess(transformed, prompt=prompt, **kwargs)

    def save(
        self,
        output_path: str,
        *args: Any,
        safe_serialization: bool = True,
        **kwargs: Any,
    ) -> None:
        super().save(
            output_path,
            *args,
            safe_serialization=safe_serialization,
            **kwargs,
        )
        (Path(output_path) / self.config_filename).write_text(
            json.dumps(
                {"text_preprocessing": self.text_preprocessing},
                indent=2,
                sort_keys=True,
            )
            + "\n",
            encoding="utf-8",
        )

    @classmethod
    def load(
        cls,
        model_name_or_path: str,
        subfolder: str = "",
        token: bool | str | None = None,
        cache_folder: str | None = None,
        revision: str | None = None,
        local_files_only: bool = False,
        **kwargs: Any,
    ) -> Self:
        hub_kwargs = {
            "subfolder": subfolder,
            "token": token,
            "cache_folder": cache_folder,
            "revision": revision,
            "local_files_only": local_files_only,
        }
        tokenizer_path = cls.load_file_path(
            model_name_or_path, filename="tokenizer.json", **hub_kwargs
        )
        config_path = cls.load_file_path(
            model_name_or_path, filename=cls.config_filename, **hub_kwargs
        )
        if tokenizer_path is None or config_path is None:
            raise FileNotFoundError(
                "PreprocessedStaticEmbedding artifact is incomplete"
            )
        tokenizer = Tokenizer.from_file(tokenizer_path)
        config = json.loads(Path(config_path).read_text(encoding="utf-8"))
        weights = cls.load_torch_weights(
            model_name_or_path=model_name_or_path,
            subfolder=subfolder,
            token=token,
            cache_folder=cache_folder,
            revision=revision,
            local_files_only=local_files_only,
        )
        model = cls(
            tokenizer,
            embedding_weights=weights["embedding.weight"],
            text_preprocessing=config["text_preprocessing"],
        )
        model.load_state_dict(weights)
        return model


# ==============================================================================
# Stable static embedding
# ==============================================================================

from typing import Literal

from sentence_transformers.sentence_transformer.modules import StaticEmbedding
from torch import nn

SSEPreprocessing = Literal["official_tokenize", "sentence_transformers"]


class SeparableDyT(nn.Module):
    """Per-dimension Dynamic Tanh from the official SSE implementation."""

    def __init__(self, hidden_dim: int, alpha_init: float = 0.5) -> None:
        super().__init__()
        self.alpha = nn.Parameter(alpha_init * torch.ones(hidden_dim))
        self.beta = nn.Parameter(torch.ones(hidden_dim))
        self.bias = nn.Parameter(torch.zeros(hidden_dim))

    def forward(self, values: torch.Tensor) -> torch.Tensor:
        return self.beta * torch.tanh(self.alpha * values + self.bias)


class StableStaticEmbedding(StaticEmbedding):
    """Sentence Transformers StaticEmbedding followed by Separable DyT.

    The module exposes both the literal token-flattening path from the official
    SSE repository and the current Sentence Transformers ``preprocess`` path.
    They are intentionally selectable so upstream API migrations can be tested
    without changing the model architecture.
    """

    config_filename = "sse_config.json"

    def __init__(
        self,
        tokenizer: Tokenizer | PreTrainedTokenizerFast,
        embedding_weights: np.ndarray | torch.Tensor | None = None,
        embedding_dim: int | None = None,
        *,
        preprocessing: SSEPreprocessing = "sentence_transformers",
        alpha_init: float = 0.5,
        train_alpha: bool = True,
        train_beta: bool = True,
        train_bias: bool = True,
        **kwargs: Any,
    ) -> None:
        if preprocessing not in {"official_tokenize", "sentence_transformers"}:
            raise ValueError(f"Unsupported SSE preprocessing mode: {preprocessing}")
        if alpha_init <= 0:
            raise ValueError("SSE alpha_init must be positive")
        super().__init__(
            tokenizer,
            embedding_weights=embedding_weights,
            embedding_dim=embedding_dim,
            **kwargs,
        )
        self.preprocessing: SSEPreprocessing = preprocessing
        self.alpha_init = float(alpha_init)
        self.dyt = SeparableDyT(self.embedding_dim, alpha_init=self.alpha_init)
        self.train_alpha = bool(train_alpha)
        self.train_beta = bool(train_beta)
        self.train_bias = bool(train_bias)
        self.dyt.alpha.requires_grad_(self.train_alpha)
        self.dyt.beta.requires_grad_(self.train_beta)
        self.dyt.bias.requires_grad_(self.train_bias)

    def official_tokenize(
        self, texts: list[str], **kwargs: Any
    ) -> dict[str, torch.Tensor]:
        """Literal flat ``input_ids``/``offsets`` algorithm used by official SSE."""
        del kwargs
        encodings = self.tokenizer.encode_batch(texts, add_special_tokens=False)
        encodings_ids = [encoding.ids for encoding in encodings]
        offsets = torch.from_numpy(
            np.cumsum([0] + [len(token_ids) for token_ids in encodings_ids[:-1]])
        )
        input_ids = torch.tensor(
            [token_id for token_ids in encodings_ids for token_id in token_ids],
            dtype=torch.long,
        )
        return {"input_ids": input_ids, "offsets": offsets}

    def sentence_transformers_preprocess(
        self,
        inputs: list[str],
        prompt: str | None = None,
        **kwargs: Any,
    ) -> dict[str, torch.Tensor]:
        return super().preprocess(inputs, prompt=prompt, **kwargs)

    def preprocess(
        self,
        inputs: list[str],
        prompt: str | None = None,
        **kwargs: Any,
    ) -> dict[str, torch.Tensor]:
        if self.preprocessing == "official_tokenize":
            if prompt:
                raise ValueError(
                    "The official SSE tokenize path does not support prompts"
                )
            return self.official_tokenize(inputs, **kwargs)
        return self.sentence_transformers_preprocess(inputs, prompt=prompt, **kwargs)

    def tokenize(self, texts: list[str], **kwargs: Any) -> dict[str, torch.Tensor]:
        if self.preprocessing == "official_tokenize":
            return self.official_tokenize(texts, **kwargs)
        return self.sentence_transformers_preprocess(texts, **kwargs)

    def forward(
        self,
        features: dict[str, torch.Tensor],
        **kwargs: Any,
    ) -> dict[str, torch.Tensor]:
        features = super().forward(features, **kwargs)
        features["sentence_embedding"] = self.dyt(features["sentence_embedding"])
        return features

    def get_sentence_embedding_dimension(self) -> int:
        return self.embedding_dim

    def save(
        self,
        output_path: str,
        *args: Any,
        safe_serialization: bool = True,
        **kwargs: Any,
    ) -> None:
        super().save(
            output_path,
            *args,
            safe_serialization=safe_serialization,
            **kwargs,
        )
        config = {
            "preprocessing": self.preprocessing,
            "alpha_init": self.alpha_init,
            "train_alpha": self.train_alpha,
            "train_beta": self.train_beta,
            "train_bias": self.train_bias,
        }
        (Path(output_path) / self.config_filename).write_text(
            json.dumps(config, indent=2, sort_keys=True) + "\n",
            encoding="utf-8",
        )

    @classmethod
    def load(
        cls,
        model_name_or_path: str,
        subfolder: str = "",
        token: bool | str | None = None,
        cache_folder: str | None = None,
        revision: str | None = None,
        local_files_only: bool = False,
        **kwargs: Any,
    ) -> Self:
        hub_kwargs = {
            "subfolder": subfolder,
            "token": token,
            "cache_folder": cache_folder,
            "revision": revision,
            "local_files_only": local_files_only,
        }
        tokenizer_path = cls.load_file_path(
            model_name_or_path,
            filename="tokenizer.json",
            **hub_kwargs,
        )
        if tokenizer_path is None:
            raise FileNotFoundError("SSE tokenizer.json was not found")
        tokenizer = Tokenizer.from_file(tokenizer_path)
        weights = cls.load_torch_weights(
            model_name_or_path=model_name_or_path,
            subfolder=subfolder,
            token=token,
            cache_folder=cache_folder,
            revision=revision,
            local_files_only=local_files_only,
        )

        config_path = cls.load_file_path(
            model_name_or_path,
            filename=cls.config_filename,
            **hub_kwargs,
        )
        config: dict[str, Any] = {}
        if config_path is not None:
            config = json.loads(Path(config_path).read_text(encoding="utf-8"))
        embedding_weights = weights["embedding.weight"]
        model = cls(
            tokenizer,
            embedding_weights=embedding_weights,
            preprocessing=config.get("preprocessing", "sentence_transformers"),
            alpha_init=float(config.get("alpha_init", 0.5)),
            train_alpha=bool(config.get("train_alpha", True)),
            train_beta=bool(config.get("train_beta", True)),
            train_bias=bool(config.get("train_bias", True)),
        )
        model.load_state_dict(weights)
        return model


# ==============================================================================
# Model construction
# ==============================================================================


from sentence_transformers import SentenceTransformer
from sentence_transformers.sentence_transformer.modules import StaticEmbedding


def _wrap_sentence_transformer(
    static_embedding: StaticEmbedding,
) -> SentenceTransformer:
    return SentenceTransformer(modules=[static_embedding])


def build_scratch_model(
    tokenizer_model: str,
    embedding_dim: int = 1024,
    max_seq_length: int = 512,
    architecture: str = "static",
    sse_preprocessing: SSEPreprocessing = "sentence_transformers",
    sse_alpha_init: float = 0.5,
    sse_train_alpha: bool = True,
    sse_train_beta: bool = True,
    sse_train_bias: bool = True,
    text_preprocessing: TextPreprocessing = "none",
) -> SentenceTransformer:
    local_path = Path(tokenizer_model)
    local_tokenizer_path = (
        local_path / "tokenizer.json" if local_path.is_dir() else local_path
    )
    if local_tokenizer_path.is_file():
        tokenizer = Tokenizer.from_file(str(local_tokenizer_path))
    else:
        # Static training needs only the tokenizer. Loading a SentenceTransformer
        # first can unnecessarily materialize multi-billion-parameter backbones
        # (for example Gemma 2) before immediately discarding them.
        tokenizer = Tokenizer.from_pretrained(tokenizer_model)
    tokenizer.enable_truncation(max_length=max_seq_length)

    if architecture == "sse":
        static_embedding = StableStaticEmbedding(
            tokenizer,
            embedding_dim=embedding_dim,
            preprocessing=sse_preprocessing,
            alpha_init=sse_alpha_init,
            train_alpha=sse_train_alpha,
            train_beta=sse_train_beta,
            train_bias=sse_train_bias,
        )
    elif architecture == "static":
        if text_preprocessing == "none":
            static_embedding = StaticEmbedding(tokenizer, embedding_dim=embedding_dim)
        else:
            static_embedding = PreprocessedStaticEmbedding(
                tokenizer,
                embedding_dim=embedding_dim,
                text_preprocessing=text_preprocessing,
            )
    else:
        raise ValueError(f"Unsupported architecture: {architecture}")
    return _wrap_sentence_transformer(static_embedding)


# ==============================================================================
# Dataset loading
# ==============================================================================

import gc
import hashlib
import os
import random
import shutil
from typing import Literal

from datasets import (
    Dataset,
    DatasetDict,
    concatenate_datasets,
    load_dataset,
    load_from_disk,
)

Manifest = dict[str, dict[str, int | str]]


def staged_cache_path(source: Path, *, stage_root: Path) -> Path:
    """Map an absolute cache file to a collision-free path under ``stage_root``."""
    resolved = source.resolve()
    return stage_root.resolve() / resolved.as_posix().lstrip("/")


def stage_dataset_cache_files(
    dataset: Dataset,
    *,
    stage_root: Path,
    force: bool = False,
) -> dict[str, int]:
    """Copy dataset Arrow files from slower filesystems into a local stage.

    Files already on the stage filesystem are intentionally left in place. Each
    cross-device copy is published atomically so interrupted staging cannot be
    mistaken for a complete cache file by a later training process.
    """
    stage_root = stage_root.resolve()
    stage_root.mkdir(parents=True, exist_ok=True)
    stage_device = stage_root.stat().st_dev
    copied_files = 0
    copied_bytes = 0
    reused_files = 0
    local_files = 0
    for cache_file in dataset.cache_files:
        source = Path(str(cache_file["filename"])).resolve()
        source_size = source.stat().st_size
        if not force and source.stat().st_dev == stage_device:
            local_files += 1
            continue
        destination = staged_cache_path(source, stage_root=stage_root)
        if destination.is_file() and destination.stat().st_size == source_size:
            reused_files += 1
            continue
        destination.parent.mkdir(parents=True, exist_ok=True)
        temporary = destination.with_name(f".{destination.name}.partial-{os.getpid()}")
        try:
            shutil.copyfile(source, temporary)
            if temporary.stat().st_size != source_size:
                raise OSError(f"staged size mismatch for {source}")
            temporary.replace(destination)
        finally:
            temporary.unlink(missing_ok=True)
        copied_files += 1
        copied_bytes += source_size
    return {
        "copied_files": copied_files,
        "copied_bytes": copied_bytes,
        "reused_files": reused_files,
        "local_files": local_files,
    }


def remap_dataset_to_stage(
    dataset: Dataset,
    *,
    stage_root: Path,
    require: bool,
) -> Dataset:
    """Reopen cross-device Arrow shards from a completed local stage."""
    if len(dataset) == 0:
        return dataset
    cache_files = dataset.cache_files
    if not cache_files:
        return dataset
    stage_root = stage_root.resolve()
    if not stage_root.is_dir():
        if require:
            raise FileNotFoundError(
                f"required stage cache does not exist: {stage_root}"
            )
        return dataset
    stage_device = stage_root.stat().st_dev
    reopened: list[Dataset] = []
    for cache_file in cache_files:
        source = Path(str(cache_file["filename"])).resolve()
        source_size = source.stat().st_size
        staged = staged_cache_path(source, stage_root=stage_root)
        if staged.is_file() and staged.stat().st_size == source_size:
            selected = staged
        elif source.stat().st_dev == stage_device:
            selected = source
        elif require:
            raise FileNotFoundError(f"required staged cache file is missing: {staged}")
        else:
            selected = source
        reopened.append(Dataset.from_file(str(selected)))
    result = reopened[0] if len(reopened) == 1 else concatenate_datasets(reopened)
    if len(result) != len(dataset):
        raise ValueError(
            "staged dataset row count mismatch: "
            f"source={len(dataset)} staged={len(result)}"
        )
    return result


def select_fraction(
    dataset: Dataset,
    *,
    fraction: float,
    seed: int,
    shared_id_column: str = "shared_id",
    shared_id_block_size: int | None = None,
) -> Dataset:
    if not 0 < fraction <= 1:
        raise ValueError(f"fraction must be in (0, 1], got {fraction}")
    if fraction == 1 or len(dataset) == 0:
        return dataset
    if shared_id_column not in dataset.column_names:
        count = max(1, int(len(dataset) * fraction))
        start = random.Random(seed).randrange(len(dataset))
        end = min(len(dataset), start + count)
        segments = [dataset.select(range(start, end))]
        wrapped_count = count - (end - start)
        if wrapped_count:
            segments.append(dataset.select(range(wrapped_count)))
        return segments[0] if len(segments) == 1 else concatenate_datasets(segments)

    if shared_id_block_size:
        ranges = [
            (start, min(len(dataset), start + shared_id_block_size))
            for start in range(0, len(dataset), shared_id_block_size)
        ]
    else:
        values = dataset[shared_id_column]
        ranges = []
        start = 0
        for index in range(1, len(values) + 1):
            if index == len(values) or values[index] != values[start]:
                ranges.append((start, index))
                start = index
    block_count = max(1, min(len(ranges), round(len(ranges) * fraction)))
    start_block = random.Random(seed).randrange(len(ranges))
    first_block_count = min(block_count, len(ranges) - start_block)
    first_begin = ranges[start_block][0]
    first_end = ranges[start_block + first_block_count - 1][1]
    segments = [dataset.select(range(first_begin, first_end))]
    wrapped_block_count = block_count - first_block_count
    if wrapped_block_count:
        segments.append(
            dataset.select(range(ranges[0][0], ranges[wrapped_block_count - 1][1]))
        )
    return segments[0] if len(segments) == 1 else concatenate_datasets(segments)


def adapt_unir_dataset(dataset: Dataset, *, task: str) -> Dataset:
    required = {"query", "pos"}
    if task != "pair":
        required.add("neg")
    missing = required.difference(dataset.column_names)
    if missing:
        raise ValueError(f"Dataset for {task!r} is missing columns: {sorted(missing)}")
    rename = {"query": "anchor", "pos": "positive"}
    keep = ["query", "pos"]
    if task != "pair":
        rename["neg"] = "negative"
        keep.append("neg")
    dataset = dataset.select_columns(keep).rename_columns(rename)
    return dataset


def adapt_supplemental_dataset(
    dataset: Dataset,
    *,
    settings: SupplementalDataSettings,
) -> Dataset:
    source_columns = [settings.query_column, settings.positive_column]
    target_columns = ["anchor", "positive"]
    if settings.task == "triplet":
        source_columns.append(settings.negative_column)
        target_columns.append("negative")
    missing = set(source_columns).difference(dataset.column_names)
    if missing:
        raise ValueError(
            f"Supplement {settings.name!r} is missing columns: {sorted(missing)}"
        )
    adapted = dataset.select_columns(source_columns)
    adapted = adapted.rename_columns(
        dict(zip(source_columns, target_columns, strict=True))
    )
    metadata_target = cast(Any, adapted)
    metadata_target._bekko_batch_deduplication = settings.batch_deduplication
    metadata_target._bekko_fingerprint_cache = settings.fingerprint_cache
    metadata_target._bekko_supplement_name = settings.name
    metadata_target._bekko_supplement_identity = {
        "dataset": settings.dataset,
        "config": settings.config,
        "revision": settings.revision,
        "split": settings.split,
        "source_rows": len(dataset),
    }
    return adapted


def dataset_row_count(datasets: DatasetDict) -> int:
    return sum(len(dataset) for dataset in datasets.values())


def resolve_pair_type(
    item: Any, *, subset_name: str
) -> Literal["query_doc", "doc_doc"]:
    metadata = getattr(item, "metadata", None)
    pair_type = getattr(metadata, "pair_type", None)
    if pair_type not in {"query_doc", "doc_doc"}:
        raise ValueError(
            f"UniR subset {subset_name!r} is missing valid pair_type metadata"
        )
    return pair_type


def load_bekko_datasets(
    *,
    spec: Path,
    tasks: tuple[str, ...],
    fraction: float,
    seed: int,
    require_cache: bool,
    shared_id_column: str,
    shared_id_block_size: int = 8192,
    supplements: tuple[SupplementalDataSettings, ...] = (),
    stage_cache_root: Path | None = None,
    require_stage_cache: bool = False,
    direct_source_cache: bool = False,
) -> tuple[
    DatasetDict,
    Manifest,
    dict[str, Literal["query_doc", "doc_doc"]],
]:
    # This public trainer deliberately accepts only the auditable direct JSON
    # manifest. It never consults UniR specs or a private dataset cache.
    if supplements:
        raise ValueError(
            "The standalone direct-Bekko trainer does not support supplemental datasets."
        )
    del tasks
    direct_manifest = _read_direct_huggingface_manifest(spec)
    if direct_manifest is None:
        raise ValueError(
            "The standalone trainer requires a direct Bekko JSON manifest "
            "containing 'dataset', 'revision', and 'configs'."
        )
    return load_direct_huggingface_datasets(
        manifest=direct_manifest,
        fraction=fraction,
        seed=seed,
        require_cache=require_cache or require_stage_cache,
        shared_id_block_size=shared_id_block_size,
        materialized_root=stage_cache_root,
        source_cache_direct=direct_source_cache,
        shared_id_column=shared_id_column,
    )


def _read_direct_huggingface_manifest(path: Path) -> dict[str, Any] | None:
    if path.suffix.lower() != ".json":
        return None
    payload = json.loads(path.read_text(encoding="utf-8"))
    if (
        not isinstance(payload, dict)
        or "dataset" not in payload
        or "configs" not in payload
    ):
        return None
    return cast(dict[str, Any], payload)


def _direct_materialized_path(
    root: Path, config_name: str, fraction: float, seed: int
) -> Path:
    identity = f"{config_name}\0{fraction:.12g}\0{seed}".encode()
    digest = hashlib.sha256(identity).hexdigest()[:16]
    readable = config_name[:80].replace("/", "_")
    return root / "direct-huggingface" / f"{readable}-{digest}"


def _prepare_direct_training_dataset(
    dataset: Dataset,
    *,
    task: str,
    shared_id_column: str,
) -> Dataset:
    """Expose model text columns while retaining shared IDs for provenance."""
    if "anchor" in dataset.column_names:
        text_columns = ["anchor", "positive"]
        rename_columns: dict[str, str] = {}
        if task == "triplet":
            text_columns.append("negative")
    else:
        text_columns = ["query", "pos"]
        rename_columns = {"query": "anchor", "pos": "positive"}
        if task == "triplet":
            text_columns.append("neg")
            rename_columns["neg"] = "negative"
    missing = set(text_columns).difference(dataset.column_names)
    if missing:
        raise ValueError(f"direct dataset is missing text columns: {sorted(missing)}")
    if shared_id_column not in dataset.column_names:
        # Backward compatibility for previously materialized smoke/search slices.
        prepared = dataset.select_columns(text_columns)
        return prepared.rename_columns(rename_columns) if rename_columns else prepared
    source_ids = dataset.select_columns([shared_id_column])
    prepared = dataset.select_columns(text_columns)
    if rename_columns:
        prepared = prepared.rename_columns(rename_columns)
    metadata_target = cast(Any, prepared)
    metadata_target._bekko_shared_id_source = source_ids
    metadata_target._bekko_shared_id_column = shared_id_column
    return prepared


def load_direct_huggingface_datasets(
    *,
    manifest: dict[str, Any],
    fraction: float,
    seed: int,
    require_cache: bool,
    shared_id_block_size: int,
    materialized_root: Path | None,
    source_cache_direct: bool = False,
    shared_id_column: str = "shared_id",
) -> tuple[DatasetDict, Manifest, dict[str, Literal["query_doc", "doc_doc"]]]:
    """Load a deterministic smoke slice directly from the published Bekko dataset.

    The released rows are already shuffled in shared-id blocks. Taking an aligned
    prefix therefore avoids scanning or downloading the remaining 99% while
    preserving batches that are safe for in-batch duplicate filtering.
    """
    if materialized_root is None and not source_cache_direct:
        raise ValueError("direct Hugging Face manifests require data.stage_cache_root")
    if source_cache_direct and fraction != 1.0:
        raise ValueError(
            "direct source-cache loading currently requires data.fraction: 1.0"
        )
    dataset_id = str(manifest["dataset"])
    revision = str(manifest["revision"])
    raw_configs = manifest["configs"]
    if not isinstance(raw_configs, list):
        raise TypeError("direct manifest configs must be a list")

    loaded: dict[str, Dataset] = {}
    data_manifest: Manifest = {}
    pair_types: dict[str, Literal["query_doc", "doc_doc"]] = {}
    config_count = len(raw_configs)
    for config_index, raw in enumerate(raw_configs, start=1):
        if not isinstance(raw, dict):
            raise TypeError("direct manifest config entries must be mappings")
        config_name = str(raw["name"])
        source_rows = int(raw["source_rows_all_splits"])
        task = str(raw["task"])
        pair_type = str(raw["pair_type"])
        if task not in {"pair", "triplet"} or pair_type not in {"query_doc", "doc_doc"}:
            raise ValueError(f"invalid direct manifest metadata for {config_name}")
        selected_rows = min(
            source_rows,
            max(
                shared_id_block_size,
                math.ceil(source_rows * fraction / shared_id_block_size)
                * shared_id_block_size,
            ),
        )
        cache_path = (
            _direct_materialized_path(
                materialized_root.resolve(), config_name, fraction, seed
            )
            if materialized_root is not None
            else None
        )
        if source_cache_direct:
            source = load_dataset(
                dataset_id,
                config_name,
                revision=revision,
                split="train",
                streaming=False,
                cache_dir=(
                    str(materialized_root.resolve())
                    if materialized_root is not None
                    else None
                ),
            )
            if not isinstance(source, Dataset):
                raise TypeError(f"direct source config is not a Dataset: {config_name}")
            if len(source) != source_rows:
                raise ValueError(
                    f"direct source row count mismatch for {config_name}: "
                    f"manifest={source_rows} actual={len(source)}"
                )
            required_columns = {"query", "pos", shared_id_column}
            if task == "triplet":
                required_columns.add("neg")
            missing = required_columns.difference(source.column_names)
            if missing:
                raise ValueError(
                    f"direct source {config_name} is missing columns: {sorted(missing)}"
                )
            dataset = _prepare_direct_training_dataset(
                source,
                task=task,
                shared_id_column=shared_id_column,
            )
            metadata_target = cast(Any, dataset)
            metadata_target._bekko_direct_source_cache = True
        elif cache_path is not None and cache_path.is_dir():
            print(
                "Loading materialized Bekko cache "
                f"{config_index}/{config_count}: {config_name}",
                flush=True,
            )
            dataset = load_from_disk(cache_path)
            if not isinstance(dataset, Dataset):
                raise TypeError(
                    f"materialized direct dataset is not a Dataset: {cache_path}"
                )
        else:
            if require_cache:
                raise FileNotFoundError(
                    f"required direct dataset slice is missing: {cache_path}"
                )
            stream = load_dataset(
                dataset_id,
                config_name,
                revision=revision,
                split="train",
                streaming=True,
            )

            def rows(
                source: Any = stream,
                count: int = selected_rows,
                row_task: str = task,
            ) -> Any:
                for row in source.take(count):
                    output = {
                        "anchor": row["query"],
                        "positive": row["pos"],
                        shared_id_column: row[shared_id_column],
                    }
                    if row_task == "triplet":
                        output["negative"] = row["neg"]
                    yield output

            dataset = Dataset.from_generator(
                rows,
                cache_dir=str(materialized_root.resolve() / "direct-builder-cache"),
            )
            del rows
            del stream
            gc.collect()
            assert cache_path is not None
            cache_path.parent.mkdir(parents=True, exist_ok=True)
            dataset.save_to_disk(cache_path)
            reopened = load_from_disk(cache_path)
            if not isinstance(reopened, Dataset):
                raise TypeError(
                    f"materialized direct dataset is not a Dataset: {cache_path}"
                )
            dataset = reopened

        dataset = _prepare_direct_training_dataset(
            dataset,
            task=task,
            shared_id_column=shared_id_column,
        )

        key = config_name
        metadata_target = cast(Any, dataset)
        metadata_target._bekko_component_name = key
        metadata_target._bekko_cache_path = str(cache_path) if cache_path else None
        loaded[key] = dataset
        pair_types[key] = pair_type
        data_manifest[key] = {
            "source_rows": source_rows,
            "selected_rows": len(dataset),
            "pair_type": pair_type,
            "source_dataset": dataset_id,
            "source_config": config_name,
            "source_revision": revision,
        }
    return DatasetDict(loaded), data_manifest, pair_types  # ty: ignore[no-matching-overload]


def load_huggingface_supplements(
    *,
    supplements: tuple[SupplementalDataSettings, ...],
    seed: int,
    shared_id_column: str,
    shared_id_block_size: int,
) -> tuple[
    dict[str, Dataset],
    Manifest,
    dict[str, Literal["query_doc", "doc_doc"]],
]:
    loaded: dict[str, Dataset] = {}
    manifest: Manifest = {}
    pair_types: dict[str, Literal["query_doc", "doc_doc"]] = {}
    for settings in supplements:
        source = load_dataset(
            settings.dataset,
            settings.config,
            revision=settings.revision,
            split=settings.split,
        )
        if not isinstance(source, Dataset):
            raise TypeError(f"Supplement {settings.name!r} did not return a Dataset")
        selected = select_fraction(
            source,
            fraction=settings.fraction,
            seed=seed,
            shared_id_column=shared_id_column,
            shared_id_block_size=shared_id_block_size,
        )
        key = f"supplement__{settings.name}".replace("/", "_")
        loaded[key] = adapt_supplemental_dataset(selected, settings=settings)
        manifest[key] = {
            "source_rows": len(source),
            "selected_rows": len(selected),
            "pair_type": settings.pair_type,
            "dataset": settings.dataset,
            "config": settings.config,
            "revision": settings.revision,
            "provenance_note": settings.provenance_note,
            "batch_deduplication": settings.batch_deduplication,
        }
        pair_types[key] = settings.pair_type
    return loaded, manifest, pair_types


def build_loss_eval_dataset(train_dataset: DatasetDict, *, limit: int = 256) -> Dataset:
    first = next(iter(train_dataset.values()))
    return first.select(range(min(limit, len(first))))


# ==============================================================================
# Representation-aware batching
# ==============================================================================

from array import array
from collections.abc import Mapping
from dataclasses import dataclass

import xxhash
from torch.utils.data import BatchSampler

REPRESENTATION_BATCH_PLAN_ATTR = "_bekko_representation_batch_plan"
TEXT_COLUMNS = ("anchor", "positive", "negative")


def _tokenizer_file(path: Path) -> Path:
    resolved = path / "tokenizer.json" if path.is_dir() else path
    if not resolved.is_file():
        raise FileNotFoundError(f"Tokenizer JSON not found: {resolved}")
    return resolved


def _token_ids_fingerprint(token_ids: list[int]) -> tuple[np.uint64, np.uint64]:
    payload = array("I", token_ids).tobytes()
    digest = xxhash.xxh3_128_digest(payload)
    values = np.frombuffer(digest, dtype=np.dtype(">u8"))
    return np.uint64(values[0]), np.uint64(values[1])


def build_model_input_fingerprints(
    dataset: Dataset,
    *,
    tokenizer_path: Path,
    text_preprocessing: TextPreprocessing,
    max_seq_length: int,
    batch_size: int = 8192,
    output_path: Path | None = None,
) -> np.ndarray:
    """Fingerprint the exact lexical representation consumed by StaticEmbedding."""
    if max_seq_length <= 0 or batch_size <= 0:
        raise ValueError("max_seq_length and batch_size must be positive")
    columns = [column for column in TEXT_COLUMNS if column in dataset.column_names]
    if len(columns) < 2:
        raise ValueError(
            "representation deduplication requires at least two text columns"
        )
    tokenizer = Tokenizer.from_file(str(_tokenizer_file(tokenizer_path)))
    tokenizer.enable_truncation(max_length=max_seq_length)
    shape = (len(dataset), len(columns), 2)
    if output_path is None:
        fingerprints = np.empty(shape, dtype=np.uint64)
    else:
        output_path.parent.mkdir(parents=True, exist_ok=True)
        fingerprints = np.lib.format.open_memmap(
            output_path,
            mode="w+",
            dtype=np.uint64,
            shape=shape,
        )

    offset = 0
    for batch in dataset.iter(batch_size=batch_size):
        row_count = len(batch[columns[0]])
        for column_index, column in enumerate(columns):
            raw_texts = [str(text) for text in batch[column]]
            unique_raw_texts = list(dict.fromkeys(raw_texts))
            transformed = [
                preprocess_text(text, text_preprocessing) for text in unique_raw_texts
            ]
            encodings = tokenizer.encode_batch(transformed, add_special_tokens=False)
            by_raw_text = {
                raw_text: _token_ids_fingerprint(encoding.ids)
                for raw_text, encoding in zip(unique_raw_texts, encodings, strict=True)
            }
            for row_index, raw_text in enumerate(raw_texts):
                fingerprints[offset + row_index, column_index] = by_raw_text[raw_text]
        offset += row_count
    if offset != len(dataset):
        raise RuntimeError(f"fingerprinted {offset} rows, expected {len(dataset)}")
    if isinstance(fingerprints, np.memmap):
        fingerprints.flush()
    return fingerprints


@dataclass(frozen=True)
class RepresentationBatchPlan:
    indices: np.ndarray
    offsets: np.ndarray
    source_rows: int
    batch_size: int
    seed: int

    def __post_init__(self) -> None:
        if self.indices.ndim != 1 or self.offsets.ndim != 1:
            raise ValueError("plan indices and offsets must be one-dimensional")
        if len(self.offsets) == 0 or int(self.offsets[0]) != 0:
            raise ValueError("plan offsets must begin at zero")
        if int(self.offsets[-1]) != len(self.indices):
            raise ValueError("plan offsets must end at the number of indices")
        if len(self.indices) != self.source_rows:
            raise ValueError("representation plan must retain every source row")

    def __len__(self) -> int:
        return len(self.offsets) - 1

    def batches(self) -> list[np.ndarray]:
        return [
            self.indices[int(start) : int(end)]
            for start, end in zip(self.offsets[:-1], self.offsets[1:], strict=True)
        ]

    def to_manifest(self) -> dict[str, int]:
        sizes = np.diff(self.offsets)
        return {
            "source_rows": self.source_rows,
            "retained_rows": len(self.indices),
            "dropped_rows": self.source_rows - len(self.indices),
            "batch_size": self.batch_size,
            "batch_count": len(self),
            "partial_batch_count": int(np.count_nonzero(sizes < self.batch_size)),
            "smallest_batch": int(sizes.min()) if len(sizes) else 0,
            "seed": self.seed,
        }

    def save(self, directory: Path) -> None:
        directory.mkdir(parents=True, exist_ok=True)
        np.save(directory / "plan_indices.npy", self.indices, allow_pickle=False)
        np.save(directory / "plan_offsets.npy", self.offsets, allow_pickle=False)
        (directory / "plan.json").write_text(
            json.dumps(self.to_manifest(), indent=2, sort_keys=True) + "\n",
            encoding="utf-8",
        )

    @classmethod
    def load(cls, directory: Path) -> RepresentationBatchPlan:
        metadata = json.loads((directory / "plan.json").read_text(encoding="utf-8"))
        return cls(
            indices=np.load(directory / "plan_indices.npy", mmap_mode="r"),
            offsets=np.load(directory / "plan_offsets.npy", mmap_mode="r"),
            source_rows=int(metadata["source_rows"]),
            batch_size=int(metadata["batch_size"]),
            seed=int(metadata["seed"]),
        )


def build_representation_batch_plan(
    fingerprints: np.ndarray,
    *,
    batch_size: int,
    seed: int,
) -> RepresentationBatchPlan:
    """Greedily color rows into full, collision-free batches without dropping rows."""
    if fingerprints.ndim != 3 or fingerprints.shape[2] != 2:
        raise ValueError("fingerprints must have shape (rows, columns, 2)")
    if fingerprints.dtype != np.uint64:
        raise ValueError("fingerprints must use uint64")
    if batch_size <= 0:
        raise ValueError("batch_size must be positive")
    row_count = len(fingerprints)
    index_dtype = np.int32 if row_count <= np.iinfo(np.int32).max else np.int64
    generator = torch.Generator().manual_seed(seed)
    shuffled = torch.randperm(row_count, generator=generator, dtype=torch.int64).numpy()
    shuffled = shuffled.astype(index_dtype, copy=False)
    next_dtype = np.int32 if row_count + 1 <= np.iinfo(np.int32).max else np.int64
    next_positions = np.arange(1, row_count + 1, dtype=next_dtype)
    if row_count:
        next_positions[-1] = -1
    planned = np.empty(row_count, dtype=index_dtype)
    offsets = [0]
    write_offset = 0
    head_position = 0 if row_count else -1

    while head_position != -1:
        batch_values: set[tuple[int, int]] = set()
        current_position = head_position
        previous_position = -1
        batch_start = write_offset
        while current_position != -1:
            next_position = int(next_positions[current_position])
            index = int(shuffled[current_position])
            row_values = {
                (int(high), int(low)) for high, low in fingerprints[index].tolist()
            }
            if not row_values.isdisjoint(batch_values):
                previous_position = current_position
                current_position = next_position
                continue

            planned[write_offset] = index
            write_offset += 1
            batch_values.update(row_values)
            if previous_position == -1:
                head_position = next_position
            else:
                next_positions[previous_position] = next_position
            if write_offset - batch_start == batch_size:
                break
            current_position = next_position

        if write_offset == batch_start:
            raise RuntimeError("unable to make progress while constructing batch plan")
        offsets.append(write_offset)

    if write_offset != row_count:
        raise RuntimeError(f"planned {write_offset} rows, expected {row_count}")
    return RepresentationBatchPlan(
        indices=planned,
        offsets=np.asarray(offsets, dtype=np.int64),
        source_rows=row_count,
        batch_size=batch_size,
        seed=seed,
    )


def attach_representation_batch_plan(
    dataset: Dataset, plan: RepresentationBatchPlan
) -> None:
    if len(dataset) != plan.source_rows:
        raise ValueError("dataset and representation batch plan row counts differ")
    setattr(dataset, REPRESENTATION_BATCH_PLAN_ATTR, plan)


def representation_batch_plan(dataset: Any) -> RepresentationBatchPlan | None:
    value = getattr(dataset, REPRESENTATION_BATCH_PLAN_ATTR, None)
    return value if isinstance(value, RepresentationBatchPlan) else None


class RepresentationNoDuplicatesBatchSampler(BatchSampler):
    """Yield a pre-audited model-input-unique plan, shuffling only batch order."""

    def __init__(
        self,
        dataset: Dataset,
        batch_size: int,
        drop_last: bool,
        valid_label_columns: list[str] | None = None,
        generator: torch.Generator | None = None,
        seed: int = 0,
    ) -> None:
        del valid_label_columns
        super().__init__(dataset, batch_size=batch_size, drop_last=drop_last)
        plan = representation_batch_plan(dataset)
        if plan is None:
            raise ValueError("dataset has no representation batch plan")
        if batch_size != plan.batch_size:
            raise ValueError(
                f"batch size {batch_size} differs from planned size {plan.batch_size}"
            )
        if drop_last:
            raise ValueError(
                "representation-safe training does not permit dropping rows"
            )
        self.plan = plan
        self.generator = generator
        self.seed = seed
        self.epoch = 0

    def set_epoch(self, epoch: int) -> None:
        self.epoch = epoch

    def __iter__(self):  # type: ignore[no-untyped-def]
        generator = self.generator or torch.Generator()
        generator.manual_seed(self.seed + self.epoch)
        order = torch.randperm(len(self.plan), generator=generator).tolist()
        batches = self.plan.batches()
        for batch_index in order:
            yield batches[batch_index].tolist()

    def __len__(self) -> int:
        return len(self.plan)


def tokenizer_sha256(tokenizer_path: Path) -> str:
    return hashlib.sha256(_tokenizer_file(tokenizer_path).read_bytes()).hexdigest()


def _identity_digest(identity: dict[str, Any]) -> str:
    payload = json.dumps(identity, sort_keys=True, separators=(",", ":")).encode()
    return hashlib.sha256(payload).hexdigest()[:20]


def _cache_identity(
    dataset: Dataset,
    *,
    tokenizer_path: Path,
    text_preprocessing: TextPreprocessing,
    max_seq_length: int,
) -> dict[str, Any]:
    return {
        "schema_version": 1,
        "supplement": getattr(dataset, "_bekko_supplement_identity", {}),
        "dataset_fingerprint": getattr(dataset, "_fingerprint", None),
        "rows": len(dataset),
        "columns": list(dataset.column_names),
        "tokenizer_sha256": tokenizer_sha256(tokenizer_path),
        "text_preprocessing": text_preprocessing,
        "max_seq_length": max_seq_length,
        "hash": "xxh3-128-token-ids-v1",
    }


def _load_or_build_fingerprints(
    dataset: Dataset,
    *,
    cache_root: Path,
    tokenizer_path: Path,
    text_preprocessing: TextPreprocessing,
    max_seq_length: int,
    fingerprint_batch_size: int,
) -> tuple[np.ndarray, Path, dict[str, Any]]:
    identity = _cache_identity(
        dataset,
        tokenizer_path=tokenizer_path,
        text_preprocessing=text_preprocessing,
        max_seq_length=max_seq_length,
    )
    cache_dir = cache_root / _identity_digest(identity)
    metadata_path = cache_dir / "fingerprints.json"
    fingerprints_path = cache_dir / "fingerprints.npy"
    if metadata_path.is_file() and fingerprints_path.is_file():
        cached_identity = json.loads(metadata_path.read_text(encoding="utf-8"))
        if cached_identity == identity:
            fingerprints = np.load(fingerprints_path, mmap_mode="r")
            expected_shape = (len(dataset), len(dataset.column_names), 2)
            if fingerprints.shape == expected_shape and fingerprints.dtype == np.uint64:
                return fingerprints, cache_dir, identity

    cache_dir.mkdir(parents=True, exist_ok=True)
    fingerprints = build_model_input_fingerprints(
        dataset,
        tokenizer_path=tokenizer_path,
        text_preprocessing=text_preprocessing,
        max_seq_length=max_seq_length,
        batch_size=fingerprint_batch_size,
        output_path=fingerprints_path,
    )
    metadata_path.write_text(
        json.dumps(identity, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    return fingerprints, cache_dir, identity


def audit_representation_batch_plan(
    fingerprints: np.ndarray,
    plan: RepresentationBatchPlan,
) -> dict[str, int]:
    collisions = 0
    audited_rows = 0
    observed_indices: list[np.ndarray] = []
    for batch in plan.batches():
        seen: set[tuple[int, int]] = set()
        observed_indices.append(batch)
        for index in batch:
            row_values = {
                (int(high), int(low)) for high, low in fingerprints[int(index)].tolist()
            }
            collisions += len(seen.intersection(row_values))
            seen.update(row_values)
            audited_rows += 1
    flattened = np.concatenate(observed_indices) if observed_indices else np.asarray([])
    unique_indices = len(np.unique(flattened))
    if (
        collisions
        or audited_rows != plan.source_rows
        or unique_indices != plan.source_rows
    ):
        raise RuntimeError(
            "representation batch audit failed: "
            f"collisions={collisions}, rows={audited_rows}, unique_indices={unique_indices}, "
            f"expected={plan.source_rows}"
        )
    return {
        "audited_rows": audited_rows,
        "unique_indices": unique_indices,
        "representation_collisions": collisions,
    }


def prepare_representation_batch_plans(
    datasets: Mapping[Any, Dataset],
    *,
    tokenizer_path: Path,
    text_preprocessing: TextPreprocessing,
    max_seq_length: int,
    batch_size: int,
    seed: int,
    default_cache_root: Path,
    fingerprint_batch_size: int = 8192,
) -> dict[str, dict[str, int | str]]:
    manifests: dict[str, dict[str, int | str]] = {}
    for raw_name, dataset in datasets.items():
        name = str(raw_name)
        if getattr(dataset, "_bekko_batch_deduplication", "text") != "model_input":
            continue
        configured_cache = getattr(dataset, "_bekko_fingerprint_cache", None)
        cache_root = (
            Path(configured_cache) if configured_cache else default_cache_root / name
        )
        fingerprints, cache_dir, identity = _load_or_build_fingerprints(
            dataset,
            cache_root=cache_root,
            tokenizer_path=tokenizer_path,
            text_preprocessing=text_preprocessing,
            max_seq_length=max_seq_length,
            fingerprint_batch_size=fingerprint_batch_size,
        )
        plan_dir = cache_dir / f"plan-bs{batch_size}-seed{seed}"
        if (plan_dir / "plan.json").is_file():
            plan = RepresentationBatchPlan.load(plan_dir)
            if plan.batch_size != batch_size or plan.seed != seed:
                raise ValueError(f"cached plan settings do not match: {plan_dir}")
        else:
            plan = build_representation_batch_plan(
                fingerprints,
                batch_size=batch_size,
                seed=seed,
            )
            plan.save(plan_dir)
        audit = audit_representation_batch_plan(fingerprints, plan)
        attach_representation_batch_plan(dataset, plan)
        manifests[name] = {
            **plan.to_manifest(),
            **audit,
            "fingerprint_cache": str(cache_dir),
            "plan_cache": str(plan_dir),
            "fingerprint_identity": _identity_digest(identity),
        }
    return manifests


# ==============================================================================
# Locality-aware batching
# ==============================================================================

import warnings
from collections.abc import Iterator

from sentence_transformers.base.sampler import NoDuplicatesBatchSampler

_EXCLUDED_COLUMNS = {"dataset_name"}


def block_local_permutation(
    num_rows: int,
    *,
    block_size: int,
    generator: torch.Generator | None,
) -> np.ndarray:
    """Shuffle blocks and rows within blocks without interleaving their pages."""
    if num_rows < 0:
        raise ValueError("num_rows must be non-negative")
    if block_size <= 0:
        raise ValueError("block_size must be positive")
    if num_rows == 0:
        return np.empty(0, dtype=np.int32)
    index_dtype = torch.int32 if num_rows <= np.iinfo(np.int32).max else torch.int64
    block_count = math.ceil(num_rows / block_size)
    block_order = torch.randperm(block_count, generator=generator, dtype=index_dtype)
    pieces: list[torch.Tensor] = []
    for block_index in block_order.tolist():
        start = int(block_index) * block_size
        length = min(block_size, num_rows - start)
        pieces.append(
            torch.randperm(length, generator=generator, dtype=index_dtype).add_(start)
        )
    return torch.cat(pieces).numpy()


class LocalityAwareNoDuplicatesBatchSampler(NoDuplicatesBatchSampler):
    """No-duplicates batching with bounded random-access working sets."""

    def __init__(
        self,
        dataset: Dataset,
        batch_size: int,
        drop_last: bool,
        valid_label_columns: list[str] | None = None,
        generator: torch.Generator | None = None,
        seed: int = 0,
        *,
        locality_block_size: int,
    ) -> None:
        super().__init__(
            dataset,
            batch_size=batch_size,
            drop_last=drop_last,
            valid_label_columns=valid_label_columns,
            generator=generator,
            seed=seed,
        )
        if locality_block_size < batch_size:
            raise ValueError("locality_block_size must be at least batch_size")
        self.locality_block_size = locality_block_size

    def __iter__(self) -> Iterator[list[int]]:
        if self.generator and self.seed is not None:
            self.generator.manual_seed(self.seed + self.epoch)
        num_rows = len(self.dataset)
        if num_rows == 0:
            return
        remaining_indices = block_local_permutation(
            num_rows,
            block_size=self.locality_block_size,
            generator=self.generator,
        )
        position_dtype = (
            np.int32 if num_rows + 1 <= np.iinfo(np.int32).max else np.int64
        )
        next_positions = np.arange(1, num_rows + 1, dtype=position_dtype)
        next_positions[-1] = -1
        head_position = 0

        while head_position != -1:
            batch_values: set[str] = set()
            batch_indices: list[int] = []
            current_position = head_position
            previous_position = -1
            full_batch = False
            while current_position != -1:
                next_position = int(next_positions[current_position])
                index = int(remaining_indices[current_position])
                try:
                    row = self.dataset[index]
                except UnicodeDecodeError as error:
                    # A malformed UTF-8 scalar cannot be collated or tokenized. Remove
                    # that row from the sampler's linked list and retain every other
                    # row instead of losing a multi-hour full-corpus run.
                    component = getattr(
                        self.dataset,
                        "_bekko_component_name",
                        getattr(self.dataset, "_fingerprint", "unknown"),
                    )
                    cache_path = getattr(self.dataset, "_bekko_cache_path", None)
                    location = f" component={component!r} row={index}"
                    if cache_path is not None:
                        location += f" cache={cache_path!r}"
                    warnings.warn(
                        f"Skipping malformed UTF-8 training row:{location}: {error}",
                        RuntimeWarning,
                        stacklevel=2,
                    )
                    if previous_position == -1:
                        head_position = next_position
                    else:
                        next_positions[previous_position] = next_position
                    current_position = next_position
                    continue
                sample_values = {
                    str(value)
                    for key, value in row.items()
                    if key not in _EXCLUDED_COLUMNS
                }
                if not sample_values.isdisjoint(batch_values):
                    previous_position = current_position
                    current_position = next_position
                    continue

                batch_indices.append(index)
                if previous_position == -1:
                    head_position = next_position
                else:
                    next_positions[previous_position] = next_position
                if len(batch_indices) == self.batch_size:
                    full_batch = True
                    yield batch_indices
                    break
                batch_values.update(sample_values)
                current_position = next_position

            if not full_batch and not self.drop_last and batch_indices:
                yield batch_indices


# ==============================================================================
# Step scheduling
# ==============================================================================

from dataclasses import dataclass


@dataclass(frozen=True)
class StepSchedule:
    total_steps: int
    eval_steps: int
    save_steps: int


def derive_step_schedule(
    *,
    dataset_rows: int,
    batch_size: int,
    epochs: float,
    target_evaluations: int,
    exact_target_evaluations: bool = False,
) -> StepSchedule:
    if dataset_rows <= 0:
        raise ValueError("dataset_rows must be positive")
    if batch_size <= 0 or epochs <= 0 or target_evaluations <= 0:
        raise ValueError("batch_size, epochs, and target_evaluations must be positive")
    total_steps = math.ceil(dataset_rows / batch_size) * math.ceil(epochs)
    # A floor interval guarantees that a non-divisible run reaches the requested
    # number of evaluations. Keep the historical ceil behavior unless a config
    # explicitly asks for an exact count.
    if exact_target_evaluations:
        eval_interval = max(1, total_steps // target_evaluations)
    else:
        eval_interval = max(1, math.ceil(total_steps / target_evaluations))
    save_interval = max(1, math.ceil(total_steps / 10))
    return StepSchedule(
        total_steps=total_steps,
        eval_steps=eval_interval,
        save_steps=save_interval,
    )


def derive_multidataset_step_schedule(
    *,
    component_rows: list[int],
    batch_size: int,
    epochs: float,
    target_evaluations: int,
    exact_target_evaluations: bool = False,
) -> StepSchedule:
    positive_rows = [rows for rows in component_rows if rows > 0]
    if not positive_rows:
        raise ValueError("component_rows must contain at least one positive row count")
    rows_as_batches = sum(math.ceil(rows / batch_size) for rows in positive_rows)
    return derive_step_schedule(
        dataset_rows=rows_as_batches * batch_size,
        batch_size=batch_size,
        epochs=epochs,
        target_evaluations=target_evaluations,
        exact_target_evaluations=exact_target_evaluations,
    )


# ==============================================================================
# Qwen3 embedding loss
# ==============================================================================

from collections.abc import Iterable

from sentence_transformers import util
from torch import Tensor, nn


def fake_quantize_embeddings(embeddings: Tensor, precision: str) -> Tensor:
    """Fake-quantize sentence embeddings with straight-through gradients."""
    if precision == "float32":
        return embeddings
    if precision == "int8":
        mins = embeddings.amin(dim=0, keepdim=True)
        maxs = embeddings.amax(dim=0, keepdim=True)
        steps = torch.clamp((maxs - mins) / 255.0, min=1e-8)
        normalized = (embeddings - mins) / steps
        quantized = normalized + (torch.round(normalized) - normalized).detach()
        quantized = torch.clamp(quantized - 128, -128, 127)
        return (quantized + 128) * steps + mins
    if precision == "binary":
        quantized = torch.sign(embeddings)
        quantized = torch.where(quantized == 0, torch.ones_like(quantized), quantized)
        return embeddings + (quantized - embeddings).detach()
    raise ValueError(f"Unsupported fake-quantization precision: {precision}")


class Qwen3EmbeddingLoss(nn.Module):
    """Qwen3-Embedding improved InfoNCE with false-negative masking."""

    def __init__(
        self,
        model: SentenceTransformer,
        *,
        temperature: float = 0.01,
        margin: float | None = 0.1,
        hard_negative_margin: float | None = None,
        bidirectional: bool = False,
        query_query_weight: float = 1.0,
        document_document_weight: float = 1.0,
        cross_document_weight: float = 1.0,
        hard_negative_weight: float = 1.0,
        qat_int8_weight: float = 0.0,
        qat_binary_weight: float = 0.0,
        qat_n_quantized_per_step: int = -1,
    ) -> None:
        super().__init__()
        if temperature <= 0:
            raise ValueError("temperature must be > 0")
        weights = {
            "query_query_weight": query_query_weight,
            "document_document_weight": document_document_weight,
            "cross_document_weight": cross_document_weight,
            "hard_negative_weight": hard_negative_weight,
        }
        for name, value in weights.items():
            if not math.isfinite(value) or value < 0:
                raise ValueError(f"{name} must be finite and non-negative")
        qat_weights = {"int8": qat_int8_weight, "binary": qat_binary_weight}
        for name, value in qat_weights.items():
            if not math.isfinite(value) or value < 0:
                raise ValueError(f"qat_{name}_weight must be finite and non-negative")
        if qat_n_quantized_per_step not in {-1, 1, 2}:
            raise ValueError("qat_n_quantized_per_step must be -1, 1, or 2")
        self.model = model
        self.temperature = temperature
        self.margin = margin
        self.hard_negative_margin = hard_negative_margin
        self.bidirectional = bidirectional
        self.query_query_weight = query_query_weight
        self.document_document_weight = document_document_weight
        self.cross_document_weight = cross_document_weight
        self.hard_negative_weight = hard_negative_weight
        self.qat_weights = qat_weights
        self.qat_n_quantized_per_step = qat_n_quantized_per_step
        self.last_loss_components: dict[str, Tensor] = {}

    @property
    def scale(self) -> float:
        return 1.0 / self.temperature

    def compute_loss_from_embeddings(self, embeddings: list[Tensor]) -> Tensor:
        if len(embeddings) < 2:
            raise ValueError(
                f"Expected at least two embedding tensors, got {len(embeddings)}"
            )
        queries, positives, *negative_sets = embeddings
        if queries.shape != positives.shape:
            raise ValueError("Query and positive embedding shapes must match")
        if any(negative.shape != queries.shape for negative in negative_sets):
            raise ValueError(
                "Every hard-negative embedding shape must match the queries"
            )

        float_loss = self._compute_loss_for_precision(queries, positives, negative_sets)
        components = {"float32": float_loss}
        enabled = [name for name, weight in self.qat_weights.items() if weight > 0]
        if 0 < self.qat_n_quantized_per_step < len(enabled):
            enabled = random.sample(enabled, self.qat_n_quantized_per_step)
        total = float_loss
        for precision in enabled:
            transformed = [
                fake_quantize_embeddings(value, precision)
                for value in (queries, positives, *negative_sets)
            ]
            quantized_loss = self._compute_loss_for_precision(
                transformed[0], transformed[1], transformed[2:]
            )
            components[precision] = quantized_loss
            total = total + self.qat_weights[precision] * quantized_loss
        dimension = queries.shape[-1]
        self.last_loss_components.update(
            {
                f"{name}_d{dimension}": value.detach()
                for name, value in components.items()
            }
        )
        return total

    def _compute_loss_for_precision(
        self,
        queries: Tensor,
        positives: Tensor,
        negative_sets: list[Tensor],
    ) -> Tensor:
        forward_loss = self._compute_directional_loss(queries, positives, negative_sets)
        if not self.bidirectional:
            return forward_loss
        reverse_loss = self._compute_directional_loss(positives, queries, [])
        return 0.5 * (forward_loss + reverse_loss)

    def _compute_directional_loss(
        self,
        queries: Tensor,
        positives: Tensor,
        negative_sets: list[Tensor],
    ) -> Tensor:
        batch_size = len(queries)
        indices = torch.arange(batch_size, device=queries.device)
        off_diagonal = ~torch.eye(batch_size, dtype=torch.bool, device=queries.device)

        query_positive = util.cos_sim(queries, positives)
        query_query = util.cos_sim(queries, queries)
        positive_positive = util.cos_sim(positives, positives)
        positive_scores = query_positive[indices, indices]
        threshold = (
            None if self.margin is None else positive_scores[:, None] + self.margin
        )

        denominator_parts = [positive_scores[:, None] * self.scale]
        for similarities, weight in zip(
            (query_query, positive_positive, query_positive),
            (
                self.query_query_weight,
                self.document_document_weight,
                self.cross_document_weight,
            ),
            strict=True,
        ):
            keep = (
                off_diagonal
                if threshold is None
                else off_diagonal & (similarities <= threshold)
            )
            denominator_parts.append(self._weighted_logits(similarities, keep, weight))

        for negatives in negative_sets:
            similarities = util.cos_sim(queries, negatives)
            if threshold is None:
                keep = torch.ones_like(similarities, dtype=torch.bool)
            else:
                keep = similarities <= threshold
            if self.hard_negative_margin is None:
                keep[indices, indices] = True
            else:
                owner_threshold = positive_scores + self.hard_negative_margin
                keep[indices, indices] = (
                    similarities[indices, indices] <= owner_threshold
                )
            denominator_parts.append(
                self._weighted_logits(similarities, keep, self.hard_negative_weight)
            )

        logits = torch.cat(denominator_parts, dim=1)
        return -(positive_scores * self.scale - torch.logsumexp(logits, dim=1)).mean()

    def _weighted_logits(
        self, similarities: Tensor, keep: Tensor, weight: float
    ) -> Tensor:
        if weight == 0:
            return torch.full_like(similarities, -torch.inf)
        return torch.where(
            keep, similarities * self.scale + math.log(weight), -torch.inf
        )

    def forward(
        self, sentence_features: Iterable[dict[str, Tensor]], labels: Tensor
    ) -> Tensor:
        del labels
        embeddings = [
            self.model(features)["sentence_embedding"] for features in sentence_features
        ]
        return self.compute_loss_from_embeddings(embeddings)

    def get_config_dict(self) -> dict[str, Any]:
        return {
            "temperature": self.temperature,
            "margin": self.margin,
            "hard_negative_margin": self.hard_negative_margin,
            "bidirectional": self.bidirectional,
            "query_query_weight": self.query_query_weight,
            "document_document_weight": self.document_document_weight,
            "cross_document_weight": self.cross_document_weight,
            "hard_negative_weight": self.hard_negative_weight,
            "qat_int8_weight": self.qat_weights["int8"],
            "qat_binary_weight": self.qat_weights["binary"],
            "qat_n_quantized_per_step": self.qat_n_quantized_per_step,
        }


# ==============================================================================
# Hakari evaluation
# ==============================================================================

import importlib
import sys
from collections import defaultdict
from collections.abc import Sequence
from types import ModuleType
from typing import Literal

from sentence_transformers.sentence_transformer.evaluation import SequentialEvaluator


def default_hakari_targets() -> list[dict[str, list[str] | str]]:
    return [
        {"dataset": "NanoBEIR-en", "splits": []},
        {"dataset": "NanoMIRACL", "splits": ["en"]},
    ]


def import_hakari(repo: str | Path) -> ModuleType:
    path = Path(repo).resolve()
    if not (path / "hakari_bench").is_dir():
        raise FileNotFoundError(f"hakari-bench checkout not found: {path}")
    path_text = str(path)
    if path_text not in sys.path:
        sys.path.insert(0, path_text)
    hakari = importlib.import_module("hakari_bench.sentence_transformers")
    _install_fast_hakari_candidate_index()
    return hakari


def _install_fast_hakari_candidate_index() -> None:
    """Use Arrow kernels when HAKARI restricts a corpus to reranking candidates.

    HAKARI's portable fallback iterates every corpus ID in Python to construct a
    candidate-only lazy corpus.  That is correct, but the first small mid-run
    evaluation can spend minutes scanning large multilingual Nano corpora.  A
    Hugging Face ``Dataset`` normally exposes its ID column as a PyArrow table,
    where ``is_in`` and ``indices_nonzero`` build the same ID-to-row map without
    a Python row loop.  The hook is deliberately local to this trainer: it does
    not alter the user's HAKARI checkout, and it falls back to HAKARI's original
    class for another dataset implementation or an older PyArrow version.

    Set ``HAKARI_BENCH_FAST_CANDIDATE_INDEX=0`` to opt out when diagnosing an
    evaluation environment.
    """

    if os.environ.get("HAKARI_BENCH_FAST_CANDIDATE_INDEX", "1") == "0":
        return
    try:
        evaluation = importlib.import_module("hakari_bench.evaluation")
        original = evaluation._LazyCorpusById
    except (AttributeError, ImportError):
        return
    if getattr(original, "_bekko_arrow_index", False):
        return

    class ArrowIndexedLazyCorpus(original):
        """Drop-in HAKARI lazy corpus whose ID index uses vectorized Arrow."""

        _bekko_arrow_index = True

        def __init__(self, dataset: Any, ids: set[str]) -> None:
            try:
                import pyarrow as pa
                import pyarrow.compute as pc

                self._dataset = dataset
                self._max_text_chars = evaluation._reranker_document_max_chars()
                ids_only = (
                    dataset.select_columns(["_id"])
                    if hasattr(dataset, "select_columns")
                    else dataset
                )
                storage = getattr(ids_only, "_data", None)
                table = getattr(storage, "table", storage)
                if table is None or not hasattr(table, "column"):
                    raise TypeError("dataset does not expose an Arrow ID column")
                id_column = table.column("_id")
                matching_rows = pc.indices_nonzero(
                    pc.is_in(id_column, value_set=pa.array(list(ids)))
                )
                matching_ids = pc.take(id_column, matching_rows).to_pylist()
                row_indices = [int(index) for index in matching_rows.to_pylist()]
                # ``Dataset.__getitem__`` is inexpensive for a single row but
                # becomes the dominant cost when HAKARI asks for tens of
                # thousands of candidate documents one by one.  Gather the
                # same Arrow rows once, then preserve the Mapping interface
                # expected by HAKARI's scorer.
                matching_texts = dataset.select(row_indices)["text"]
                self._index_by_id = {
                    str(corpus_id): int(row_index)
                    for corpus_id, row_index in zip(
                        matching_ids, row_indices, strict=True
                    )
                }
                self._text_by_id = {
                    str(corpus_id): str(text)
                    for corpus_id, text in zip(
                        matching_ids, matching_texts, strict=True
                    )
                }

            except Exception:  # noqa: BLE001 - preserve HAKARI's safe fallback.
                # Keep exact HAKARI semantics when its internal Dataset layout
                # changes or Arrow cannot provide the required kernel.
                super().__init__(dataset, ids)

        def __getitem__(self, key: str) -> str:
            if not hasattr(self, "_text_by_id"):
                return super().__getitem__(key)
            text = self._text_by_id[key]
            if self._max_text_chars is not None and len(text) > self._max_text_chars:
                return text[: self._max_text_chars]
            return text

    evaluation._LazyCorpusById = ArrowIndexedLazyCorpus


def build_hakari_evaluator(
    *,
    hakari_repo: str | Path,
    mode: Literal["mid", "full"],
    batch_size: int,
    query_seed: int,
    embedding_variants: list[dict[str, Any]] | None = None,
    target_specs: Sequence[Mapping[str, Any]] | None = None,
    query_limit: int = 10,
    candidate_ranking: str = "reranking_hybrid",
    rerank_top_k: int = 100,
) -> Any:
    hakari = import_hakari(hakari_repo)
    is_mid = mode == "mid"
    if is_mid:
        # Candidate reranking evaluation only monitors training.  Static v1
        # encodes at most 512 tokens, so retaining arbitrarily long candidate
        # text merely spends CPU tokenizing suffixes that cannot reach the
        # model.  This HAKARI setting is read by its candidate-only lazy corpus
        # and does not affect the separate full HAKARI command after training.
        os.environ.setdefault("HAKARI_RERANKER_DOCUMENT_MAX_CHARS", "8192")
    raw_specs = list(target_specs or default_hakari_targets())
    grouped_specs: dict[int | None, list[Mapping[str, Any]]] = defaultdict(list)
    for spec in raw_specs:
        raw_limit = spec.get("query_limit", query_limit)
        if is_mid and not isinstance(raw_limit, (int, str)):
            raise TypeError("evaluation target query_limit must be an integer")
        limit = int(cast(int | str, raw_limit)) if is_mid else None
        grouped_specs[limit].append(spec)

    evaluators = []
    for group_index, (limit, specs) in enumerate(grouped_specs.items()):
        targets = []
        for spec in specs:
            kwargs: dict[str, Any] = {"splits": list(spec.get("splits", []))}
            if spec.get("collection") is not None:
                kwargs["collection"] = spec["collection"]
            else:
                kwargs["dataset"] = spec["dataset"]
            targets.append(hakari.HakariNanoTarget(**kwargs))
        if len(grouped_specs) == 1:
            name = "mid_eval" if is_mid else "full_eval"
        else:
            labels = [
                str(spec.get("dataset") or spec.get("collection")) for spec in specs
            ]
            slug = "_".join(label.lower().replace("-", "_") for label in labels)
            name = f"mid_eval_{slug}" if is_mid else f"full_eval_{group_index}"
        evaluators.append(
            hakari.HakariNanoEmbeddingEvaluator(
                targets=targets,
                name=name,
                batch_size=batch_size,
                metrics=("nDCG@10", "mAP@10"),
                query_limit=limit,
                query_sample_seed=query_seed,
                corpus_policy="sampled_candidates" if is_mid else "full",
                candidate_ranking=candidate_ranking if is_mid else None,
                rerank_top_k=rerank_top_k if is_mid else None,
                cache_datasets=True,
                embedding_variants=embedding_variants,
            )
        )
    if len(evaluators) == 1:
        return evaluators[0]
    return SequentialEvaluator(
        evaluators,
        main_score_function=lambda scores: sum(scores) / len(scores),
    )


def build_dimension_quantization_variants(
    hakari_repo: str | Path,
    *,
    dimensions: tuple[int, ...],
    precisions: tuple[Literal["int8", "binary"], ...] = (),
) -> list[dict[str, Any]]:
    if not dimensions:
        return []
    if any(dim <= 0 for dim in dimensions):
        raise ValueError("dimensions must be positive")
    import_hakari(hakari_repo)
    truncate = "truncate:" + ",".join(str(dim) for dim in dimensions)
    precision_spec = ",".join(precisions)
    values = [truncate]
    cross_values: list[list[str]] = []
    if precisions:
        values.append(precision_spec)
        cross_values.append([truncate, precision_spec])
    variants_module = importlib.import_module("hakari_bench.embedding_variants")
    return variants_module.parse_embedding_variants(values, cross_values)


def write_evaluation_result(
    path: Path,
    metrics: dict[str, float],
    metadata: dict[str, Any],
    *,
    variants: dict[str, dict[str, float]] | None = None,
) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    payload: dict[str, Any] = {"metadata": metadata, "metrics": metrics}
    if variants is not None:
        payload["variants"] = variants
    path.write_text(
        json.dumps(payload, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )


def evaluate_hakari_with_variants(
    evaluator: Any,
    model: Any,
    *,
    output_path: str | None,
) -> tuple[dict[str, float], dict[str, dict[str, float]]]:
    """Preserve Hakari embedding-variant metrics that its ST adapter omits."""
    hakari_evaluation = importlib.import_module("hakari_bench.evaluation")
    base_task_metrics: list[dict[str, float]] = []
    variant_task_metrics: dict[str, list[dict[str, float]]] = {}
    for task in evaluator.tasks:
        dataset = evaluator._load_dataset(task)
        evaluation = hakari_evaluation.evaluate_dense_task(
            model=model,
            dataset=dataset,
            batch_size=evaluator.batch_size,
            show_progress=evaluator.show_progress_bar,
            query_prompt=evaluator.query_prompt,
            corpus_prompt=evaluator.document_prompt,
            query_prompt_name=evaluator.query_prompt_name,
            corpus_prompt_name=evaluator.document_prompt_name,
            query_task=evaluator.query_encode_task,
            corpus_task=evaluator.document_encode_task,
            truncate_dim=evaluator.truncate_dim,
            truncate_sparse_query_max_dims=evaluator.sparse_query_max_active_dims,
            truncate_sparse_docs_max_dims=evaluator.sparse_document_max_active_dims,
            embedding_variants=evaluator.embedding_variants,
            aggregate_metric=evaluator.aggregate_metric_suffix,
            score_device=evaluator.retrieval_score_device,
            rerank_top_n=evaluator.rerank_top_k,
            candidate_ranking_name=evaluator.candidate_ranking,
            metric_names=evaluator.metric_suffixes,
        )
        base_task_metrics.append(evaluation.metrics)
        for embedding_evaluation in evaluation.embedding_evaluations:
            name = str(embedding_evaluation["name"])
            if name == "base":
                continue
            variant_task_metrics.setdefault(name, []).append(
                dict(embedding_evaluation["metrics"])
            )

    base = evaluator._finalize(
        model=model,
        metrics=base_task_metrics,
        output_path=output_path,
        epoch=-1,
        steps=-1,
    )
    expected_tasks = len(evaluator.tasks)
    variants = {
        name: evaluator._finalize(
            model=model, metrics=task_metrics, output_path=None, epoch=-1, steps=-1
        )
        for name, task_metrics in variant_task_metrics.items()
        if len(task_metrics) == expected_tasks
    }
    return base, variants


# ==============================================================================
# Training orchestration and CLI
# ==============================================================================

import argparse
import subprocess
import time


def materialize_config_shard(
    config: TrainConfig,
    *,
    shard_index: int,
    num_shards: int,
) -> dict[str, int]:
    if num_shards <= 0 or not 0 <= shard_index < num_shards:
        raise ValueError(
            f"invalid materialization shard: index={shard_index} count={num_shards}"
        )
    manifest = _read_direct_huggingface_manifest(config.data.spec)
    if manifest is None:
        raise ValueError("parallel materialization requires a direct Bekko manifest")
    raw_configs = manifest["configs"]
    if not isinstance(raw_configs, list):
        raise TypeError("direct manifest configs must be a list")
    manifest["configs"] = raw_configs[shard_index::num_shards]
    datasets, data_manifest, _ = load_direct_huggingface_datasets(
        manifest=manifest,
        fraction=config.data.fraction,
        seed=config.data.seed,
        require_cache=False,
        shared_id_block_size=config.data.shared_id_block_size,
        materialized_root=config.data.stage_cache_root,
        source_cache_direct=False,
        shared_id_column=config.data.shared_id_column,
    )
    return {
        "shard_index": shard_index,
        "configs": len(datasets),
        "selected_rows": sum(
            int(item["selected_rows"]) for item in data_manifest.values()
        ),
    }


def materialize_named_configs(
    config: TrainConfig, *, config_names: list[str]
) -> dict[str, int | list[str]]:
    """Materialize selected manifest configs for targeted cache recovery."""

    if not config_names:
        raise ValueError("at least one materialization config name is required")
    manifest = _read_direct_huggingface_manifest(config.data.spec)
    if manifest is None:
        raise ValueError("named materialization requires a direct Bekko manifest")
    raw_configs = manifest["configs"]
    if not isinstance(raw_configs, list):
        raise TypeError("direct manifest configs must be a list")
    requested = set(config_names)
    selected = [raw for raw in raw_configs if str(raw["name"]) in requested]
    found = {str(raw["name"]) for raw in selected}
    missing = sorted(requested.difference(found))
    if missing:
        raise ValueError(f"unknown direct manifest configs: {missing}")
    manifest["configs"] = selected
    datasets, data_manifest, _ = load_direct_huggingface_datasets(
        manifest=manifest,
        fraction=config.data.fraction,
        seed=config.data.seed,
        require_cache=False,
        shared_id_block_size=config.data.shared_id_block_size,
        materialized_root=config.data.stage_cache_root,
        source_cache_direct=False,
        shared_id_column=config.data.shared_id_column,
    )
    return {
        "config_names": [str(raw["name"]) for raw in selected],
        "configs": len(datasets),
        "selected_rows": sum(
            int(item["selected_rows"]) for item in data_manifest.values()
        ),
    }


def _materialized_cache_is_complete(
    cache_path: Path,
    *,
    expected_rows: int,
    task: str,
    shared_id_column: str,
    allow_fewer_rows: bool = False,
) -> bool:
    """Validate a saved Dataset from its small metadata files only.

    Restarting a full run must not reopen and project every Arrow shard merely to
    discover that materialization already finished.  The files written last by
    ``Dataset.save_to_disk`` describe the expected row count, schema, and Arrow
    members, so they provide a cheap and deterministic completion check.
    """
    try:
        state = json.loads((cache_path / "state.json").read_text(encoding="utf-8"))
        info = json.loads(
            (cache_path / "dataset_info.json").read_text(encoding="utf-8")
        )
        train_info = info["splits"]["train"]
        actual_rows = int(train_info["num_examples"])
        rows_match = (
            0 < actual_rows <= expected_rows
            if allow_fewer_rows
            else actual_rows == expected_rows
        )
        if not rows_match:
            return False
        required_columns = {"anchor", "positive", shared_id_column}
        if task == "triplet":
            required_columns.add("negative")
        if not required_columns.issubset(info["features"]):
            return False
        data_files = state["_data_files"]
        if not isinstance(data_files, list) or not data_files:
            return False
        for item in data_files:
            filename = item["filename"]
            path = cache_path / filename
            if not path.is_file() or path.stat().st_size == 0:
                return False
    except (KeyError, OSError, TypeError, ValueError, json.JSONDecodeError):
        return False
    return True


def all_direct_materialized_caches_complete(config: TrainConfig) -> bool:
    """Return whether every configured direct Bekko cache is complete."""
    if config.data.stage_cache_root is None:
        return False
    manifest = _read_direct_huggingface_manifest(config.data.spec)
    if manifest is None:
        return False
    raw_configs = manifest.get("configs")
    if not isinstance(raw_configs, list) or not raw_configs:
        return False
    root = config.data.stage_cache_root.resolve()
    for raw in raw_configs:
        if not isinstance(raw, dict):
            return False
        source_rows = int(raw["source_rows_all_splits"])
        expected_rows = min(
            source_rows,
            max(
                config.data.shared_id_block_size,
                math.ceil(
                    source_rows
                    * config.data.fraction
                    / config.data.shared_id_block_size
                )
                * config.data.shared_id_block_size,
            ),
        )
        cache_path = _direct_materialized_path(
            root,
            str(raw["name"]),
            config.data.fraction,
            config.data.seed,
        )
        if not _materialized_cache_is_complete(
            cache_path,
            expected_rows=expected_rows,
            task=str(raw["task"]),
            shared_id_column=config.data.shared_id_column,
            # The public manifest records an aligned all-splits upper bound.
            # A full train split can legitimately end before that bound.
            allow_fewer_rows=config.data.fraction == 1.0,
        ):
            return False
    return True


def materialize_in_parallel(config: TrainConfig) -> None:
    workers = config.data.materialize_workers
    if workers == 1 or config.data.direct_source_cache:
        return
    if all_direct_materialized_caches_complete(config):
        print(
            "All direct Bekko caches are complete; skipping materialization workers.",
            flush=True,
        )
        return

    def command(index: int) -> list[str]:
        return [
            sys.executable,
            str(Path(__file__).resolve()),
            "--config",
            str(config.source_path),
            "--materialize-shard",
            str(index),
            "--materialize-num-shards",
            str(workers),
        ]

    pending = list(range(workers))
    failures: list[tuple[int, int]] = []
    for attempt in range(1, config.data.materialize_max_attempts + 1):
        processes = {index: subprocess.Popen(command(index)) for index in pending}
        failures = [
            (index, return_code)
            for index, process in processes.items()
            if (return_code := process.wait()) != 0
        ]
        if not failures:
            return
        if attempt < config.data.materialize_max_attempts:
            pending = [index for index, _ in failures]
            print(
                "Materialization attempt "
                f"{attempt}/{config.data.materialize_max_attempts} failed for "
                f"shards {failures}; retrying those shards after "
                f"{config.data.materialize_retry_delay_seconds:g}s.",
                flush=True,
            )
            time.sleep(config.data.materialize_retry_delay_seconds)
    raise RuntimeError(
        "parallel materialization exhausted "
        f"{config.data.materialize_max_attempts} attempts: {failures}"
    )


from typing import Literal

from sentence_transformers import (
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
)
from sentence_transformers.sentence_transformer.losses import (
    MatryoshkaLoss,
    MultipleNegativesRankingLoss,
)
from sentence_transformers.sentence_transformer.training_args import (
    BatchSamplers,
    MultiDatasetBatchSamplers,
)
from torch import nn
from torch.optim import Optimizer
from torch.utils.data import BatchSampler
from transformers import (
    TrainerCallback,
    TrainerControl,
    TrainerState,
    TrainingArguments,
)


class LossComponentLoggingCallback(TrainerCallback):
    """Expose per-dimension float/QAT loss components in trainer and W&B logs."""

    def __init__(self, losses: list[Qwen3EmbeddingLoss]) -> None:
        self.losses = losses

    def on_log(
        self,
        args: TrainingArguments,
        state: TrainerState,
        control: TrainerControl,
        logs: dict[str, float] | None = None,
        **kwargs: object,
    ) -> None:
        del args, state, control, kwargs
        if logs is None:
            return
        for loss_index, loss in enumerate(self.losses):
            prefix = (
                "loss_component"
                if len(self.losses) == 1
                else f"loss_{loss_index}_component"
            )
            for name, value in loss.last_loss_components.items():
                logs[f"{prefix}/{name}"] = value.item()


class EvaluationProgressCallback(TrainerCallback):
    """Print every sampled Hakari result as one machine-readable progress line."""

    def __init__(self, target_evaluations: int) -> None:
        if target_evaluations <= 0:
            raise ValueError("target_evaluations must be positive")
        self.target_evaluations = target_evaluations
        self.completed_evaluations = 0

    def on_evaluate(
        self,
        args: TrainingArguments,
        state: TrainerState,
        control: TrainerControl,
        metrics: dict[str, float] | None = None,
        **kwargs: object,
    ) -> None:
        del args, control, kwargs
        self.completed_evaluations += 1
        numeric_metrics = {
            key: float(value)
            for key, value in (metrics or {}).items()
            if isinstance(value, (int, float, np.number))
        }
        payload = {
            "event": "mid_eval_result",
            "evaluation": self.completed_evaluations,
            "target_evaluations": self.target_evaluations,
            "global_step": int(state.global_step),
            "max_steps": int(state.max_steps),
            "training_progress_percent": (
                100.0 * state.global_step / state.max_steps if state.max_steps else 0.0
            ),
            "metrics": numeric_metrics,
        }
        print("MID_EVAL_RESULT " + json.dumps(payload, sort_keys=True), flush=True)


def split_sse_optimizer_groups(
    optimizer: Optimizer,
    model: nn.Module,
    *,
    sse_learning_rate: float,
) -> None:
    """Move trainable SeparableDyT parameters into LR-specific optimizer groups."""
    dyt_parameter_ids = {
        id(parameter)
        for name, parameter in model.named_parameters()
        if parameter.requires_grad and (name.startswith("dyt.") or ".dyt." in name)
    }
    if not dyt_parameter_ids:
        return
    new_groups: list[dict[str, object]] = []
    for group in optimizer.param_groups:
        regular = [
            parameter
            for parameter in group["params"]
            if id(parameter) not in dyt_parameter_ids
        ]
        dyt = [
            parameter
            for parameter in group["params"]
            if id(parameter) in dyt_parameter_ids
        ]
        if regular:
            new_groups.append({**group, "params": regular})
        if dyt:
            dyt_group = {**group, "params": dyt, "lr": sse_learning_rate}
            if "initial_lr" in dyt_group:
                dyt_group["initial_lr"] = sse_learning_rate
            new_groups.append(dyt_group)
    optimizer.param_groups[:] = new_groups


class SSESentenceTransformerTrainer(SentenceTransformerTrainer):
    def __init__(
        self,
        *args: object,
        sse_learning_rate: float | None = None,
        locality_block_size: int | None = None,
        **kwargs: object,
    ) -> None:
        self.sse_learning_rate = sse_learning_rate
        self.locality_block_size = locality_block_size
        super().__init__(*args, **kwargs)

    def create_optimizer(self, model: nn.Module | None = None) -> Optimizer:
        optimizer = super().create_optimizer(model=model)
        if self.sse_learning_rate is not None:
            split_sse_optimizer_groups(
                optimizer,
                self.model if model is None else model,
                sse_learning_rate=self.sse_learning_rate,
            )
        return optimizer

    def get_eval_dataloader(self, eval_dataset: Any = None) -> Any:
        """Build evaluation loaders without forking after CUDA initialization.

        Training benefits substantially from parallel reads over the materialized
        Arrow collection. Evaluation is different: its small loss dataset is first
        opened after CUDA and tokenizer thread pools are active. Forking workers at
        that point can hang before the first batch. Keep the configured workers for
        training, but construct each evaluation loader in the trainer process.
        """
        workers = self.args.dataloader_num_workers
        prefetch = self.args.dataloader_prefetch_factor
        persistent = self.args.dataloader_persistent_workers
        try:
            self.args.dataloader_num_workers = 0
            self.args.dataloader_prefetch_factor = None
            self.args.dataloader_persistent_workers = False
            return super().get_eval_dataloader(eval_dataset)
        finally:
            self.args.dataloader_num_workers = workers
            self.args.dataloader_prefetch_factor = prefetch
            self.args.dataloader_persistent_workers = persistent

    def get_batch_sampler(
        self,
        dataset: Dataset,
        batch_size: int,
        drop_last: bool,
        valid_label_columns: list[str] | None = None,
        generator: torch.Generator | None = None,
        seed: int = 0,
    ) -> BatchSampler | None:
        if representation_batch_plan(dataset) is not None:
            return RepresentationNoDuplicatesBatchSampler(
                dataset,
                batch_size=batch_size,
                drop_last=drop_last,
                valid_label_columns=valid_label_columns,
                generator=generator,
                seed=seed,
            )
        if self.locality_block_size is not None:
            return LocalityAwareNoDuplicatesBatchSampler(
                dataset,
                batch_size=batch_size,
                drop_last=drop_last,
                valid_label_columns=valid_label_columns,
                generator=generator,
                seed=seed,
                locality_block_size=self.locality_block_size,
            )
        return super().get_batch_sampler(
            dataset,
            batch_size=batch_size,
            drop_last=drop_last,
            valid_label_columns=valid_label_columns,
            generator=generator,
            seed=seed,
        )


def train(config: TrainConfig) -> Path:
    _seed_everything(config.training.seed)
    output_dir = config.run.output_dir.resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    train_dataset, data_manifest, pair_types = load_bekko_datasets(
        spec=config.data.spec,
        tasks=config.data.tasks,
        fraction=config.data.fraction,
        seed=config.data.seed,
        require_cache=config.data.require_cache,
        shared_id_column=config.data.shared_id_column,
        shared_id_block_size=config.data.shared_id_block_size,
        supplements=config.data.supplements,
        stage_cache_root=config.data.stage_cache_root,
        require_stage_cache=config.data.require_stage_cache,
        direct_source_cache=config.data.direct_source_cache,
    )
    model = build_scratch_model(
        config.model.tokenizer_model,
        embedding_dim=config.model.embedding_dim,
        max_seq_length=config.model.max_seq_length,
        architecture=config.model.architecture,
        sse_preprocessing=config.model.sse_preprocessing,
        sse_alpha_init=config.model.sse_alpha_init,
        sse_train_alpha=config.model.sse_train_alpha,
        sse_train_beta=config.model.sse_train_beta,
        sse_train_bias=config.model.sse_train_bias,
        text_preprocessing=config.model.text_preprocessing,
    )
    representation_manifests = prepare_representation_batch_plans(
        train_dataset,
        tokenizer_path=Path(config.model.tokenizer_model),
        text_preprocessing=config.model.text_preprocessing,
        max_seq_length=config.model.max_seq_length,
        batch_size=config.training.batch_size,
        seed=config.training.seed,
        default_cache_root=output_dir / "representation-cache",
    )
    for name, manifest in representation_manifests.items():
        data_manifest[name].update(
            {f"representation_{key}": value for key, value in manifest.items()}
        )
    component_batches = []
    for name, dataset in train_dataset.items():
        plan = representation_batch_plan(dataset)
        component_batches.append(
            len(plan)
            if plan is not None
            else math.ceil(len(dataset) / config.training.batch_size)
        )
    schedule = derive_multidataset_step_schedule(
        component_rows=[
            count * config.training.batch_size for count in component_batches
        ],
        batch_size=config.training.batch_size,
        epochs=config.training.epochs,
        target_evaluations=config.training.target_evaluations,
        exact_target_evaluations=config.training.exact_target_evaluations,
    )
    (output_dir / "data_manifest.json").write_text(
        json.dumps(data_manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    loss = build_training_loss(
        model,
        settings=config.loss,
        matryoshka_dims=config.model.matryoshka_dims,
        matryoshka_weights=config.model.matryoshka_weights,
        pair_types=pair_types,
    )
    evaluator = build_hakari_evaluator(
        hakari_repo=config.evaluation.hakari_repo,
        mode="mid",
        batch_size=config.evaluation.batch_size,
        query_seed=config.evaluation.query_seed,
        target_specs=[
            {
                "dataset": target.dataset,
                "collection": target.collection,
                "splits": list(target.splits),
                "query_limit": target.query_limit,
            }
            for target in config.evaluation.targets
        ]
        or None,
        query_limit=config.evaluation.query_limit,
        candidate_ranking=config.evaluation.candidate_ranking,
        rerank_top_k=config.evaluation.rerank_top_k,
    )

    if config.tracking.enabled:
        os.environ["WANDB_PROJECT"] = config.tracking.project
    report_to = ["wandb"] if config.tracking.enabled else []
    args = SentenceTransformerTrainingArguments(
        output_dir=str(output_dir),
        num_train_epochs=config.training.epochs,
        per_device_train_batch_size=config.training.batch_size,
        per_device_eval_batch_size=config.training.eval_batch_size,
        learning_rate=config.training.learning_rate,
        adam_beta2=config.training.adam_beta2,
        adam_epsilon=config.training.adam_epsilon,
        weight_decay=config.training.weight_decay,
        lr_scheduler_type="cosine",
        warmup_steps=config.training.warmup_ratio,
        bf16=True,
        fp16=False,
        batch_sampler=BatchSamplers.NO_DUPLICATES,
        multi_dataset_batch_sampler=MultiDatasetBatchSamplers.PROPORTIONAL,
        eval_strategy="steps",
        eval_steps=schedule.eval_steps,
        save_strategy="steps" if config.training.save_checkpoints else "no",
        save_steps=schedule.save_steps,
        save_total_limit=config.training.save_total_limit,
        logging_steps=max(1, math.ceil(schedule.eval_steps / 5)),
        logging_first_step=True,
        dataloader_num_workers=config.training.dataloader_num_workers,
        dataloader_prefetch_factor=config.training.dataloader_prefetch_factor,
        report_to=report_to,
        run_name=config.run.name,
        seed=config.training.seed,
        data_seed=config.data.seed,
        load_best_model_at_end=False,
    )
    trainer = SSESentenceTransformerTrainer(
        model=model,
        args=args,
        train_dataset=train_dataset,
        # Intermediate evaluation is retrieval-only. Passing a training-derived
        # eval dataset makes Trainer run an additional loss-evaluation DataLoader
        # before NanoMIRACL/MNanoBEIR. Besides being outside the intended metric,
        # that loader can deadlock while training workers hold the very large
        # memory-mapped Arrow collection. With an evaluator and no eval dataset,
        # SentenceTransformers uses an empty loss loader and directly runs the
        # small fixed-candidate retrieval evaluator.
        eval_dataset=None,
        loss=loss,
        evaluator=evaluator,
        callbacks=[
            LossComponentLoggingCallback(_qwen_losses(loss)),
            EvaluationProgressCallback(config.training.target_evaluations),
        ],
        sse_learning_rate=config.training.sse_learning_rate,
        locality_block_size=config.training.dataloader_locality_block_size,
    )
    # SentenceTransformers 5.6 records heterogeneous DatasetDict metadata in a
    # shape its model-card template cannot render. Keep checkpoints quiet and
    # generate the final card separately once the run metadata is complete.
    model.model_card_data.train_datasets = []
    model.model_card_data.eval_datasets = []
    model.model_card_data.generate_widget_examples = False
    trainer.train()
    final_dir = output_dir / "final"
    model.save_pretrained(str(final_dir))
    set_final_tokenizer_max_length(final_dir, config.model.final_max_seq_length)
    return final_dir


def set_final_tokenizer_max_length(model_dir: Path, max_length: int) -> None:
    """Expand only the exported tokenizer after bounded-length training.

    Training and checkpoints retain ``model.max_seq_length``. The final model is
    rewritten after serialization so inference can consume much longer inputs
    without allocating unbounded tokenizer output.
    """

    if max_length <= 0:
        raise ValueError("final tokenizer max_length must be positive")
    tokenizer_path = model_dir / "tokenizer.json"
    if not tokenizer_path.is_file():
        raise FileNotFoundError(f"final tokenizer not found: {tokenizer_path}")
    tokenizer = Tokenizer.from_file(str(tokenizer_path))
    tokenizer.enable_truncation(max_length=max_length)
    temporary_path = tokenizer_path.with_name(tokenizer_path.name + ".tmp")
    tokenizer.save(str(temporary_path))
    os.replace(temporary_path, tokenizer_path)

    config_path = model_dir / "tokenizer_config.json"
    if config_path.is_file():
        payload = json.loads(config_path.read_text(encoding="utf-8"))
        if not isinstance(payload, dict):
            raise TypeError(f"tokenizer config must be a mapping: {config_path}")
        payload["model_max_length"] = max_length
        temporary_config = config_path.with_name(config_path.name + ".tmp")
        temporary_config.write_text(
            json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8"
        )
        os.replace(temporary_config, config_path)


def _qwen_losses(loss: nn.Module | dict[str, nn.Module]) -> list[Qwen3EmbeddingLoss]:
    wrappers = list(loss.values()) if isinstance(loss, dict) else [loss]
    result: list[Qwen3EmbeddingLoss] = []
    for wrapper in wrappers:
        if isinstance(wrapper, Qwen3EmbeddingLoss):
            result.append(wrapper)
        elif isinstance(wrapper, MatryoshkaLoss) and isinstance(
            wrapper.loss, Qwen3EmbeddingLoss
        ):
            result.append(wrapper.loss)
    return result


def build_full_hakari_command(config: TrainConfig, final_dir: Path) -> list[str]:
    """Build the resumable full-evaluation command for the saved final model."""

    settings = config.evaluation
    results_dir = (
        settings.full_results_dir
        if settings.full_results_dir is not None
        else config.run.output_dir / "full-eval" / "hakari-results"
    ).resolve()
    command = [
        "uv",
        "run",
        "--project",
        str(settings.hakari_repo.resolve()),
        "hakari-bench",
        "evaluate",
        "dense",
        "--model",
        str(final_dir.resolve()),
        "--model-alias",
        settings.full_model_alias or config.run.name,
        "--dtype",
        settings.full_dtype,
        "--device",
        settings.full_device,
        "--retrieval-score-device",
        settings.full_retrieval_score_device,
        "--batch-size",
        str(settings.batch_size),
        "--results-dir",
        str(results_dir),
        "--result-format",
        settings.full_result_format,
        "--show-progress",
    ]
    truncate_dims = sorted(
        (dim for dim in config.model.matryoshka_dims if dim < config.model.embedding_dim),
        reverse=True,
    )
    if truncate_dims:
        command.extend(
            ["--embedding-variant", "truncate:" + ",".join(map(str, truncate_dims))]
        )
    if settings.full_all:
        command.append("--all")
    else:
        for target in settings.targets:
            if target.collection is not None:
                command.extend(["--collection", target.collection])
            elif target.dataset is not None:
                command.extend(["--dataset", target.dataset])
            for split in target.splits:
                command.extend(["--split", split])
    if settings.full_overwrite:
        command.append("--overwrite")
    return command


def run_full_hakari_evaluation(config: TrainConfig, final_dir: Path) -> Path | None:
    """Run all-query/full-corpus Hakari evaluation after successful training."""

    if not config.evaluation.full_after_training:
        return None
    command = build_full_hakari_command(config, final_dir)
    manifest_path = config.run.output_dir.resolve() / "full-eval" / "run.json"
    manifest_path.parent.mkdir(parents=True, exist_ok=True)
    started_at = time.time()
    payload: dict[str, Any] = {
        "status": "running",
        "started_at_unix": started_at,
        "command": command,
        "model": str(final_dir.resolve()),
    }
    manifest_path.write_text(
        json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    print("FULL_EVAL_START " + json.dumps(payload, sort_keys=True), flush=True)
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    try:
        subprocess.run(command, check=True)
    except BaseException as error:
        payload.update(
            {
                "status": "failed",
                "finished_at_unix": time.time(),
                "error": f"{type(error).__name__}: {error}",
            }
        )
        manifest_path.write_text(
            json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8"
        )
        print("FULL_EVAL_FAILED " + json.dumps(payload, sort_keys=True), flush=True)
        raise
    payload.update(
        {
            "status": "completed",
            "finished_at_unix": time.time(),
            "runtime_seconds": time.time() - started_at,
        }
    )
    manifest_path.write_text(
        json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    print("FULL_EVAL_COMPLETE " + json.dumps(payload, sort_keys=True), flush=True)
    return manifest_path


def build_training_loss(
    model: SentenceTransformer,
    *,
    settings: LossSettings,
    matryoshka_dims: tuple[int, ...],
    matryoshka_weights: tuple[float, ...],
    pair_types: Mapping[str, Literal["query_doc", "doc_doc"]],
) -> nn.Module | dict[str, nn.Module]:
    def make_loss(*, bidirectional: bool) -> nn.Module:
        if settings.name == "qwen3_improved_infonce":
            base_loss = Qwen3EmbeddingLoss(
                model,
                temperature=settings.temperature,
                margin=settings.margin,
                hard_negative_margin=settings.hard_negative_margin,
                bidirectional=bidirectional,
                query_query_weight=settings.query_query_weight,
                document_document_weight=settings.document_document_weight,
                cross_document_weight=settings.cross_document_weight,
                hard_negative_weight=settings.hard_negative_weight,
                qat_int8_weight=settings.qat_int8_weight,
                qat_binary_weight=settings.qat_binary_weight,
                qat_n_quantized_per_step=settings.qat_n_quantized_per_step,
            )
        else:
            # Keep ``loss.temperature`` meaningful for both objective families.
            # Sentence Transformers exposes MNRL temperature as its reciprocal
            # ``scale`` argument instead of accepting temperature directly.
            base_loss = MultipleNegativesRankingLoss(
                model, scale=1.0 / settings.temperature
            )
        if not matryoshka_dims:
            return base_loss
        return MatryoshkaLoss(
            model,
            base_loss,
            matryoshka_dims=list(matryoshka_dims),
            matryoshka_weights=list(matryoshka_weights),
        )

    if settings.pair_type_routing == "shared":
        return make_loss(bidirectional=False)
    return {
        key: make_loss(
            bidirectional=pair_type == "doc_doc" and settings.doc_doc_bidirectional
        )
        for key, pair_type in pair_types.items()
    }


def _seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def configure_live_console_logging() -> None:
    """Flush trainer metrics promptly when output is redirected to a log file."""

    for stream in (sys.stdout, sys.stderr):
        reconfigure = getattr(stream, "reconfigure", None)
        if reconfigure is not None:
            reconfigure(line_buffering=True)


def main() -> None:
    configure_live_console_logging()
    parser = argparse.ArgumentParser(
        description="Train multilingual static embeddings directly from public Bekko data"
    )
    parser.add_argument(
        "--config",
        type=Path,
        required=True,
        help="YAML experiment configuration (paths are resolved from the working directory)",
    )
    parser.add_argument("--materialize-shard", type=int, default=None)
    parser.add_argument("--materialize-num-shards", type=int, default=None)
    parser.add_argument(
        "--materialize-config",
        action="append",
        default=[],
        help="Materialize only this manifest config; repeat for multiple configs",
    )
    args = parser.parse_args()
    config = load_config(args.config)
    if args.materialize_config:
        if args.materialize_shard is not None:
            parser.error("--materialize-config cannot be combined with --materialize-shard")
        result = materialize_named_configs(
            config, config_names=list(args.materialize_config)
        )
        print(json.dumps(result, sort_keys=True), flush=True)
        os._exit(0)
    if args.materialize_shard is not None:
        if args.materialize_num_shards is None:
            parser.error("--materialize-shard requires --materialize-num-shards")
        result = materialize_config_shard(
            config,
            shard_index=args.materialize_shard,
            num_shards=args.materialize_num_shards,
        )
        print(json.dumps(result, sort_keys=True), flush=True)
        # CPython 3.13 can trip a PyArrow finalizer after every durable Arrow
        # slice has already been closed. Shard workers have no further work.
        os._exit(0)
    materialize_in_parallel(config)
    final_dir = train(config)
    print(final_dir, flush=True)
    run_full_hakari_evaluation(config, final_dir)


if __name__ == "__main__":
    main()