opencode-skills-collection 3.1.12 → 3.1.13

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

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