tribunal-kit 9.2.1 → 9.2.4

package/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())

package/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()