tribunal-kit 9.2.3 → 9.2.4

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

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