📦 deps(thirdparty): update snapshots

antigravity-awesome-skills/skills/train-sentence-transformers/scripts/mine_hard_negatives.py

@@ -0,0 +1,199 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "sentence-transformers[train]>=5.0",
+#     "datasets>=2.19.0",
+# ]
+# ///
+"""Mine hard negatives as a pre-training step for contrastive losses.
+
+Hard negatives are the single highest-leverage lever for retrieval quality.
+This script is a thin, CLI-friendly wrapper around
+`sentence_transformers.util.mine_hard_negatives`.
+
+Typical workflow:
+1. Start from a dataset of (anchor, positive) pairs.
+2. Pick a retriever model (can be your current base model or a stronger one).
+3. Run this script to produce a new dataset with N mined negatives per anchor.
+4. Train with MultipleNegativesRankingLoss or CachedMultipleNegativesRankingLoss.
+
+Usage:
+    python mine_hard_negatives.py \\
+        --dataset sentence-transformers/gooaq \\
+        --model sentence-transformers/all-MiniLM-L6-v2 \\
+        --num-negatives 5 \\
+        --output-path data/gooaq-hard-negatives
+
+    # Mine from a separate document corpus (recommended for production):
+    python mine_hard_negatives.py \\
+        --dataset sentence-transformers/gooaq \\
+        --model sentence-transformers/all-MiniLM-L6-v2 \\
+        --corpus-dataset sentence-transformers/wikipedia-en-passages \\
+        --corpus-column text \\
+        --num-negatives 5 \\
+        --output-path data/gooaq-hn-wiki
+
+    # With a cross-encoder as an "oracle" to filter negatives by score:
+    python mine_hard_negatives.py \\
+        --dataset sentence-transformers/gooaq \\
+        --model sentence-transformers/all-MiniLM-L6-v2 \\
+        --cross-encoder cross-encoder/ms-marco-MiniLM-L-6-v2 \\
+        --num-negatives 5 \\
+        --max-score 0.9 \\
+        --relative-margin 0.05 \\
+        --output-path data/gooaq-hn-filtered
+
+    # Push the mined dataset to the Hub:
+    python mine_hard_negatives.py \\
+        --dataset sentence-transformers/gooaq --model ... --num-negatives 5 \\
+        --push-to-hub your-username/gooaq-hard-negatives
+
+Key options:
+    --num-negatives     How many hard negatives to mine per anchor (default 3).
+    --range-min/max     Which retrieval-rank window to sample from (default 0..100).
+    --sampling-strategy "top" (rank-1 hardest) or "random" (within the window).
+    --relative-margin   Require that negative_score < positive_score * (1 - margin).
+    --max-score         Filter candidates above this score (likely false negatives).
+    --cross-encoder     Use a cross-encoder to re-score candidates before filtering.
+    --corpus-dataset    Mine from a separate document pool instead of the input
+                        dataset's positives. Recommended for production: typical
+                        retrieval corpora (Wikipedia, MSMARCO passages) are far
+                        larger than your training-pair pool, giving harder negatives.
+
+See the `mine_hard_negatives` API reference for full semantics and all flags.
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import sys
+
+from datasets import load_dataset
+
+from sentence_transformers import CrossEncoder, SentenceTransformer
+from sentence_transformers.util import mine_hard_negatives
+
+logging.basicConfig(format="%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO)
+for _noisy in ("httpx", "httpcore", "huggingface_hub", "urllib3", "filelock", "fsspec"):
+    logging.getLogger(_noisy).setLevel(logging.WARNING)
+
+
+def build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--dataset", required=True)
+    p.add_argument("--subset", default=None)
+    p.add_argument("--split", default="train")
+    p.add_argument("--model", required=True, help="Retriever / bi-encoder used to score candidates")
+    p.add_argument("--cross-encoder", default=None, help="Optional CrossEncoder to re-score and filter")
+    p.add_argument("--anchor-column", default=None)
+    p.add_argument("--positive-column", default=None)
+    p.add_argument(
+        "--num-negatives",
+        type=int,
+        default=3,
+        help="Number of hard negatives to mine per anchor. Default 3 matches the library.",
+    )
+    p.add_argument("--range-min", type=int, default=0)
+    p.add_argument("--range-max", type=int, default=100)
+    p.add_argument("--sampling-strategy", choices=["top", "random"], default="top")
+    p.add_argument(
+        "--max-score", type=float, default=None, help="Drop candidates scoring above this (likely false negatives)"
+    )
+    p.add_argument("--min-score", type=float, default=None)
+    p.add_argument("--absolute-margin", type=float, default=None)
+    p.add_argument("--relative-margin", type=float, default=None)
+    p.add_argument(
+        "--output-format",
+        choices=["triplet", "n-tuple", "labeled-pair", "labeled-list"],
+        default="triplet",
+    )
+    p.add_argument("--include-positives", action="store_true")
+    p.add_argument("--output-scores", action="store_true")
+    p.add_argument("--batch-size", type=int, default=32)
+    p.add_argument("--use-faiss", action="store_true")
+    p.add_argument(
+        "--corpus-dataset",
+        default=None,
+        help="Optional Hub dataset id or local path for a separate document pool to mine from. "
+        "If unset, mines negatives from the input dataset's positives.",
+    )
+    p.add_argument("--corpus-subset", default=None, help="Subset of --corpus-dataset (optional)")
+    p.add_argument("--corpus-split", default="train", help="Split of --corpus-dataset (default 'train')")
+    p.add_argument(
+        "--corpus-column", default="text", help="Text column to extract from --corpus-dataset (default 'text')"
+    )
+    p.add_argument("--output-path", default=None, help="Local directory to save the mined dataset to")
+    p.add_argument("--push-to-hub", default=None, help="Hub repo id to push the mined dataset to (optional)")
+    p.add_argument("--private", action="store_true", help="Push as a private repo")
+    return p
+
+
+def main() -> int:
+    args = build_parser().parse_args()
+
+    dataset = (
+        load_dataset(args.dataset, args.subset, split=args.split)
+        if args.subset
+        else load_dataset(args.dataset, split=args.split)
+    )
+    print(f"Loaded {len(dataset):,} rows from {args.dataset} (split={args.split})")
+
+    model = SentenceTransformer(args.model)
+    cross_encoder = CrossEncoder(args.cross_encoder) if args.cross_encoder else None
+
+    corpus = None
+    if args.corpus_dataset:
+        corpus_ds = (
+            load_dataset(args.corpus_dataset, args.corpus_subset, split=args.corpus_split)
+            if args.corpus_subset
+            else load_dataset(args.corpus_dataset, split=args.corpus_split)
+        )
+        if args.corpus_column not in corpus_ds.column_names:
+            raise SystemExit(
+                f"--corpus-column '{args.corpus_column}' not in {args.corpus_dataset} columns: "
+                f"{corpus_ds.column_names}"
+            )
+        corpus = list(corpus_ds[args.corpus_column])
+        print(f"Loaded corpus: {len(corpus):,} documents from {args.corpus_dataset}.{args.corpus_column}")
+
+    mined = mine_hard_negatives(
+        dataset=dataset,
+        model=model,
+        corpus=corpus,
+        cross_encoder=cross_encoder,
+        anchor_column_name=args.anchor_column,
+        positive_column_name=args.positive_column,
+        num_negatives=args.num_negatives,
+        range_min=args.range_min,
+        range_max=args.range_max,
+        sampling_strategy=args.sampling_strategy,
+        max_score=args.max_score,
+        min_score=args.min_score,
+        absolute_margin=args.absolute_margin,
+        relative_margin=args.relative_margin,
+        output_format=args.output_format,
+        include_positives=args.include_positives,
+        output_scores=args.output_scores,
+        batch_size=args.batch_size,
+        use_faiss=args.use_faiss,
+    )
+    print(f"Mined dataset: {len(mined):,} rows | columns: {mined.column_names}")
+
+    if args.output_path:
+        mined.save_to_disk(args.output_path)
+        print(f"Saved to {args.output_path}")
+
+    if args.push_to_hub:
+        mined.push_to_hub(args.push_to_hub, private=args.private)
+        print(f"Pushed to https://huggingface.co/datasets/{args.push_to_hub}")
+
+    if not args.output_path and not args.push_to_hub:
+        print("No --output-path or --push-to-hub provided; nothing persisted.")
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

antigravity-awesome-skills/skills/train-sentence-transformers/scripts/train_cross_encoder_distillation_example.py

@@ -0,0 +1,254 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "sentence-transformers[train]>=5.0",
+#     "datasets>=2.19.0",
+#     "accelerate>=0.26.0",
+#     "trackio",
+# ]
+# ///
+"""CrossEncoder distillation: train a small reranker from a stronger one's scores.
+
+Uses MarginMSELoss on `(query, positive, negative, score_diff)` where
+`score_diff = teacher(q, pos) - teacher(q, neg)`. Workhorse of MS MARCO-style
+distillation: the student learns the teacher's score gaps without ever calling
+the teacher at training time (precomputed score diffs).
+
+For Embedding-MSE / Listwise-KL variants and the broader pattern map, see the
+sibling `train_sentence_transformer_distillation_example.py` docstring.
+
+CRITICAL: `activation_fn=nn.Identity()` is mandatory. The default `Sigmoid` (with
+`num_labels=1`) saturates raw logits >5 to ~1.0 inside `predict()` at eval time,
+silently collapsing eval ranking (training loss stays healthy while nDCG drops
+from e.g. ~0.59 to ~0.14). See `../references/troubleshooting.md` ("CrossEncoder
+eval nDCG crashes after distillation / listwise / pairwise training").
+
+Data: this script uses `sentence-transformers/msmarco` (`bert-ensemble-margin-mse`
+subset), which has precomputed teacher score diffs per (q, pos, neg) row. To
+distill from your own teacher, replace the dataset loading with a one-time
+teacher pass over your (q, pos, neg) triples and store `score_diff = teacher_pos
+- teacher_neg` as the label.
+
+Run locally:
+    pip install "sentence-transformers[train]>=5.0"
+    python train_cross_encoder_distillation_example.py
+
+Multi-GPU:
+    accelerate launch train_cross_encoder_distillation_example.py
+
+Hugging Face Jobs: paste this file's contents as the `script` in hf_jobs(...).
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import os
+from contextlib import nullcontext
+
+import torch
+import torch.nn as nn
+from datasets import load_dataset, load_from_disk
+from transformers import EarlyStoppingCallback
+
+from sentence_transformers import (
+    CrossEncoder,
+    CrossEncoderModelCardData,
+    CrossEncoderTrainer,
+    CrossEncoderTrainingArguments,
+)
+from sentence_transformers.cross_encoder.evaluation import CrossEncoderNanoBEIREvaluator
+from sentence_transformers.cross_encoder.losses import MarginMSELoss
+
+
+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
+
+
+MODEL_NAME = "microsoft/MiniLM-L12-H384-uncased"
+DATASET_NAME = "sentence-transformers/msmarco"
+DATASET_SUBSET = "bert-ensemble-margin-mse"
+TRAIN_SIZE = 100_000
+EVAL_SIZE = 5_000
+OUTPUT_DIR = "models/minilm-msmarco-distilled"
+RUN_NAME = "minilm-msmarco-distilled"
+DATA_CACHE = f"data/{RUN_NAME}-resolved"
+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")
+
+
+def load_resolved_dataset():
+    """Load (query, positive, negative, score) rows. The MSMARCO subset is keyed by
+    passage_id / query_id; resolve to text once and cache to disk so reruns skip the work."""
+    if os.path.isdir(DATA_CACHE):
+        logging.info(f"Loading cached resolved dataset from {DATA_CACHE}")
+        return load_from_disk(DATA_CACHE)
+
+    logging.info(f"Resolving {DATASET_NAME}/{DATASET_SUBSET} ids -> text (one-time, cached)")
+    corpus_ds = load_dataset(DATASET_NAME, "corpus", split="train")
+    corpus = dict(zip(corpus_ds["passage_id"], corpus_ds["passage"]))
+    queries_ds = load_dataset(DATASET_NAME, "queries", split="train")
+    queries = dict(zip(queries_ds["query_id"], queries_ds["query"]))
+    raw = load_dataset(DATASET_NAME, DATASET_SUBSET, split="train").select(range(TRAIN_SIZE + EVAL_SIZE))
+
+    def id_to_text(batch):
+        return {
+            "query": [queries[qid] for qid in batch["query_id"]],
+            "positive": [corpus[pid] for pid in batch["positive_id"]],
+            "negative": [corpus[pid] for pid in batch["negative_id"]],
+            "score": batch["score"],
+        }
+
+    resolved = raw.map(id_to_text, batched=True, remove_columns=["query_id", "positive_id", "negative_id"])
+    resolved.save_to_disk(DATA_CACHE)
+    return resolved
+
+
+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 = CrossEncoder(cli.eval_only)
+        evaluator = CrossEncoderNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"])
+        with autocast_ctx():
+            evaluator(model)
+        return
+
+    logging.info(f"Loading base model: {MODEL_NAME}")
+    model = CrossEncoder(
+        MODEL_NAME,
+        num_labels=1,
+        activation_fn=nn.Identity(),  # Mandatory for distillation losses.
+        model_card_data=CrossEncoderModelCardData(
+            language="en",
+            license="apache-2.0",
+            model_name=f"{MODEL_NAME.split('/')[-1]} reranker distilled from MS MARCO ensemble",
+        ),
+    )
+
+    resolved = load_resolved_dataset()
+    if SMOKE_TEST:
+        logging.info("SMOKE_TEST=1: trimmed dataset; will run max_steps=1 and skip Hub push")
+        resolved = resolved.select(range(min(70, len(resolved))))
+    eval_size = 20 if SMOKE_TEST else EVAL_SIZE
+    split = resolved.train_test_split(test_size=eval_size, seed=12)
+    train_dataset = split["train"]
+    eval_dataset = split["test"]
+    logging.info(f"  train: {len(train_dataset):,} rows | eval: {len(eval_dataset):,} rows")
+    logging.info(f"  columns: {train_dataset.column_names}")
+
+    loss = MarginMSELoss(model)
+
+    evaluator = CrossEncoderNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"])
+    logging.info("Baseline evaluation:")
+    with autocast_ctx():
+        # Must run before deriving metric_key: evaluator(model) mutates primary_metric to add the name_ prefix.
+        baseline_eval = evaluator(model)[evaluator.primary_metric]
+    metric_key = f"eval_{evaluator.primary_metric}"
+
+    args = CrossEncoderTrainingArguments(
+        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=8e-6,  # Lower than typical 2e-5; distillation regression converges faster
+        weight_decay=0.01,
+        warmup_steps=0.1,
+        lr_scheduler_type="linear",
+        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 = CrossEncoderTrainer(
+        model=model,
+        args=args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        loss=loss,
+        evaluator=evaluator,
+        callbacks=[EarlyStoppingCallback(early_stopping_patience=3)],  # CE rerankers peak mid-training
+    )
+    if not SMOKE_TEST:
+        log_trackio_dashboard()
+    trainer.train()
+
+    logging.info("Post-training evaluation:")
+    with autocast_ctx():
+        score = evaluator(model)[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)
+        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()

antigravity-awesome-skills/skills/train-sentence-transformers/scripts/train_cross_encoder_example.py

@@ -0,0 +1,253 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "sentence-transformers[train]>=5.0",
+#     "datasets>=2.19.0",
+#     "accelerate>=0.26.0",
+#     "trackio",
+# ]
+# ///
+"""Production-ready cross-encoder (reranker) training template.
+
+Demonstrates:
+- BinaryCrossEntropyLoss on (query, passage, label) pointwise data
+- Hard-negative mining (`mine_hard_negatives` with `output_format="labeled-pair"`)
+  to produce the labeled training data BCE needs, starting from (question, answer)
+  pairs
+- `pos_weight=num_negatives` to offset the positive/negative imbalance
+- CrossEncoderNanoBEIREvaluator for retrieval reranking metrics
+- load_best_model_at_end on the retrieval metric
+- Auto model card + optional Hub push
+
+Run locally:
+    pip install "sentence-transformers[train]>=5.0"
+    python train_cross_encoder_example.py
+
+Multi-GPU:
+    accelerate launch train_cross_encoder_example.py
+
+Hugging Face Jobs: paste this file's contents as the `script` in hf_jobs(...).
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import os
+from contextlib import nullcontext
+
+import torch
+from datasets import load_dataset, load_from_disk
+from transformers import EarlyStoppingCallback
+
+from sentence_transformers import (
+    CrossEncoder,
+    CrossEncoderModelCardData,
+    CrossEncoderTrainer,
+    CrossEncoderTrainingArguments,
+    SentenceTransformer,
+)
+from sentence_transformers.cross_encoder.evaluation import CrossEncoderNanoBEIREvaluator
+from sentence_transformers.cross_encoder.losses import BinaryCrossEntropyLoss
+from sentence_transformers.util import mine_hard_negatives
+
+
+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
+
+
+MODEL_NAME = "microsoft/MiniLM-L12-H384-uncased"
+DATASET_NAME = "sentence-transformers/gooaq"
+RETRIEVER_NAME = "sentence-transformers/static-retrieval-mrl-en-v1"
+TRAIN_SIZE = 100_000
+EVAL_SIZE = 1_000
+NUM_NEGATIVES = 5
+OUTPUT_DIR = "models/minilm-gooaq-ce"
+RUN_NAME = "minilm-gooaq-ce"
+HARD_NEG_CACHE = f"data/{RUN_NAME}-hard-negatives"  # delete this dir to remine
+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 = CrossEncoder(cli.eval_only)
+        evaluator = CrossEncoderNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"])
+        with autocast_ctx():
+            evaluator(model)
+        return
+
+    logging.info(f"Loading base model: {MODEL_NAME}")
+    model = CrossEncoder(
+        MODEL_NAME,
+        num_labels=1,
+        model_card_data=CrossEncoderModelCardData(
+            language="en",
+            license="apache-2.0",
+            model_name=f"{MODEL_NAME.split('/')[-1]} reranker finetuned on GooAQ",
+        ),
+    )
+
+    logging.info(f"Loading dataset: {DATASET_NAME}")
+    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")
+    pairs = load_dataset(DATASET_NAME, split="train").select(range(train_size + eval_size))
+
+    if os.path.isdir(HARD_NEG_CACHE):
+        logging.info(f"Loading cached mined hard negatives from {HARD_NEG_CACHE}")
+        labeled = load_from_disk(HARD_NEG_CACHE)
+    else:
+        logging.info(f"Mining hard negatives with {RETRIEVER_NAME}")
+        retriever = SentenceTransformer(RETRIEVER_NAME)
+        labeled = mine_hard_negatives(
+            dataset=pairs,
+            model=retriever,
+            num_negatives=NUM_NEGATIVES,
+            range_min=10,  # skip the top-10 (likely to contain true positives)
+            range_max=100,
+            sampling_strategy="top",
+            output_format="labeled-pair",
+            use_faiss=True,
+        )
+        labeled.save_to_disk(HARD_NEG_CACHE)
+        logging.info(f"Saved mined dataset to {HARD_NEG_CACHE} (delete to remine)")
+        del retriever
+        torch.cuda.empty_cache()
+    # EVAL_SIZE here counts labeled-pair rows, not distinct queries: each
+    # query contributes 1 positive + NUM_NEGATIVES negatives, so e.g. 1000
+    # rows is approximately 1000 / (1 + NUM_NEGATIVES) distinct queries.
+    split = labeled.train_test_split(test_size=eval_size, seed=12)
+    train_dataset = split["train"]
+    eval_dataset = split["test"]
+    logging.info(f"  train: {len(train_dataset):,} rows | eval: {len(eval_dataset):,} rows")
+    logging.info(f"  columns: {train_dataset.column_names}")
+
+    # pos_weight = negatives / positives, derived from the actual label distribution
+    # so it stays correct if rows get filtered or the mining ratio drifts.
+    n_pos = sum(1 for label in train_dataset["label"] if label > 0.5)
+    n_neg = len(train_dataset) - n_pos
+    pos_weight_value = n_neg / max(n_pos, 1)
+    logging.info(f"  positives: {n_pos:,} | negatives: {n_neg:,} | pos_weight: {pos_weight_value:.2f}")
+    loss = BinaryCrossEntropyLoss(model, pos_weight=torch.tensor(pos_weight_value))
+
+    evaluator = CrossEncoderNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"])
+    logging.info("Baseline evaluation:")
+    with autocast_ctx():
+        # Must run before deriving metric_key: evaluator(model) mutates primary_metric to add the name_ prefix.
+        baseline_eval = evaluator(model)[evaluator.primary_metric]
+    metric_key = f"eval_{evaluator.primary_metric}"
+
+    args = CrossEncoderTrainingArguments(
+        output_dir=OUTPUT_DIR,
+        num_train_epochs=1,
+        max_steps=1 if SMOKE_TEST else -1,
+        per_device_train_batch_size=64,
+        per_device_eval_batch_size=64,
+        learning_rate=2e-5,
+        weight_decay=0.01,
+        warmup_steps=0.1,
+        lr_scheduler_type="linear",
+        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,
+    )
+
+    # EarlyStoppingCallback earns its keep for cross-encoders: CE rerankers
+    # typically peak mid-training and then degrade, so stopping at the best
+    # eval checkpoint is load-bearing (unlike bi-encoders, which tend to
+    # plateau rather than regress).
+    trainer = CrossEncoderTrainer(
+        model=model,
+        args=args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        loss=loss,
+        evaluator=evaluator,
+        callbacks=[EarlyStoppingCallback(early_stopping_patience=3)],
+    )
+    if not SMOKE_TEST:
+        log_trackio_dashboard()
+    trainer.train()
+
+    logging.info("Post-training evaluation:")
+    with autocast_ctx():
+        score = evaluator(model)[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)
+        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()

antigravity-awesome-skills/skills/train-sentence-transformers/scripts/train_cross_encoder_listwise_example.py

@@ -0,0 +1,282 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "sentence-transformers[train]>=5.0",
+#     "datasets>=2.19.0",
+#     "accelerate>=0.26.0",
+#     "trackio",
+# ]
+# ///
+"""CrossEncoder listwise training with LambdaLoss.
+
+LambdaLoss is the state-of-the-art listwise ranking loss — it optimizes a
+surrogate of nDCG via weighted pairwise comparisons over a per-query candidate
+list. Use this when you have multiple candidates per query with graded
+relevance, and you want a stronger ranker than pointwise BCE.
+
+Data shape: `(query, [doc_1, ..., doc_K], [score_1, ..., score_K])` per row.
+This script builds it via `mine_hard_negatives(..., output_format="labeled-list")`
+starting from `(question, answer)` pairs: each row gets the positive plus K
+hard negatives, with binary scores (1 for positive, 0 for negatives).
+
+CRITICAL: `activation_fn=nn.Identity()` is mandatory for LambdaLoss / ListNet /
+ListMLE / PListMLE / RankNet / MarginMSE / MSE — anything that's not
+`BinaryCrossEntropyLoss` or `CrossEntropyLoss`. The default `Sigmoid` (with
+`num_labels=1`) saturates raw logits >5 to ~1.0 inside `predict()`, silently
+collapsing eval ranking. See `../references/troubleshooting.md` ("CrossEncoder
+eval nDCG crashes after distillation / listwise / pairwise training").
+
+OOM recovery for LambdaLoss: drop `mini_batch_size` first (chunking inside the
+loss preserves the K-list semantic), then `per_device_train_batch_size` paired
+with `gradient_accumulation_steps`, then reduce K (the per-query candidate-list
+length) only as a last resort. Lowering K changes the experiment.
+
+Run locally:
+    pip install "sentence-transformers[train]>=5.0"
+    python train_cross_encoder_listwise_example.py
+
+Multi-GPU:
+    accelerate launch train_cross_encoder_listwise_example.py
+
+Hugging Face Jobs: paste this file's contents as the `script` in hf_jobs(...).
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import os
+from contextlib import nullcontext
+
+import torch
+import torch.nn as nn
+from datasets import load_dataset, load_from_disk
+from transformers import EarlyStoppingCallback
+
+from sentence_transformers import (
+    CrossEncoder,
+    CrossEncoderModelCardData,
+    CrossEncoderTrainer,
+    CrossEncoderTrainingArguments,
+    SentenceTransformer,
+)
+from sentence_transformers.base.evaluation import SequentialEvaluator
+from sentence_transformers.cross_encoder.evaluation import (
+    CrossEncoderNanoBEIREvaluator,
+    CrossEncoderRerankingEvaluator,
+)
+from sentence_transformers.cross_encoder.losses import LambdaLoss
+from sentence_transformers.util import mine_hard_negatives
+
+
+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
+
+
+MODEL_NAME = "answerdotai/ModernBERT-base"
+DATASET_NAME = "sentence-transformers/gooaq"
+RETRIEVER_NAME = "sentence-transformers/static-retrieval-mrl-en-v1"
+TRAIN_SIZE = 100_000
+EVAL_SIZE = 1_000
+NUM_NEGATIVES = 7  # K-1 negatives + 1 positive per row
+EVAL_RERANK_DEPTH = 30  # candidates per query in the in-domain eval set
+OUTPUT_DIR = "models/modernbert-gooaq-lambda"
+RUN_NAME = "modernbert-gooaq-lambda"
+HARD_NEG_CACHE = f"data/{RUN_NAME}-hard-negatives"
+HARD_EVAL_CACHE = f"data/{RUN_NAME}-hard-eval"
+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")
+
+
+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 = CrossEncoder(cli.eval_only)
+        evaluator = CrossEncoderNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"])
+        with autocast_ctx():
+            evaluator(model)
+        return
+
+    logging.info(f"Loading base model: {MODEL_NAME}")
+    model = CrossEncoder(
+        MODEL_NAME,
+        num_labels=1,
+        activation_fn=nn.Identity(),  # Mandatory for LambdaLoss; Sigmoid would saturate eval logits.
+        model_card_data=CrossEncoderModelCardData(
+            language="en",
+            license="apache-2.0",
+            model_name=f"{MODEL_NAME.split('/')[-1]} reranker trained with LambdaLoss on GooAQ",
+        ),
+    )
+    # ModernBERT defaults to max_seq_length=8192, which allocates activation memory
+    # for 8192-token sequences regardless of input length. Pin to a (q, doc) cap.
+    model.max_seq_length = 512
+
+    full_dataset = load_dataset(DATASET_NAME, split="train").select(range(TRAIN_SIZE))
+    split = full_dataset.train_test_split(test_size=EVAL_SIZE, seed=12)
+    train_pairs, eval_pairs = split["train"], split["test"]
+
+    if os.path.isdir(HARD_NEG_CACHE) and os.path.isdir(HARD_EVAL_CACHE):
+        logging.info("Loading cached mined hard-negative datasets")
+        hard_train = load_from_disk(HARD_NEG_CACHE)
+        hard_eval = load_from_disk(HARD_EVAL_CACHE)
+    else:
+        logging.info(f"Mining hard negatives with {RETRIEVER_NAME}")
+        retriever = SentenceTransformer(RETRIEVER_NAME)
+        hard_train = mine_hard_negatives(
+            train_pairs,
+            retriever,
+            num_negatives=NUM_NEGATIVES,
+            range_min=10,
+            range_max=100,
+            sampling_strategy="top",
+            output_format="labeled-list",  # Listwise: (query, [docs], [scores])
+            use_faiss=True,
+            batch_size=4096,
+        )
+        hard_eval = mine_hard_negatives(
+            eval_pairs,
+            retriever,
+            corpus=full_dataset["answer"],
+            num_negatives=EVAL_RERANK_DEPTH,
+            output_format="n-tuple",
+            use_faiss=True,
+            batch_size=4096,
+        )
+        hard_train.save_to_disk(HARD_NEG_CACHE)
+        hard_eval.save_to_disk(HARD_EVAL_CACHE)
+        del retriever
+        torch.cuda.empty_cache()
+    if SMOKE_TEST:
+        logging.info("SMOKE_TEST=1: trimmed mined datasets; will run max_steps=1 and skip Hub push")
+        hard_train = hard_train.select(range(min(50, len(hard_train))))
+        hard_eval = hard_eval.select(range(min(20, len(hard_eval))))
+    logging.info(f"  train: {len(hard_train):,} rows | columns: {hard_train.column_names}")
+
+    loss = LambdaLoss(model=model, mini_batch_size=16)  # mini_batch_size: drop first if OOM
+
+    nano_beir = CrossEncoderNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"])
+    # Pure reranker quality: positive is in `documents` and `always_rerank_positives=True` (default).
+    in_domain = CrossEncoderRerankingEvaluator(
+        samples=[
+            {
+                "query": row["question"],
+                "positive": [row["answer"]],
+                "documents": [row["answer"]] + [row[col] for col in hard_eval.column_names[2:]],
+            }
+            for row in hard_eval
+        ],
+        batch_size=64,
+        name="gooaq-dev",
+    )
+    evaluator = SequentialEvaluator([in_domain, nano_beir])
+    logging.info("Baseline evaluation:")
+    with autocast_ctx():
+        baseline_eval = evaluator(model)[in_domain.primary_metric]
+
+    args = CrossEncoderTrainingArguments(
+        output_dir=OUTPUT_DIR,
+        num_train_epochs=1,
+        max_steps=1 if SMOKE_TEST else -1,
+        per_device_train_batch_size=64,
+        per_device_eval_batch_size=64,
+        learning_rate=2e-5,
+        weight_decay=0.01,
+        warmup_steps=0.1,
+        lr_scheduler_type="linear",
+        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=f"eval_{in_domain.primary_metric}",  # in-domain reranker > NanoBEIR
+        greater_is_better=True,
+        report_to="none" if SMOKE_TEST else "trackio",
+        run_name=RUN_NAME,
+        seed=12,
+    )
+
+    trainer = CrossEncoderTrainer(
+        model=model,
+        args=args,
+        train_dataset=hard_train,
+        loss=loss,
+        evaluator=evaluator,
+        callbacks=[EarlyStoppingCallback(early_stopping_patience=3)],
+    )
+    if not SMOKE_TEST:
+        log_trackio_dashboard()
+    trainer.train()
+
+    logging.info("Post-training evaluation:")
+    with autocast_ctx():
+        score = evaluator(model)[in_domain.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)
+        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()

antigravity-awesome-skills/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()

antigravity-awesome-skills/skills/train-sentence-transformers/scripts/train_sentence_transformer_example.py

@@ -0,0 +1,227 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "sentence-transformers[train]>=5.0",
+#     "datasets>=2.19.0",
+#     "accelerate>=0.26.0",
+#     "trackio",
+# ]
+# ///
+"""Production-ready bi-encoder (SentenceTransformer) training template.
+
+This script demonstrates a recommended setup:
+- MultipleNegativesRankingLoss on (anchor, positive, negative) triplets
+- NanoBEIREvaluator for retrieval metrics during training
+- BatchSamplers.NO_DUPLICATES (critical for MNRL)
+- 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]>=5.0"
+    python train_sentence_transformer_example.py
+
+    # Or with uv (no explicit install needed)
+    uv run train_sentence_transformer_example.py
+
+    # Multi-GPU
+    accelerate launch train_sentence_transformer_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 (
+    SentenceTransformer,
+    SentenceTransformerModelCardData,
+    SentenceTransformerTrainer,
+    SentenceTransformerTrainingArguments,
+)
+from sentence_transformers.base.sampler import BatchSamplers
+from sentence_transformers.sentence_transformer.evaluation import NanoBEIREvaluator
+from sentence_transformers.sentence_transformer.losses import MultipleNegativesRankingLoss
+
+
+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
+
+
+MODEL_NAME = "microsoft/mpnet-base"
+DATASET_NAME = "sentence-transformers/all-nli"
+DATASET_SUBSET = "triplet"
+TRAIN_SIZE = 50_000
+EVAL_SIZE = 1_000
+OUTPUT_DIR = "models/mpnet-base-all-nli"
+RUN_NAME = "mpnet-base-all-nli"
+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)
+        evaluator = NanoBEIREvaluator()
+        with autocast_ctx():
+            evaluator(model)
+        return
+
+    logging.info(f"Loading base model: {MODEL_NAME}")
+    model = SentenceTransformer(
+        MODEL_NAME,
+        model_card_data=SentenceTransformerModelCardData(
+            language="en",
+            license="apache-2.0",
+            model_name=f"{MODEL_NAME.split('/')[-1]} finetuned on AllNLI",
+        ),
+    )
+
+    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
+    train_dataset = load_dataset(DATASET_NAME, DATASET_SUBSET, split="train").select(range(train_size))
+    eval_dataset = load_dataset(DATASET_NAME, DATASET_SUBSET, split="dev").select(range(eval_size))
+    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")
+
+    loss = MultipleNegativesRankingLoss(model)
+
+    evaluator = NanoBEIREvaluator()
+    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.
+        baseline_eval = evaluator(model)[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=64,
+        per_device_eval_batch_size=64,
+        learning_rate=2e-5,
+        weight_decay=0.01,
+        warmup_steps=0.1,
+        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 = SentenceTransformerTrainer(
+        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:")
+    with autocast_ctx():
+        score = evaluator(model)[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()

antigravity-awesome-skills/skills/train-sentence-transformers/scripts/train_sentence_transformer_make_multilingual_example.py

@@ -0,0 +1,308 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "sentence-transformers[train]>=5.0",
+#     "datasets>=2.19.0",
+#     "accelerate>=0.26.0",
+#     "trackio",
+# ]
+# ///
+"""Multilingual teacher-student distillation: extend an English bi-encoder to other languages.
+
+The trick (Reimers & Gurevych 2020, https://huggingface.co/papers/2004.09813):
+the teacher embeds English; the multilingual student is trained so that BOTH
+the English sentence AND its translation map to the SAME teacher embedding.
+Cross-lingual retrieval works out of the box because translations sit near
+each other in the joint space.
+
+Use when you have a strong English bi-encoder and want a multilingual version
+but lack in-language supervised data. If you DO have in-language labeled data,
+train directly with MNRL / CoSENTLoss on it; that usually wins on in-language
+tasks.
+
+Data: parallel `(english, non_english)` pairs. `sentence-transformers/parallel-sentences-*`
+covers many corpora (talks, europarl, tatoeba, wikimatrix, opensubtitles, jw300,
+news-commentary, ...) with `{src}-{tgt}` subsets. ~500k pairs per language is
+plenty.
+
+Picks:
+- Teacher: any English bi-encoder you want a multilingual copy of
+  (all-mpnet-base-v2, all-MiniLM-L6-v2, BAAI/bge-base-en-v1.5, intfloat/e5-base-v2).
+- Student: must be multilingual (xlm-roberta-base, paraphrase-multilingual-MiniLM-L12-v2,
+  microsoft/mdeberta-v3-base).
+- Student dim must match teacher dim, otherwise add a PCA-init Dense projection
+  (see train_sentence_transformer_distillation_example.py).
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import os
+from contextlib import nullcontext
+
+import numpy as np
+import torch
+from datasets import DatasetDict, load_dataset
+
+from sentence_transformers import (
+    SentenceTransformer,
+    SentenceTransformerModelCardData,
+    SentenceTransformerTrainer,
+    SentenceTransformerTrainingArguments,
+)
+from sentence_transformers.sentence_transformer.evaluation import (
+    MSEEvaluator,
+    SequentialEvaluator,
+    TranslationEvaluator,
+)
+from sentence_transformers.sentence_transformer.losses import MSELoss
+from sentence_transformers.sentence_transformer.modules import Normalize
+
+
+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 = "FacebookAI/xlm-roberta-base"
+
+PARALLEL_DATASET = "sentence-transformers/parallel-sentences-talks"
+SOURCE_LANGUAGE = "en"
+TARGET_LANGUAGES = ("de", "es", "fr", "it")
+
+MAX_SENTENCES_PER_LANGUAGE = 200_000
+EVAL_SENTENCES_PER_LANGUAGE = 1_000
+STUDENT_MAX_SEQ_LENGTH = 128
+
+OUTPUT_DIR = "models/xlm-roberta-multilingual-from-mpnet"
+RUN_NAME = "xlm-roberta-multilingual-from-mpnet"
+
+TEACHER_ENCODE_BATCH_SIZE = 256
+TRAIN_BATCH_SIZE = 64
+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 load_parallel_data() -> tuple[DatasetDict, DatasetDict]:
+    """Load (en, target_lang) parallel pairs for each target language as a DatasetDict."""
+    train_dict = DatasetDict()
+    eval_dict = DatasetDict()
+    for tgt in TARGET_LANGUAGES:
+        subset = f"{SOURCE_LANGUAGE}-{tgt}"
+        try:
+            train_ds = load_dataset(PARALLEL_DATASET, subset, split="train")
+        except Exception as exc:
+            logging.error(f"Could not load {PARALLEL_DATASET}/{subset}: {exc}")
+            continue
+        if len(train_ds) > MAX_SENTENCES_PER_LANGUAGE:
+            train_ds = train_ds.select(range(MAX_SENTENCES_PER_LANGUAGE))
+
+        try:
+            eval_ds = load_dataset(PARALLEL_DATASET, subset, split="dev").select(range(EVAL_SENTENCES_PER_LANGUAGE))
+        except Exception:
+            split = train_ds.train_test_split(test_size=EVAL_SENTENCES_PER_LANGUAGE, shuffle=True, seed=12)
+            train_ds, eval_ds = split["train"], split["test"]
+
+        train_dict[subset] = train_ds
+        eval_dict[subset] = eval_ds
+    if not train_dict:
+        raise SystemExit(f"No language subsets loaded from {PARALLEL_DATASET}. Check TARGET_LANGUAGES.")
+    return train_dict, eval_dict
+
+
+def build_evaluator(eval_dict: DatasetDict, teacher: SentenceTransformer) -> SequentialEvaluator:
+    """Per-language MSE + TranslationEvaluator. `main_score_function` averages
+    translation accuracies only; MSE (`negative_mse * 100`) is on a different
+    scale and would break the verdict threshold if mixed in."""
+    sub_evaluators = []
+    for subset, ds in eval_dict.items():
+        sub_evaluators.append(
+            MSEEvaluator(
+                source_sentences=ds["english"],
+                target_sentences=ds["non_english"],
+                name=subset,
+                teacher_model=teacher,
+                batch_size=TEACHER_ENCODE_BATCH_SIZE,
+            )
+        )
+        sub_evaluators.append(
+            TranslationEvaluator(
+                source_sentences=ds["english"],
+                target_sentences=ds["non_english"],
+                name=subset,
+                batch_size=TEACHER_ENCODE_BATCH_SIZE,
+            )
+        )
+    # Sub-evaluators alternate MSE / Translation per language; scores[1::2] are the translation accuracies.
+    return SequentialEvaluator(sub_evaluators, main_score_function=lambda scores: float(np.mean(scores[1::2])))
+
+
+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}")
+        student = SentenceTransformer(cli.eval_only)
+        teacher = SentenceTransformer(TEACHER_MODEL_NAME)
+        _, eval_dict = load_parallel_data()
+        evaluator = build_evaluator(eval_dict, teacher)
+        with autocast_ctx():
+            evaluator(student)
+        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=[SOURCE_LANGUAGE, *TARGET_LANGUAGES],
+            license="apache-2.0",
+            model_name=f"{STUDENT_MODEL_NAME.split('/')[-1]} multilingual from {TEACHER_MODEL_NAME.split('/')[-1]}",
+        ),
+    )
+    student.max_seq_length = STUDENT_MAX_SEQ_LENGTH
+    # 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()}). MSELoss requires matching dims. "
+            "Either pick a student with matching dim, or add a Dense projection layer "
+            "(see train_sentence_transformer_distillation_example.py 'MISMATCHED EMBEDDING DIMS')."
+        )
+
+    logging.info("Loading parallel data")
+    train_dict, eval_dict = load_parallel_data()
+    if SMOKE_TEST:
+        logging.info("SMOKE_TEST=1: trimming each language subset; will run max_steps=1 and skip Hub push")
+        train_dict = DatasetDict({k: v.select(range(min(50, len(v)))) for k, v in train_dict.items()})
+        eval_dict = DatasetDict({k: v.select(range(min(20, len(v)))) for k, v in eval_dict.items()})
+
+    def attach_teacher_label(batch):
+        return {
+            "english": batch["english"],
+            "non_english": batch["non_english"],
+            "label": teacher.encode(batch["english"], batch_size=TEACHER_ENCODE_BATCH_SIZE, show_progress_bar=False),
+        }
+
+    column_names = list(train_dict.values())[0].column_names
+    logging.info("Encoding training English sentences with teacher (cached on disk if you save_to_disk)")
+    train_dict = train_dict.map(attach_teacher_label, batched=True, batch_size=10_000, remove_columns=column_names)
+    eval_dict = eval_dict.map(attach_teacher_label, batched=True, batch_size=10_000, remove_columns=column_names)
+
+    loss = MSELoss(model=student)
+
+    evaluator = build_evaluator(eval_dict, teacher)
+    logging.info("Student baseline (before training):")
+    with autocast_ctx():
+        baseline_eval = evaluator(student)["sequential_score"]
+
+    args = SentenceTransformerTrainingArguments(
+        output_dir=OUTPUT_DIR,
+        num_train_epochs=3,
+        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=2e-5,
+        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="eval_sequential_score",
+        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_dict,
+        eval_dataset=eval_dict,
+        loss=loss,
+        evaluator=evaluator,
+    )
+    if not SMOKE_TEST:
+        log_trackio_dashboard()
+    trainer.train()
+
+    logging.info("Final student evaluation:")
+    with autocast_ctx():
+        score = evaluator(student)["sequential_score"]
+    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()

