Hugging Face's logo

hotchpotch
/
bekko-static-embedding-v1

Sentence Similarity
sentence-transformers
ONNX
Safetensors
OpenVINO
static-embedding
embeddings
feature-extraction
matryoshka
retrieval
Model card Files
xet
 Community
bekko-static-embedding-v1 / training_scripts /train_model.py
hotchpotch's picture
hotchpotch
Release multilingual v1 with evaluated FP16 static weights
1bf933f verified
Raw History Blame Contribute Delete
139 kB
#!/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()