tribunal-kit 9.2.1 → 9.2.4

package/skills/train-sentence-transformers/scripts/train_multi_vector_encoder_example.py

@@ -0,0 +1,240 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "sentence-transformers[train]>=6.0",
+#     "datasets>=2.19.0",
+#     "accelerate>=0.26.0",
+#     "trackio",
+# ]
+# ///
+"""Production-ready multi-vector (ColBERT / late-interaction) training template.
+
+This script demonstrates a recommended setup:
+- MultiVectorMultipleNegativesRankingLoss on (anchor, positive, negative) triplets, scored with MaxSim
+- MultiVectorNanoBEIREvaluator for retrieval metrics during training
+- BatchSamplers.NO_DUPLICATES (critical for MNRL-family)
+- load_best_model_at_end with a retrieval metric
+- Auto model card + optional Hub push
+
+Runs identically in two modes:
+
+    # Local
+    pip install "sentence-transformers[train]>=6.0"
+    python train_multi_vector_encoder_example.py
+
+    # Or with uv (no explicit install needed)
+    uv run train_multi_vector_encoder_example.py
+
+    # Multi-GPU
+    accelerate launch train_multi_vector_encoder_example.py
+
+    # Hugging Face Jobs (paste the entire file contents as `script`)
+    hf_jobs("uv", {
+        "script": "<contents of this file>",
+        "flavor": "a10g-large",
+        "timeout": "3h",
+        "secrets": {"HF_TOKEN": "$HF_TOKEN"},
+    })
+
+Adjust MODEL_NAME, DATASET_NAME, OUTPUT_DIR, RUN_NAME at the top of the script.
+Default Hub push: at end of run, public, under your authenticated user as
+`{user}/{RUN_NAME}`, wrapped in try/except. To skip the push, comment out the
+push_to_hub call. For HF Jobs (ephemeral env), also enable in-trainer push:
+add `push_to_hub=True`, `hub_model_id=RUN_NAME`, `hub_strategy="every_save"`
+to TrainingArguments.
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import os
+from contextlib import nullcontext
+
+import torch
+from datasets import load_dataset
+
+from sentence_transformers import (
+    MultiVectorEncoder,
+    MultiVectorEncoderModelCardData,
+    MultiVectorEncoderTrainer,
+    MultiVectorEncoderTrainingArguments,
+)
+from sentence_transformers.base.sampler import BatchSamplers
+from sentence_transformers.multi_vector_encoder.evaluation import MultiVectorNanoBEIREvaluator
+from sentence_transformers.multi_vector_encoder.losses import MultiVectorMultipleNegativesRankingLoss
+
+
+def autocast_ctx():
+    """bf16/fp16 autocast for evaluator calls outside the trainer (which has its own autocast)."""
+    if not torch.cuda.is_available():
+        return nullcontext()
+    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+    return torch.autocast("cuda", dtype=dtype)
+
+
+def log_trackio_dashboard():
+    """Surface the Trackio dashboard URL so the user can watch training live."""
+    try:
+        from huggingface_hub import whoami
+
+        hf_user = whoami().get("name")
+        if hf_user:
+            logging.info(
+                f"Trackio dashboard (live training progress): https://huggingface.co/spaces/{hf_user}/trackio"
+            )
+    except Exception:
+        pass
+
+
+# Bare HF encoder. MVE grafts a fresh 128-dim token projection on top.
+MODEL_NAME = "answerdotai/ModernBERT-base"
+DATASET_NAME = "sentence-transformers/msmarco-bm25"
+DATASET_SUBSET = "triplet"
+TRAIN_SIZE = 50_000
+EVAL_SIZE = 1_000
+OUTPUT_DIR = "models/modernbert-base-msmarco-colbert"
+RUN_NAME = "modernbert-base-msmarco-colbert"
+SMOKE_TEST = os.environ.get("SMOKE_TEST") == "1"
+
+
+def setup_logging():
+    """Configure logging + TF32. Tees to logs/{RUN_NAME}.log and silences HTTP spam."""
+    os.makedirs("logs", exist_ok=True)
+    logging.basicConfig(
+        format="%(asctime)s - %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+        level=logging.INFO,
+        handlers=[logging.StreamHandler(), logging.FileHandler(f"logs/{RUN_NAME}.log")],
+        force=True,
+    )
+    for noisy in ("httpx", "httpcore", "huggingface_hub", "urllib3", "filelock", "fsspec"):
+        logging.getLogger(noisy).setLevel(logging.WARNING)
+    if torch.cuda.is_available():
+        torch.set_float32_matmul_precision("high")  # TF32 on Ampere+, no quality loss
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--eval-only", type=str, default=None, help="Skip training, load this saved model and run only the evaluator."
+    )
+    cli, _ = parser.parse_known_args()
+
+    setup_logging()
+
+    if cli.eval_only:
+        logging.info(f"Eval-only mode: loading model from {cli.eval_only}")
+        model = MultiVectorEncoder(cli.eval_only)
+        evaluator = MultiVectorNanoBEIREvaluator(batch_size=16)
+        with autocast_ctx():
+            evaluator(model)
+        return
+
+    logging.info(f"Loading base model: {MODEL_NAME}")
+    # Load in fp32, autocast bf16 during forward. torch_dtype=bfloat16 breaks the optimizer's precision.
+    model = MultiVectorEncoder(
+        MODEL_NAME,
+        model_card_data=MultiVectorEncoderModelCardData(
+            language="en",
+            license="apache-2.0",
+            model_name=f"ColBERT {MODEL_NAME.split('/')[-1]} trained on MS MARCO",
+        ),
+        model_kwargs={"torch_dtype": "float32"},
+    )
+
+    logging.info(f"Loading dataset: {DATASET_NAME} ({DATASET_SUBSET})")
+    train_size = 50 if SMOKE_TEST else TRAIN_SIZE
+    eval_size = 20 if SMOKE_TEST else EVAL_SIZE
+    full_dataset = load_dataset(DATASET_NAME, DATASET_SUBSET, split="train").select(range(train_size + eval_size))
+    dataset_dict = full_dataset.train_test_split(test_size=eval_size, seed=12)
+    train_dataset = dataset_dict["train"]
+    eval_dataset = dataset_dict["test"]
+    if SMOKE_TEST:
+        logging.info("SMOKE_TEST=1: trimmed dataset, will run max_steps=1 and skip Hub push")
+    logging.info(f"  train: {len(train_dataset):,} examples")
+    logging.info(f"  eval:  {len(eval_dataset):,} examples")
+
+    # MNRL with in-batch negatives + explicit hard negatives. scale=1.0 (default) is correct for
+    # unnormalized MaxSim: do not copy scale=20.0 from bi-encoder MNRL. Length-normalized MeanMaxSim
+    # scoring instead wants a scale of roughly the average query length.
+    loss = MultiVectorMultipleNegativesRankingLoss(model=model)
+
+    # Cheap in-training evaluator on 3 datasets (drives `load_best_model_at_end`).
+    # The end-of-run evaluator below covers the full 13-dataset suite for the VERDICT delta.
+    evaluator = MultiVectorNanoBEIREvaluator(dataset_names=["msmarco", "nq", "fiqa2018"], batch_size=16)
+    test_evaluator = MultiVectorNanoBEIREvaluator(show_progress_bar=True, batch_size=16)
+    logging.info("Baseline evaluation (before training):")
+    with autocast_ctx():
+        # Must run before deriving metric_key: evaluator(model) mutates primary_metric to add the name_ prefix.
+        evaluator(model)  # populates primary_metric on the training-time evaluator
+        # Baseline on the full suite so the VERDICT delta is apples-to-apples with the post-training score.
+        baseline_eval = test_evaluator(model)[test_evaluator.primary_metric]
+    metric_key = f"eval_{evaluator.primary_metric}"
+
+    args = MultiVectorEncoderTrainingArguments(
+        output_dir=OUTPUT_DIR,
+        num_train_epochs=1,
+        max_steps=1 if SMOKE_TEST else -1,
+        per_device_train_batch_size=32,
+        per_device_eval_batch_size=32,
+        learning_rate=3e-5,
+        weight_decay=0.01,
+        warmup_steps=0.05,
+        lr_scheduler_type="linear",
+        bf16=True,
+        batch_sampler=BatchSamplers.NO_DUPLICATES,
+        eval_strategy="steps",
+        eval_steps=0.1,
+        save_strategy="steps",
+        save_steps=0.1,
+        save_total_limit=2,
+        logging_steps=0.01,
+        logging_first_step=True,
+        load_best_model_at_end=True,
+        metric_for_best_model=metric_key,
+        greater_is_better=True,
+        report_to="none" if SMOKE_TEST else "trackio",
+        run_name=RUN_NAME,
+        seed=12,
+    )
+
+    trainer = MultiVectorEncoderTrainer(
+        model=model,
+        args=args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        loss=loss,
+        evaluator=evaluator,
+    )
+    if not SMOKE_TEST:
+        log_trackio_dashboard()
+    trainer.train()
+
+    logging.info("Post-training evaluation (full NanoBEIR suite):")
+    with autocast_ctx():
+        score = test_evaluator(model)[test_evaluator.primary_metric]
+    delta = score - baseline_eval
+    verdict = "WIN" if delta >= 0.005 else "MARGINAL" if delta >= 0 else "REGRESSION"
+    logging.info(f"VERDICT: {verdict} | score={score:.4f} | baseline={baseline_eval:.4f} | delta={delta:+.4f}")
+
+    final_dir = f"{OUTPUT_DIR}/final"
+    model.save_pretrained(final_dir)
+    logging.info(f"Saved final model to {final_dir}")
+
+    if SMOKE_TEST:
+        logging.info("SMOKE_TEST=1: skipping Hub push")
+        return
+
+    try:
+        commit_url = model.push_to_hub(RUN_NAME)  # public by default, uses your authenticated user
+        logging.info(f"Pushed model to {commit_url.rsplit('/commit/', 1)[0]}")
+    except Exception:
+        import traceback
+
+        logging.error(f"Hub push failed:\n{traceback.format_exc()}")
+
+
+if __name__ == "__main__":
+    main()

package/skills/train-sentence-transformers/scripts/train_sentence_transformer_distillation_example.py

@@ -0,0 +1,303 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "sentence-transformers[train]>=5.0",
+#     "datasets>=2.19.0",
+#     "accelerate>=0.26.0",
+#     "trackio",
+# ]
+# ///
+"""Distillation template: train a small student to match a stronger teacher's embeddings.
+
+This script implements **embedding MSE** (Pattern 1 below): pre-compute teacher
+embeddings, train the student to minimize MSE against them. Yields a smaller /
+faster student typically reaching 95-99% of teacher quality.
+
+Three distillation patterns:
+- Pattern 1 (this script): `(text, teacher_embedding)` + `MSELoss`. Cheapest,
+  most data-efficient. Use when student + teacher are both bi-encoders with the
+  same output dim and you have a pile of unlabeled text.
+- Pattern 2: `(query, positive, negative, score_diff)` + `MarginMSELoss` from a
+  CrossEncoder teacher's score differences. Workhorse of ms-marco distillation.
+  See `../references/losses_sentence_transformer.md` (MarginMSELoss section).
+- Pattern 3: `(query, positive, neg_1, ..., neg_n, labels)` + `DistillKLDivLoss`
+  to preserve the full teacher distribution. More data-hungry. Natural fit when
+  distilling from an ensemble of rerankers.
+
+Mismatched dims: if the student's dim is smaller than the teacher's, MSELoss
+fails. Add a PCA-init `Dense` projection so the teacher matches the student:
+
+    from sklearn.decomposition import PCA
+    from sentence_transformers.sentence_transformer.modules import Dense
+    pca = PCA(n_components=student.get_embedding_dimension())
+    pca.fit(teacher.encode(sentences[:20_000], convert_to_numpy=True))
+    dense = Dense(
+        in_features=teacher.get_embedding_dimension(),
+        out_features=student.get_embedding_dimension(),
+        bias=False, activation_function=torch.nn.Identity(),
+    )
+    dense.linear.weight = torch.nn.Parameter(torch.from_numpy(pca.components_).float())
+    teacher.add_module("dense", dense)
+
+Distilling to a CrossEncoder student: construct with `activation_fn=nn.Identity()`
+or eval ranking collapses silently. Every non-BCE CE loss expects raw logits
+during training, but the model's `activation_fn` runs at eval time inside
+`predict()`. Default `Sigmoid` (when `num_labels=1`) saturates raw logits >5 to
+~1.0, dropping nDCG from e.g. ~0.59 to ~0.14 with healthy-looking training loss.
+Applies to all CE distillation / listwise / pairwise losses. See SKILL.md
+Directive 7 ([CE]).
+
+Layer pruning shortcut for Pattern 1: copy the teacher, delete layers (often
+keeps 99%+ of quality at a fraction of the layers), then distill with MSELoss:
+
+    from copy import deepcopy
+    student = deepcopy(teacher)
+    layers = student.transformers_model.encoder.layer  # BERT/MPNet/DistilBERT
+    student.transformers_model.encoder.layer = torch.nn.ModuleList(
+        [layers[0], layers[3], layers[6], layers[9]]
+    )
+    student.transformers_model.config.num_hidden_layers = 4
+
+Tips: pre-compute teacher outputs once and cache (`dataset.save_to_disk`). LR
+1e-4 (higher than the usual 2e-5: the target is dense regression). 1 epoch is
+usually enough. The student inherits the teacher's weaknesses, so pick a
+teacher strong on YOUR task. If the teacher expects an instruction prefix,
+include it during teacher encoding so the student's target matches inference.
+
+For multilingual student distillation (extend an English teacher to other
+languages without in-language supervised data), see `train_sentence_transformer_make_multilingual_example.py`.
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import os
+from contextlib import nullcontext
+
+import torch
+from datasets import Dataset, load_dataset
+
+from sentence_transformers import (
+    SentenceTransformer,
+    SentenceTransformerModelCardData,
+    SentenceTransformerTrainer,
+    SentenceTransformerTrainingArguments,
+)
+from sentence_transformers.sentence_transformer.evaluation import EmbeddingSimilarityEvaluator
+from sentence_transformers.sentence_transformer.losses import MSELoss
+from sentence_transformers.sentence_transformer.modules import Normalize
+from sentence_transformers.util.similarity import SimilarityFunction
+
+
+def autocast_ctx():
+    """bf16/fp16 autocast for evaluator calls outside the trainer (which has its own autocast)."""
+    if not torch.cuda.is_available():
+        return nullcontext()
+    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+    return torch.autocast("cuda", dtype=dtype)
+
+
+def log_trackio_dashboard():
+    """Surface the Trackio dashboard URL so the user can watch training live."""
+    try:
+        from huggingface_hub import whoami
+
+        hf_user = whoami().get("name")
+        if hf_user:
+            logging.info(
+                f"Trackio dashboard (live training progress): https://huggingface.co/spaces/{hf_user}/trackio"
+            )
+    except Exception:
+        pass
+
+
+TEACHER_MODEL_NAME = "sentence-transformers/all-mpnet-base-v2"
+STUDENT_MODEL_NAME = "distilbert/distilbert-base-uncased"
+
+CORPUS_DATASET = "sentence-transformers/all-nli"
+CORPUS_SUBSET = "pair"
+
+TRAIN_SIZE = 50_000
+EVAL_SIZE = 1_000
+OUTPUT_DIR = "models/distilbert-distilled-from-mpnet"
+RUN_NAME = "distilbert-distill-from-mpnet"
+
+TEACHER_ENCODE_BATCH_SIZE = 256
+TRAIN_BATCH_SIZE = 128
+SMOKE_TEST = os.environ.get("SMOKE_TEST") == "1"
+
+
+def setup_logging():
+    """Configure logging + TF32. Tees to logs/{RUN_NAME}.log and silences HTTP spam."""
+    os.makedirs("logs", exist_ok=True)
+    logging.basicConfig(
+        format="%(asctime)s - %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+        level=logging.INFO,
+        handlers=[logging.StreamHandler(), logging.FileHandler(f"logs/{RUN_NAME}.log")],
+        force=True,
+    )
+    for noisy in ("httpx", "httpcore", "huggingface_hub", "urllib3", "filelock", "fsspec"):
+        logging.getLogger(noisy).setLevel(logging.WARNING)
+    if torch.cuda.is_available():
+        torch.set_float32_matmul_precision("high")  # TF32 on Ampere+, no quality loss
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--eval-only", type=str, default=None, help="Skip training; load this saved model and run only the evaluator."
+    )
+    cli, _ = parser.parse_known_args()
+
+    setup_logging()
+
+    if cli.eval_only:
+        logging.info(f"Eval-only mode: loading model from {cli.eval_only}")
+        model = SentenceTransformer(cli.eval_only)
+        stsb = load_dataset("sentence-transformers/stsb", split="validation")
+        evaluator = EmbeddingSimilarityEvaluator(
+            sentences1=stsb["sentence1"],
+            sentences2=stsb["sentence2"],
+            scores=stsb["score"],
+            main_similarity=SimilarityFunction.COSINE,
+            name="sts-dev",
+        )
+        with autocast_ctx():
+            evaluator(model)
+        return
+
+    logging.info(f"Loading teacher: {TEACHER_MODEL_NAME}")
+    teacher = SentenceTransformer(TEACHER_MODEL_NAME)
+
+    logging.info(f"Loading student: {STUDENT_MODEL_NAME}")
+    student = SentenceTransformer(
+        STUDENT_MODEL_NAME,
+        model_card_data=SentenceTransformerModelCardData(
+            language="en",
+            license="apache-2.0",
+            model_name=f"{STUDENT_MODEL_NAME.split('/')[-1]} distilled from {TEACHER_MODEL_NAME.split('/')[-1]}",
+        ),
+    )
+    # Match the teacher's final Normalize. MSELoss against unit-norm targets fights student
+    # outputs at norm ~5-10 and can silently regress
+    if any(isinstance(m, Normalize) for m in teacher) and not any(isinstance(m, Normalize) for m in student):
+        student.append(Normalize())
+
+    if student.get_embedding_dimension() != teacher.get_embedding_dimension():
+        raise SystemExit(
+            f"Student dim ({student.get_embedding_dimension()}) != teacher dim "
+            f"({teacher.get_embedding_dimension()}). Plain MSELoss requires matching dims. "
+            "Either pick a student with matching dim, or add a Dense projection layer "
+            "(PCA-initialized from teacher embeddings). See the 'MISMATCHED EMBEDDING DIMS' "
+            "section in this script's docstring."
+        )
+
+    logging.info(f"Loading corpus: {CORPUS_DATASET} ({CORPUS_SUBSET})")
+    train_size = 50 if SMOKE_TEST else TRAIN_SIZE
+    eval_size = 20 if SMOKE_TEST else EVAL_SIZE
+    if SMOKE_TEST:
+        logging.info("SMOKE_TEST=1: trimmed dataset; will run max_steps=1 and skip Hub push")
+    raw = load_dataset(CORPUS_DATASET, CORPUS_SUBSET, split="train")
+    sentences = list(dict.fromkeys(s for row in raw for s in (row["anchor"], row["positive"]) if isinstance(s, str)))
+    sentences = sentences[: train_size + eval_size]
+    train_sentences = sentences[:train_size]
+    eval_sentences = sentences[train_size : train_size + eval_size]
+
+    logging.info(f"Encoding {len(train_sentences):,} training sentences with the teacher (may take a while)")
+    teacher_train = teacher.encode(
+        train_sentences, batch_size=TEACHER_ENCODE_BATCH_SIZE, convert_to_numpy=True, show_progress_bar=True
+    )
+
+    logging.info(f"Encoding {len(eval_sentences):,} eval sentences with the teacher")
+    teacher_eval = teacher.encode(
+        eval_sentences, batch_size=TEACHER_ENCODE_BATCH_SIZE, convert_to_numpy=True, show_progress_bar=True
+    )
+
+    train_dataset = Dataset.from_dict({"sentence": train_sentences, "label": teacher_train.tolist()})
+    eval_dataset = Dataset.from_dict({"sentence": eval_sentences, "label": teacher_eval.tolist()})
+
+    logging.info(f"Building training dataset ({len(train_dataset):,}) and eval dataset ({len(eval_dataset):,})")
+
+    loss = MSELoss(model=student)
+
+    logging.info("Setting up STS-B evaluator for quality tracking")
+    stsb = load_dataset("sentence-transformers/stsb", split="validation")
+    evaluator = EmbeddingSimilarityEvaluator(
+        sentences1=stsb["sentence1"],
+        sentences2=stsb["sentence2"],
+        scores=stsb["score"],
+        main_similarity=SimilarityFunction.COSINE,
+        name="sts-dev",
+    )
+    logging.info("Teacher performance:")
+    evaluator(teacher)
+    logging.info("Student performance before distillation:")
+    # Must run before deriving metric_key: evaluator(model) mutates primary_metric to add the name_ prefix.
+    baseline_eval = evaluator(student)[evaluator.primary_metric]
+    metric_key = f"eval_{evaluator.primary_metric}"
+
+    args = SentenceTransformerTrainingArguments(
+        output_dir=OUTPUT_DIR,
+        num_train_epochs=1,
+        max_steps=1 if SMOKE_TEST else -1,
+        per_device_train_batch_size=TRAIN_BATCH_SIZE,
+        per_device_eval_batch_size=TRAIN_BATCH_SIZE,
+        learning_rate=1e-4,
+        weight_decay=0.01,
+        warmup_steps=0.1,
+        bf16=True,
+        eval_strategy="steps",
+        eval_steps=0.1,
+        save_strategy="steps",
+        save_steps=0.1,
+        save_total_limit=2,
+        logging_steps=0.01,
+        logging_first_step=True,
+        load_best_model_at_end=True,
+        metric_for_best_model=metric_key,
+        greater_is_better=True,
+        report_to="none" if SMOKE_TEST else "trackio",
+        run_name=RUN_NAME,
+        seed=12,
+    )
+
+    trainer = SentenceTransformerTrainer(
+        model=student,
+        args=args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        loss=loss,
+        evaluator=evaluator,
+    )
+    if not SMOKE_TEST:
+        log_trackio_dashboard()
+    trainer.train()
+
+    logging.info("Student performance after distillation:")
+    score = evaluator(student)[evaluator.primary_metric]
+    delta = score - baseline_eval
+    verdict = "WIN" if delta >= 0.005 else "MARGINAL" if delta >= 0 else "REGRESSION"
+    logging.info(f"VERDICT: {verdict} | score={score:.4f} | baseline={baseline_eval:.4f} | delta={delta:+.4f}")
+
+    final_dir = f"{OUTPUT_DIR}/final"
+    student.save_pretrained(final_dir)
+    logging.info(f"Saved to {final_dir}")
+
+    if SMOKE_TEST:
+        logging.info("SMOKE_TEST=1: skipping Hub push")
+        return
+
+    try:
+        commit_url = student.push_to_hub(RUN_NAME)
+        logging.info(f"Pushed model to {commit_url.rsplit('/commit/', 1)[0]}")
+    except Exception:
+        import traceback
+
+        logging.error(f"Hub push failed:\n{traceback.format_exc()}")
+
+
+if __name__ == "__main__":
+    main()