antigravity-awesome-skills/skills/train-sentence-transformers/scripts/train_sentence_transformer_matryoshka_example.py

@@ -0,0 +1,212 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "sentence-transformers[train]>=5.0",
+#     "datasets>=2.19.0",
+#     "accelerate>=0.26.0",
+#     "trackio",
+# ]
+# ///
+"""Matryoshka (MRL) training: train once, deploy at multiple embedding dimensions.
+
+MatryoshkaLoss wraps a base loss and optimizes it at several truncated dimensions
+simultaneously. At inference, load with `truncate_dim=<target>` to get that size
+with ~95% of full-dim quality.
+
+Typical use: train at [768, 512, 256, 128, 64], deploy at 128 for 6x smaller
+index + 6x faster ANN with minimal quality loss.
+
+Run locally:
+    pip install "sentence-transformers[train]>=5.0"
+    python train_sentence_transformer_matryoshka_example.py
+"""
+
+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 (
+    SentenceTransformer,
+    SentenceTransformerTrainer,
+    SentenceTransformerTrainingArguments,
+)
+from sentence_transformers.base.sampler import BatchSamplers
+from sentence_transformers.sentence_transformer.evaluation import (
+    EmbeddingSimilarityEvaluator,
+    NanoBEIREvaluator,
+    SequentialEvaluator,
+)
+from sentence_transformers.sentence_transformer.losses import MatryoshkaLoss, MultipleNegativesRankingLoss
+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
+
+
+MODEL_NAME = "microsoft/mpnet-base"
+MATRYOSHKA_DIMS = [768, 512, 256, 128, 64]
+OUTPUT_DIR = "models/mpnet-matryoshka"
+RUN_NAME = "mpnet-matryoshka"
+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)
+        evaluator = NanoBEIREvaluator()
+        with autocast_ctx():
+            evaluator(model)
+        return
+
+    model = SentenceTransformer(MODEL_NAME)
+
+    train_size = 50 if SMOKE_TEST else 50_000
+    eval_size = 20 if SMOKE_TEST else 1_000
+    if SMOKE_TEST:
+        logging.info("SMOKE_TEST=1: trimmed dataset; will run max_steps=1 and skip Hub push")
+    train_dataset = load_dataset("sentence-transformers/all-nli", "triplet", split="train").select(range(train_size))
+    eval_dataset = load_dataset("sentence-transformers/all-nli", "triplet", split="dev").select(range(eval_size))
+
+    inner_loss = MultipleNegativesRankingLoss(model)
+    loss = MatryoshkaLoss(model, inner_loss, matryoshka_dims=MATRYOSHKA_DIMS)
+
+    stsb = load_dataset("sentence-transformers/stsb", split="validation")
+    per_dim_evaluators = [
+        EmbeddingSimilarityEvaluator(
+            sentences1=stsb["sentence1"],
+            sentences2=stsb["sentence2"],
+            scores=stsb["score"],
+            main_similarity=SimilarityFunction.COSINE,
+            name=f"sts-dev-{dim}",
+            truncate_dim=dim,
+        )
+        for dim in MATRYOSHKA_DIMS
+    ]
+    evaluator = SequentialEvaluator(
+        [*per_dim_evaluators, NanoBEIREvaluator()],
+        main_score_function=lambda scores: scores[0],
+    )
+    logging.info("Baseline evaluation (before training):")
+    with autocast_ctx():
+        # Must run before deriving metric_key: each sub-evaluator mutates its primary_metric to add the name_ prefix.
+        baseline_result = evaluator(model)
+    # Drive on the first per-dim evaluator's metric (matches main_score_function above).
+    metric_key = f"eval_{per_dim_evaluators[0].primary_metric}"
+    baseline_eval = baseline_result[per_dim_evaluators[0].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=128,
+        per_device_eval_batch_size=128,
+        learning_rate=2e-5,
+        weight_decay=0.01,
+        warmup_steps=0.1,
+        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 = SentenceTransformerTrainer(
+        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:")
+    with autocast_ctx():
+        score = evaluator(model)[per_dim_evaluators[0].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 to {final_dir}")
+    logging.info(f"To use at a specific dimension, load with: SentenceTransformer({final_dir!r}, truncate_dim=128)")
+
+    if SMOKE_TEST:
+        logging.info("SMOKE_TEST=1: skipping Hub push")
+        return
+
+    try:
+        commit_url = model.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()

antigravity-awesome-skills/skills/train-sentence-transformers/scripts/train_sentence_transformer_multi_dataset_example.py

@@ -0,0 +1,258 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "sentence-transformers[train]>=5.0",
+#     "datasets>=2.19.0",
+#     "accelerate>=0.26.0",
+#     "trackio",
+# ]
+# ///
+"""Multi-dataset / multi-task training: train one model on several datasets at once.
+
+Pass `train_dataset` (and optionally `eval_dataset`) as dicts. Pass `loss` as
+either a dict keyed by the same names (one loss per dataset, this script's
+default — Variant A) or a single loss instance applied to every dataset
+(Variant B). Dict keys are arbitrary but must match exactly across all three
+dicts; they show up in log output as `loss_all-nli=...`, `loss_stsb=...`.
+
+Three reasons to use multi-dataset training:
+- Multi-task: combine datasets with different signals (retrieval + STS +
+  classification) into a single general-purpose embedder.
+- Data augmentation: add a supplementary dataset (STS labels alongside your
+  main retrieval pairs) as a regularizer.
+- Domain coverage: train on several domains at once rather than sequentially.
+
+Variant A (this script): different shapes per dataset, so each needs its own
+matching loss. Pass `loss` as a dict; the trainer dispatches per-dataset and
+mixing loss arities (MNRL with 3 inputs + CoSENTLoss with 2+label) is fine.
+
+Variant B: same shape, same loss, but you want each mini-batch drawn from a
+single domain (so MNRL in-batch negatives stay in-domain and remain genuinely
+hard). Pass ONE loss and a dict of datasets:
+
+    train_datasets = {"medical": medical_pairs, "legal": legal_pairs, "code": code_pairs}
+    loss = MultipleNegativesRankingLoss(model)
+    trainer = SentenceTransformerTrainer(model=model, args=args,
+        train_dataset=train_datasets, loss=loss, ...)
+
+The multi-dataset batch sampler draws each batch from a single dataset, so a
+3-domain MNRL run gets in-domain negatives by construction. Counter-intuitive
+benefit: DatasetDict can outperform `concatenate_datasets` even with losses
+that don't share across the batch (e.g. LambdaLoss in cross-encoder training).
+
+Multi-dataset samplers:
+- `PROPORTIONAL` (default): sample from each dataset in proportion to its size.
+  Every row is seen ~once per epoch. Bias toward the largest dataset.
+- `ROUND_ROBIN`: alternate evenly; training stops when the SMALLEST is
+  exhausted. Equal screen-time per task.
+Common pattern: `PROPORTIONAL` for 1 epoch, then `ROUND_ROBIN` for a second
+if a smaller task's loss is still decreasing.
+
+Per-dataset prompts (bi-encoder, sparse-encoder): pass `prompts={"all-nli": "",
+"stsb": "Represent ...: ", "msmarco": {"query": "query: ", "positive":
+"passage: ", ...}}` to TrainingArguments. The nested per-column form works for
+bi-encoder and sparse-encoder; cross-encoders support single-value or
+per-dataset only. See `../references/prompts_and_instructions.md`.
+
+Eval metric aggregation: with a dict `eval_dataset`, each dataset's loss is
+logged separately (`eval_loss_all-nli`, `eval_loss_stsb`). The evaluator runs
+on the full model, so its metrics aren't per-dataset unless you wrap a
+`SequentialEvaluator` with per-dataset sub-evaluators. Set
+`metric_for_best_model` to a single evaluator metric, NOT a per-dataset loss.
+
+Gotchas: keys must match EXACTLY across all three dicts (train/eval/loss) or
+training fails at step 0; `NO_DUPLICATES` + `PROPORTIONAL` works (deduplicates
+within each batch regardless of source dataset); `ROUND_ROBIN` with uneven
+dataset sizes means `num_train_epochs=N` is N passes over the SMALLEST — use
+`PROPORTIONAL` or `max_steps` if you want N passes over the largest.
+"""
+
+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 (
+    SentenceTransformer,
+    SentenceTransformerTrainer,
+    SentenceTransformerTrainingArguments,
+)
+from sentence_transformers.base.sampler import BatchSamplers
+from sentence_transformers.sentence_transformer.evaluation import (
+    EmbeddingSimilarityEvaluator,
+    NanoBEIREvaluator,
+)
+from sentence_transformers.sentence_transformer.losses import (
+    CoSENTLoss,
+    MultipleNegativesRankingLoss,
+)
+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
+
+
+RUN_NAME = "mpnet-nli-stsb"
+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)
+        evaluator = NanoBEIREvaluator()
+        with autocast_ctx():
+            evaluator(model)
+        return
+
+    model = SentenceTransformer("microsoft/mpnet-base")
+
+    if SMOKE_TEST:
+        logging.info("SMOKE_TEST=1: trimmed datasets; will run max_steps=1 and skip Hub push")
+    nli_train_size = 50 if SMOKE_TEST else 50_000
+    nli_eval_size = 20 if SMOKE_TEST else 500
+    nli_train = load_dataset("sentence-transformers/all-nli", "triplet", split="train").select(range(nli_train_size))
+    stsb_train = load_dataset("sentence-transformers/stsb", split="train")
+    if SMOKE_TEST:
+        stsb_train = stsb_train.select(range(min(50, len(stsb_train))))
+    nli_eval = load_dataset("sentence-transformers/all-nli", "triplet", split="dev").select(range(nli_eval_size))
+    stsb_eval = load_dataset("sentence-transformers/stsb", split="validation")
+    if SMOKE_TEST:
+        stsb_eval = stsb_eval.select(range(min(20, len(stsb_eval))))
+
+    train_datasets = {"all-nli": nli_train, "stsb": stsb_train}
+    eval_datasets = {"all-nli": nli_eval, "stsb": stsb_eval}
+
+    losses = {
+        "all-nli": MultipleNegativesRankingLoss(model),
+        "stsb": CoSENTLoss(model),
+    }
+
+    evaluator = EmbeddingSimilarityEvaluator(
+        sentences1=stsb_eval["sentence1"],
+        sentences2=stsb_eval["sentence2"],
+        scores=stsb_eval["score"],
+        main_similarity=SimilarityFunction.COSINE,
+        name="sts-dev",
+    )
+    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.
+        baseline_eval = evaluator(model)[evaluator.primary_metric]
+    metric_key = f"eval_{evaluator.primary_metric}"
+
+    # multi_dataset_batch_sampler defaults to PROPORTIONAL (samples each dataset
+    # in proportion to its size). To force equal alternation between datasets:
+    #   from sentence_transformers.base.sampler import MultiDatasetBatchSamplers
+    #   ... multi_dataset_batch_sampler=MultiDatasetBatchSamplers.ROUND_ROBIN ...
+    args = SentenceTransformerTrainingArguments(
+        output_dir="models/mpnet-nli-stsb",
+        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=2e-5,
+        weight_decay=0.01,
+        warmup_steps=0.1,
+        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="mpnet-nli-stsb",
+        seed=12,
+    )
+
+    trainer = SentenceTransformerTrainer(
+        model=model,
+        args=args,
+        train_dataset=train_datasets,
+        eval_dataset=eval_datasets,
+        loss=losses,
+        evaluator=evaluator,
+    )
+    if not SMOKE_TEST:
+        log_trackio_dashboard()
+    trainer.train()
+
+    logging.info("Post-training evaluation:")
+    with autocast_ctx():
+        score = evaluator(model)[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}")
+
+    model.save_pretrained("models/mpnet-nli-stsb/final")
+
+    if SMOKE_TEST:
+        logging.info("SMOKE_TEST=1: skipping Hub push")
+        return
+
+    try:
+        commit_url = model.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()

antigravity-awesome-skills/skills/train-sentence-transformers/scripts/train_sentence_transformer_static_embedding_example.py

@@ -0,0 +1,273 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "sentence-transformers[train]>=5.0",
+#     "datasets>=2.19.0",
+#     "accelerate>=0.26.0",
+#     "trackio",
+#     "tokenizers>=0.20",
+#     "model2vec",  # only needed when WARMSTART=True (StaticEmbedding.from_model2vec)
+# ]
+# ///
+"""Train a StaticEmbedding model on a contrastive dataset.
+
+StaticEmbedding is a token-bag model: a per-token embedding table averaged over
+the tokens of an input. No transformer, no attention. Inference is ~20x faster
+on GPU and ~80x faster on CPU than a small encoder, with surprisingly competitive
+quality on retrieval benchmarks when trained on >=1M contrastive pairs.
+
+Two init paths via `WARMSTART` constant:
+- `WARMSTART=False` (default): random init. Use with >=1M contrastive pairs;
+  reaches a higher ceiling than warm-start when given enough data. The default
+  dataset below (GooAQ, ~3M pairs) is comfortably in this regime.
+- `WARMSTART=True`: `StaticEmbedding.from_model2vec(...)` — distil from a
+  model2vec checkpoint. Flip to True if you swap in a smaller dataset (<1M
+  pairs); converges faster and reaches better quality at lower data scales.
+
+Demonstrates:
+- MultipleNegativesRankingLoss wrapped in MatryoshkaLoss for nested embedding dims
+- Large batch size (1024+) with a high LR (~2e-1 for random init, ~5e-2 for warm-
+  start) since the loss surface for a token-bag is much flatter than for a
+  pretrained encoder
+- BatchSamplers.NO_DUPLICATES (load-bearing for in-batch negatives with duplicated
+  anchors)
+- NanoBEIREvaluator at full embedding dim
+- Auto model card + optional Hub push
+
+Run locally (CPU works for inference, but training needs a GPU for batch=1024+):
+    pip install "sentence-transformers[train]>=5.0"
+    python train_sentence_transformer_static_embedding_example.py
+
+Multi-GPU:
+    accelerate launch train_sentence_transformer_static_embedding_example.py
+
+Hugging Face Jobs: paste this file's contents as the `script` in hf_jobs(...).
+
+References:
+- HF blog post: https://huggingface.co/blog/static-embeddings
+- Module docs: sentence_transformers.sentence_transformer.modules.StaticEmbedding
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import os
+from contextlib import nullcontext
+
+import datasets
+import torch
+from datasets import load_dataset
+from tokenizers import Tokenizer
+
+from sentence_transformers import (
+    SentenceTransformer,
+    SentenceTransformerModelCardData,
+    SentenceTransformerTrainer,
+    SentenceTransformerTrainingArguments,
+)
+from sentence_transformers.base.sampler import BatchSamplers
+from sentence_transformers.sentence_transformer.evaluation import NanoBEIREvaluator
+from sentence_transformers.sentence_transformer.losses import (
+    MatryoshkaLoss,
+    MultipleNegativesRankingLoss,
+)
+from sentence_transformers.sentence_transformer.modules import StaticEmbedding
+
+
+def autocast_ctx():
+    """bf16/fp16 autocast for evaluator calls outside the trainer."""
+    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
+
+
+TOKENIZER_NAME = "google-bert/bert-base-uncased"
+EMBEDDING_DIM = 1024
+MATRYOSHKA_DIMS = [1024, 512, 256, 128, 64, 32]  # ordered largest-first per MatryoshkaLoss
+
+# False: random init (recommended for >=1M pairs; reaches a higher ceiling).
+# True: warm-start from a model2vec checkpoint (recommended for <1M pairs).
+# Default False because the example dataset (GooAQ, ~3M pairs) is well above the threshold.
+WARMSTART = False
+WARMSTART_MODEL2VEC = "minishlab/potion-base-8M"
+
+OUTPUT_DIR = "models/static-embedding-bert-uncased"
+RUN_NAME = "static-embedding-bert-uncased"
+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")
+
+
+def load_pair_dataset() -> datasets.Dataset:
+    """Load a contrastive-pair dataset for training.
+
+    StaticEmbedding starts from random initialization, so it needs *a lot* of
+    contrastive signal to converge. GooAQ alone provides ~3M (question, answer)
+    pairs, comfortably over the >=1M threshold below which a warm-start would
+    beat random init. For stronger production models, concatenate more sources
+    (NaturalQuestions, MSMARCO, MIRACL, etc.) and shuffle, in the same family of
+    sources used in `sentence-transformers/static-retrieval-mrl-en-v1`.
+    """
+    return (
+        load_dataset("sentence-transformers/gooaq", split="train")
+        .rename_columns({"question": "anchor", "answer": "positive"})
+        .select_columns(["anchor", "positive"])
+    )
+
+
+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)
+        evaluator = NanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"])
+        with autocast_ctx():
+            evaluator(model)
+        return
+
+    if WARMSTART:
+        logging.info(f"Warm-starting StaticEmbedding from model2vec: {WARMSTART_MODEL2VEC}")
+        # `StaticEmbedding.from_distillation("<bi-encoder>", vocabulary=...)` is the
+        # alternative warm-start path (distil from a stronger teacher's vectors); pick
+        # one. model2vec is faster to load and converges quickly on smaller datasets.
+        static_embedding = StaticEmbedding.from_model2vec(WARMSTART_MODEL2VEC)
+    else:
+        logging.info(f"Random-init StaticEmbedding from {TOKENIZER_NAME} tokenizer (dim={EMBEDDING_DIM})")
+        tokenizer = Tokenizer.from_pretrained(TOKENIZER_NAME)
+        static_embedding = StaticEmbedding(tokenizer, embedding_dim=EMBEDDING_DIM)
+    model = SentenceTransformer(
+        modules=[static_embedding],
+        model_card_data=SentenceTransformerModelCardData(
+            language="en",
+            license="apache-2.0",
+            model_name=f"Static embedding ({EMBEDDING_DIM}d) trained on contrastive pairs",
+        ),
+    )
+
+    logging.info("Loading + concatenating training datasets")
+    full = load_pair_dataset()
+    if SMOKE_TEST:
+        logging.info("SMOKE_TEST=1: trimmed dataset; will run max_steps=1 and skip Hub push")
+        full = full.select(range(min(200, len(full))))
+    eval_size = 20 if SMOKE_TEST else 10_000
+    split = full.train_test_split(test_size=eval_size, seed=12)
+    train_dataset = split["train"]
+    eval_dataset = split["test"]
+    logging.info(f"  train: {len(train_dataset):,} rows | eval: {len(eval_dataset):,} rows")
+    logging.info(f"  columns: {train_dataset.column_names}")
+
+    inner = MultipleNegativesRankingLoss(model)
+    loss = MatryoshkaLoss(model, inner, matryoshka_dims=MATRYOSHKA_DIMS)
+
+    evaluator = NanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"])
+    logging.info("Baseline evaluation (random init scores near zero; warm-start scores 0.3+):")
+    with autocast_ctx():
+        # Must run before deriving metric_key: evaluator(model) mutates primary_metric to add the name_ prefix.
+        baseline_eval = evaluator(model)[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=2048,
+        per_device_eval_batch_size=2048,
+        learning_rate=5e-2
+        if WARMSTART
+        else 2e-1,  # warm-start needs less LR; both far higher than encoder fine-tuning
+        weight_decay=0.0,  # weight decay on a token-bag is usually harmful
+        warmup_steps=0.1,
+        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 = SentenceTransformerTrainer(
+        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:")
+    with autocast_ctx():
+        score = evaluator(model)[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)
+        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()

antigravity-awesome-skills/skills/train-sentence-transformers/scripts/train_sentence_transformer_with_lora_example.py

@@ -0,0 +1,275 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "sentence-transformers[train]>=5.0",
+#     "peft>=0.7.0",
+#     "datasets>=2.19.0",
+#     "accelerate>=0.26.0",
+#     "trackio",
+# ]
+# ///
+"""LoRA / PEFT adapter training: memory-efficient fine-tuning.
+
+Instead of training all parameters, LoRA injects small low-rank adapter
+matrices and trains only those (~50-100x fewer trainable params). Works out
+of the box with sentence-transformers via `peft`.
+
+Use LoRA when: the base is a large decoder (Qwen3, Llama, Mistral, Gemma) and
+a full fine-tune is VRAM-prohibitive; you want multiple task-specific adapters
+on one base; you're nudging an existing strong retriever (E5-mistral,
+Qwen3-Embedding) on domain data. Skip LoRA for small encoders (BERT-base,
+MiniLM — full fine-tune is tractable and usually better) or tiny datasets
+(<1k pairs — adapter rank becomes the bottleneck).
+
+This template defaults to BERT-base for portability. Swap `MODEL_NAME` to a
+decoder backbone for the use case LoRA actually shines on.
+
+Architecture variants:
+- Bi-encoder (this script): `TaskType.FEATURE_EXTRACTION`.
+- Sparse-encoder: same pattern; `SparseEncoder` supports `add_adapter`. See
+  `examples/sparse_encoder/training/peft/train_splade_gooaq_peft.py`.
+- Cross-encoder: use `TaskType.SEQ_CLS` for `num_labels >= 1`. Community
+  examples are sparse; smoke-test with `max_steps=1` first.
+
+Key hyperparameters:
+- `r` (rank): 8-128. Bigger = more capacity + memory. 64 is a strong default.
+- `lora_alpha`: typically 2 x r (some teams use 1 x r for stability).
+- `lora_dropout`: 0.05-0.1; raise to 0.1 for small datasets.
+- `target_modules=None` auto-picks attention modules; pass
+  `["q_proj", "k_proj", "v_proj", "o_proj"]` (attention) or
+  `["gate_proj", "up_proj", "down_proj"]` (MLP) for explicit control.
+- `modules_to_save=["pooler"]` for CLS-pooled bases — the pooler Dense should
+  be trained too, not adapted.
+- LR is HIGHER than full fine-tune: 1e-4 to 5e-4 for LoRA vs. 2e-5 full.
+
+Rough memory savings on a 0.6B base (bf16, batch 64, seq 128): full fine-tune
+~24 GB, LoRA r=64 ~10 GB (~12M trainable, 2%), LoRA r=16 ~8 GB (~3M, 0.5%).
+Bigger savings on 7B+ models.
+
+Saving / sharing: `model.save_pretrained("dir")` writes ONLY the adapter (few
+MB) plus a reference to the base model. Loaders call the same one-liner;
+`peft` is invoked and the base downloaded on demand. For a merged model that
+loads without `peft` (needed for vLLM-style servers), call
+`model.transformers_model.merge_and_unload()` then `save_pretrained` /
+`push_to_hub`.
+
+Swapping adapters at inference (the main multi-task deployment win):
+    model = SentenceTransformer("base-model")
+    model.load_adapter("adapter-a", adapter_name="a")
+    model.load_adapter("adapter-b", adapter_name="b")
+    model.set_adapter("a"); emb_a = model.encode([...])
+
+QLoRA (4-bit base + LoRA) for 7B+ on consumer GPUs:
+    from transformers import BitsAndBytesConfig
+    bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
+                             bnb_4bit_compute_dtype=torch.bfloat16)
+    model = SentenceTransformer("Qwen/Qwen3-Embedding-7B",
+                                 model_kwargs={"quantization_config": bnb})
+    model.add_adapter(LoraConfig(r=64, lora_alpha=128, ...))
+`pip install bitsandbytes` first. Linux-only for bitsandbytes (Windows: the
+fork or WSL).
+
+Known issues:
+- LoRA PEFT on Qwen2.5-VL / paligemma / gemma3 / internvl / aya_vision under
+  transformers v5: `AutoModel.from_pretrained(peft_path)` crashes with
+  `KeyError: 'qwen2_vl'`. Pin transformers to 4.x or wait for the upstream fix.
+- `gradient_checkpointing=True` + LoRA: usually works; if you hit "None of
+  the inputs have requires_grad=True", call
+  `model.transformers_model.enable_input_require_grads()` after `add_adapter`.
+- `add_adapter` before pooling: when building from scratch (not loading a
+  pre-assembled checkpoint), call `add_adapter` AFTER
+  `SentenceTransformer(modules=[...])` is complete.
+
+Common gotchas: LR still at 2e-5 (LoRA needs higher); forgetting to merge for
+vLLM-style servers (they don't load `peft`); `r=8` too small for retrievers
+trained on millions of pairs (try 32 or 64); `modules_to_save` missing the
+pooler on CLS-pooled bases.
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import os
+from contextlib import nullcontext
+
+import torch
+from datasets import load_dataset
+from peft import LoraConfig, TaskType
+
+from sentence_transformers import (
+    SentenceTransformer,
+    SentenceTransformerModelCardData,
+    SentenceTransformerTrainer,
+    SentenceTransformerTrainingArguments,
+)
+from sentence_transformers.base.sampler import BatchSamplers
+from sentence_transformers.sentence_transformer.evaluation import NanoBEIREvaluator
+from sentence_transformers.sentence_transformer.losses import CachedMultipleNegativesRankingLoss
+
+
+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
+
+
+MODEL_NAME = "google-bert/bert-base-uncased"
+OUTPUT_DIR = "models/bert-base-gooaq-lora"
+RUN_NAME = "bert-base-gooaq-lora"
+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)
+        evaluator = NanoBEIREvaluator()
+        with autocast_ctx():
+            evaluator(model)
+        return
+
+    model = SentenceTransformer(
+        MODEL_NAME,
+        model_card_data=SentenceTransformerModelCardData(
+            language="en",
+            license="apache-2.0",
+            model_name=f"{MODEL_NAME.split('/')[-1]} LoRA adapter on GooAQ",
+        ),
+    )
+
+    peft_config = LoraConfig(
+        task_type=TaskType.FEATURE_EXTRACTION,
+        inference_mode=False,
+        r=64,
+        lora_alpha=128,
+        lora_dropout=0.1,
+    )
+    model.add_adapter(peft_config)
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    total = sum(p.numel() for p in model.parameters())
+    logging.info(f"trainable params: {trainable:,} / {total:,} ({100 * trainable / total:.2f}%)")
+
+    full = load_dataset("sentence-transformers/gooaq", split="train")
+    if SMOKE_TEST:
+        logging.info("SMOKE_TEST=1: trimmed dataset; will run max_steps=1 and skip Hub push")
+        full = full.select(range(min(200, len(full))))
+    eval_size = 20 if SMOKE_TEST else 10_000
+    train_cap = 50 if SMOKE_TEST else 500_000
+    split = full.train_test_split(test_size=eval_size, seed=12)
+    train_dataset = split["train"].select(range(min(train_cap, len(split["train"]))))
+    eval_dataset = split["test"]
+    logging.info(f"train={len(train_dataset):,} eval={len(eval_dataset):,}")
+
+    loss = CachedMultipleNegativesRankingLoss(model, mini_batch_size=32)
+
+    evaluator = NanoBEIREvaluator()
+    logging.info("Baseline:")
+    with autocast_ctx():
+        # Must run before deriving metric_key: evaluator(model) mutates primary_metric to add the name_ prefix.
+        baseline_eval = evaluator(model)[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=512,
+        per_device_eval_batch_size=512,
+        learning_rate=1e-4,
+        weight_decay=0.01,
+        warmup_steps=0.1,
+        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 = SentenceTransformerTrainer(
+        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:")
+    with autocast_ctx():
+        score = evaluator(model)[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}")
+
+    model.save_pretrained(f"{OUTPUT_DIR}/final")
+
+    if SMOKE_TEST:
+        logging.info("SMOKE_TEST=1: skipping Hub push")
+        return
+
+    try:
+        commit_url = model.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()

antigravity-awesome-skills/skills/train-sentence-transformers/scripts/train_sparse_encoder_distillation_example.py

@@ -0,0 +1,262 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "sentence-transformers[train]>=5.0",
+#     "datasets>=2.19.0",
+#     "accelerate>=0.26.0",
+#     "trackio",
+# ]
+# ///
+"""SPLADE distillation from a cross-encoder teacher.
+
+Trains a SPLADE sparse retriever to match a stronger cross-encoder's score
+gaps via `SparseMarginMSELoss` wrapped in `SpladeLoss` (the FLOPS regularizer
+is non-negotiable; without it embeddings collapse to dense).
+
+Data shape: `(query, positive, negative, score_diff)` where
+`score_diff = teacher(q, pos) - teacher(q, neg)`. This script uses
+`sentence-transformers/msmarco` (`bert-ensemble-margin-mse` subset) which has
+precomputed teacher score diffs. To distill from your own cross-encoder
+teacher, run a one-time teacher pass over your (q, pos, neg) triples and
+store the per-row score diff.
+
+Why distill SPLADE from a cross-encoder: SPLADE alone is hard to train from
+contrastive labels because the FLOPS regularizer fights early-training signal;
+distilling from a strong cross-encoder gives the model a dense regression
+target and reaches stronger nDCG faster than MNRL-only.
+
+Run locally:
+    pip install "sentence-transformers[train]>=5.0"
+    python train_sparse_encoder_distillation_example.py
+
+Multi-GPU:
+    accelerate launch train_sparse_encoder_distillation_example.py
+
+Hugging Face Jobs: paste this file's contents as the `script` in hf_jobs(...).
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import os
+from contextlib import nullcontext
+
+import torch
+from datasets import load_dataset, load_from_disk
+
+from sentence_transformers import (
+    SparseEncoder,
+    SparseEncoderModelCardData,
+    SparseEncoderTrainer,
+    SparseEncoderTrainingArguments,
+)
+from sentence_transformers.sparse_encoder.evaluation import SparseNanoBEIREvaluator
+from sentence_transformers.sparse_encoder.losses import SparseMarginMSELoss, SpladeLoss
+
+
+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
+
+
+MODEL_NAME = "Luyu/co-condenser-marco"  # MS MARCO-tuned MLM base; very strong starting point
+DATASET_NAME = "sentence-transformers/msmarco"
+DATASET_SUBSET = "bert-ensemble-margin-mse"
+TRAIN_SIZE = 100_000
+EVAL_SIZE = 5_000
+OUTPUT_DIR = "models/splade-msmarco-distilled"
+RUN_NAME = "splade-msmarco-distilled"
+DATA_CACHE = f"data/{RUN_NAME}-resolved"
+
+QUERY_REGULARIZER_WEIGHT = 0.1  # higher than contrastive recipe; distillation tolerates more sparsity pressure
+DOCUMENT_REGULARIZER_WEIGHT = 0.08
+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")
+
+
+def load_resolved_dataset():
+    """Load (query, positive, negative, score) rows. The MSMARCO subset is keyed by
+    passage_id / query_id; resolve to text once and cache to disk so reruns skip the work."""
+    if os.path.isdir(DATA_CACHE):
+        logging.info(f"Loading cached resolved dataset from {DATA_CACHE}")
+        return load_from_disk(DATA_CACHE)
+
+    logging.info(f"Resolving {DATASET_NAME}/{DATASET_SUBSET} ids -> text (one-time, cached)")
+    corpus_ds = load_dataset(DATASET_NAME, "corpus", split="train")
+    corpus = dict(zip(corpus_ds["passage_id"], corpus_ds["passage"]))
+    queries_ds = load_dataset(DATASET_NAME, "queries", split="train")
+    queries = dict(zip(queries_ds["query_id"], queries_ds["query"]))
+    raw = load_dataset(DATASET_NAME, DATASET_SUBSET, split="train").select(range(TRAIN_SIZE + EVAL_SIZE))
+
+    def id_to_text(batch):
+        return {
+            "query": [queries[qid] for qid in batch["query_id"]],
+            "positive": [corpus[pid] for pid in batch["positive_id"]],
+            "negative": [corpus[pid] for pid in batch["negative_id"]],
+            "score": batch["score"],
+        }
+
+    resolved = raw.map(id_to_text, batched=True, remove_columns=["query_id", "positive_id", "negative_id"])
+    resolved.save_to_disk(DATA_CACHE)
+    return resolved
+
+
+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 = SparseEncoder(cli.eval_only)
+        evaluator = SparseNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"])
+        with autocast_ctx():
+            evaluator(model)
+        return
+
+    logging.info(f"Loading base model: {MODEL_NAME}")
+    model = SparseEncoder(
+        MODEL_NAME,
+        model_card_data=SparseEncoderModelCardData(
+            language="en",
+            license="apache-2.0",
+            model_name=f"SPLADE from {MODEL_NAME.split('/')[-1]} distilled from MS MARCO ensemble",
+        ),
+    )
+    model.max_seq_length = 256
+
+    resolved = load_resolved_dataset()
+    if SMOKE_TEST:
+        logging.info("SMOKE_TEST=1: trimmed dataset; will run max_steps=1 and skip Hub push")
+        resolved = resolved.select(range(min(70, len(resolved))))
+    eval_size = 20 if SMOKE_TEST else EVAL_SIZE
+    split = resolved.train_test_split(test_size=eval_size, seed=12)
+    train_dataset = split["train"]
+    eval_dataset = split["test"]
+    logging.info(f"  train: {len(train_dataset):,} rows | eval: {len(eval_dataset):,} rows")
+    logging.info(f"  columns: {train_dataset.column_names}")
+
+    loss = SpladeLoss(
+        model=model,
+        loss=SparseMarginMSELoss(model=model),
+        query_regularizer_weight=QUERY_REGULARIZER_WEIGHT,
+        document_regularizer_weight=DOCUMENT_REGULARIZER_WEIGHT,
+    )
+
+    evaluator = SparseNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"])
+    logging.info("Baseline evaluation (fill-mask base scores near zero, confirms pipeline):")
+    with autocast_ctx():
+        # Must run before deriving metric_key: evaluator(model) mutates primary_metric to add the name_ prefix.
+        baseline_result = evaluator(model)
+        baseline_eval = baseline_result[evaluator.primary_metric]
+    metric_key = f"eval_{evaluator.primary_metric}"
+
+    args = SparseEncoderTrainingArguments(
+        output_dir=OUTPUT_DIR,
+        num_train_epochs=1,
+        max_steps=1 if SMOKE_TEST else -1,
+        per_device_train_batch_size=16,
+        per_device_eval_batch_size=16,
+        learning_rate=2e-5,
+        weight_decay=0.01,
+        warmup_steps=0.1,
+        lr_scheduler_type="linear",
+        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 = SparseEncoderTrainer(
+        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:")
+    with autocast_ctx():
+        result = evaluator(model)
+        score = result[evaluator.primary_metric]
+    delta = score - baseline_eval
+    verdict = "WIN" if delta >= 0.005 else "MARGINAL" if delta >= 0 else "REGRESSION"
+    # Active-dim keys come back name-prefixed (e.g. "NanoBEIR_..._query_active_dims"); suffix-match for compat.
+    qad = next((v for k, v in result.items() if k.endswith("query_active_dims")), "n/a")
+    cad = next((v for k, v in result.items() if k.endswith("corpus_active_dims")), "n/a")
+    logging.info(
+        f"VERDICT: {verdict} | score={score:.4f} | baseline={baseline_eval:.4f} | delta={delta:+.4f} "
+        f"| query_active={qad} corpus_active={cad}"
+    )
+
+    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)
+        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()

antigravity-awesome-skills/skills/train-sentence-transformers/scripts/train_sparse_encoder_example.py

@@ -0,0 +1,237 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "sentence-transformers[train]>=5.0",
+#     "datasets>=2.19.0",
+#     "accelerate>=0.26.0",
+#     "trackio",
+# ]
+# ///
+"""Production-ready sparse-encoder (SPLADE) training template.
+
+Demonstrates:
+- SpladeLoss wrapping SparseMultipleNegativesRankingLoss
+- FLOPS regularization (`query_regularizer_weight` / `document_regularizer_weight`)
+- SparseNanoBEIREvaluator for sparse retrieval metrics
+- load_best_model_at_end on the retrieval metric
+
+Base model must expose a masked-LM head; any `AutoModelForMaskedLM`-compatible
+checkpoint works (DistilBERT, BERT, MiniLM MLM variants, existing SPLADE models).
+
+Run locally:
+    pip install "sentence-transformers[train]>=5.0"
+    python train_sparse_encoder_example.py
+
+Multi-GPU:
+    accelerate launch train_sparse_encoder_example.py
+
+Hugging Face Jobs: paste this file's contents as the `script` in hf_jobs(...).
+"""
+
+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 (
+    SparseEncoder,
+    SparseEncoderModelCardData,
+    SparseEncoderTrainer,
+    SparseEncoderTrainingArguments,
+)
+from sentence_transformers.base.sampler import BatchSamplers
+from sentence_transformers.sparse_encoder.evaluation import SparseNanoBEIREvaluator
+from sentence_transformers.sparse_encoder.losses import (
+    SparseMultipleNegativesRankingLoss,
+    SpladeLoss,
+)
+
+
+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
+
+
+MODEL_NAME = "distilbert/distilbert-base-uncased"
+DATASET_NAME = "sentence-transformers/gooaq"
+TRAIN_SIZE = 100_000
+EVAL_SIZE = 1_000
+OUTPUT_DIR = "models/distilbert-splade-gooaq"
+RUN_NAME = "distilbert-splade-gooaq"
+
+QUERY_REGULARIZER_WEIGHT = 5e-5
+DOCUMENT_REGULARIZER_WEIGHT = 3e-5
+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 = SparseEncoder(cli.eval_only)
+        evaluator = SparseNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"])
+        with autocast_ctx():
+            evaluator(model)
+        return
+
+    logging.info(f"Loading base model: {MODEL_NAME}")
+    # Prompts are optional for SPLADE: most BERT-style MLM bases don't need them.
+    # If you're starting from a CSR (`Transformer + Pooling + SparseAutoEncoder`)
+    # base like `tomaarsen/csr-mxbai-embed-large-v1-nq` that *was* trained with
+    # prompts, mirror them here to preserve quality:
+    #   prompts={"query": "Represent this sentence for similarity: ", "document": ""},
+    #   default_prompt_name="document",
+    model = SparseEncoder(
+        MODEL_NAME,
+        model_card_data=SparseEncoderModelCardData(
+            language="en",
+            license="apache-2.0",
+            model_name=f"SPLADE from {MODEL_NAME.split('/')[-1]} trained on GooAQ",
+        ),
+    )
+
+    logging.info(f"Loading dataset: {DATASET_NAME}")
+    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")
+    full = load_dataset(DATASET_NAME, split="train")
+    split = full.train_test_split(test_size=eval_size, seed=12)
+    train_dataset = split["train"].select(range(min(train_size, len(split["train"]))))
+    eval_dataset = split["test"]
+    logging.info(f"  train: {len(train_dataset):,} rows | eval: {len(eval_dataset):,} rows")
+    logging.info(f"  columns: {train_dataset.column_names}")
+
+    loss = SpladeLoss(
+        model=model,
+        loss=SparseMultipleNegativesRankingLoss(model=model),
+        query_regularizer_weight=QUERY_REGULARIZER_WEIGHT,
+        document_regularizer_weight=DOCUMENT_REGULARIZER_WEIGHT,
+    )
+
+    evaluator = SparseNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"])
+    logging.info("Baseline evaluation:")
+    with autocast_ctx():
+        # Must run before deriving metric_key: evaluator(model) mutates primary_metric to add the name_ prefix.
+        baseline_result = evaluator(model)
+        baseline_eval = baseline_result[evaluator.primary_metric]
+    metric_key = f"eval_{evaluator.primary_metric}"
+
+    args = SparseEncoderTrainingArguments(
+        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=2e-5,
+        weight_decay=0.01,
+        warmup_steps=0.1,
+        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 = SparseEncoderTrainer(
+        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:")
+    with autocast_ctx():
+        result = evaluator(model)
+        score = result[evaluator.primary_metric]
+    delta = score - baseline_eval
+    verdict = "WIN" if delta >= 0.005 else "MARGINAL" if delta >= 0 else "REGRESSION"
+    # Active-dim keys come back name-prefixed (e.g. "NanoBEIR_..._query_active_dims"); suffix-match for compat.
+    qad = next((v for k, v in result.items() if k.endswith("query_active_dims")), "n/a")
+    cad = next((v for k, v in result.items() if k.endswith("corpus_active_dims")), "n/a")
+    logging.info(
+        f"VERDICT: {verdict} | score={score:.4f} | baseline={baseline_eval:.4f} | delta={delta:+.4f} "
+        f"| query_active={qad} corpus_active={cad}"
+    )
+
+    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)
+        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()