Upload train_reranker.py with huggingface_hub

train_reranker.py

@@ -0,0 +1,127 @@
+# /// script
+# requires-python = ">=3.11"
+# dependencies = [
+#     "sentence-transformers[train]>=4.0",
+#     "datasets",
+#     "torch>=2.4",
+#     "transformers>=4.48",
+#     "trackio",
+# ]
+# ///
+"""
+Soft-Label Cross-Encoder Reranker Training
+
+Trains a reranker using continuous relevance scores (soft labels).
+Dataset format: {"query": "...", "text": "...", "score": 0.0-1.0}
+"""
+
+import logging
+import os
+from collections import defaultdict
+from datasets import load_dataset
+from sentence_transformers.cross_encoder import (
+    CrossEncoder,
+    CrossEncoderTrainer,
+    CrossEncoderTrainingArguments,
+)
+from sentence_transformers.cross_encoder.evaluation import CrossEncoderNanoBEIREvaluator
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Configuration
+DATASET_NAME = os.environ.get("DATASET_NAME", "amanwithaplan/arcade-reranker-data")
+HUB_MODEL_ID = os.environ.get("HUB_MODEL_ID", "amanwithaplan/arcade-reranker")
+BASE_MODEL = os.environ.get("BASE_MODEL", "Alibaba-NLP/gte-reranker-modernbert-base")
+NUM_EPOCHS = int(os.environ.get("NUM_EPOCHS", "5"))
+BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "16"))
+LEARNING_RATE = float(os.environ.get("LEARNING_RATE", "2e-5"))
+MAX_SEQ_LENGTH = int(os.environ.get("MAX_SEQ_LENGTH", "512"))
+RUN_NAME = os.environ.get("RUN_NAME", "reranker-03130903")
+
+
+def main():
+    logger.info(f"Configuration:")
+    logger.info(f"  Dataset: {DATASET_NAME}")
+    logger.info(f"  Base model: {BASE_MODEL}")
+    logger.info(f"  Epochs: {NUM_EPOCHS}")
+    logger.info(f"  Run name: {RUN_NAME}")
+
+    model = CrossEncoder(BASE_MODEL, max_length=MAX_SEQ_LENGTH)
+
+    logger.info(f"Loading dataset: {DATASET_NAME}")
+    dataset = load_dataset(DATASET_NAME, split="train")
+
+    # Log dataset composition
+    if "type" in dataset.column_names:
+        type_counts = defaultdict(int)
+        for item in dataset:
+            type_counts[item["type"]] += 1
+        logger.info(f"Dataset composition: {dict(type_counts)}")
+
+    logger.info(f"Total examples: {len(dataset)}")
+
+    # Rename columns for CrossEncoderTrainer
+    dataset = dataset.rename_columns({
+        "query": "sentence1",
+        "text": "sentence2",
+        "score": "label"
+    })
+
+    # Split for evaluation
+    eval_size = min(400, int(len(dataset) * 0.15))
+    splits = dataset.train_test_split(test_size=eval_size, seed=42)
+    train_dataset = splits["train"]
+    eval_dataset = splits["test"]
+
+    logger.info(f"Train: {len(train_dataset)}, Eval: {len(eval_dataset)}")
+
+    # NanoBEIR for benchmark comparison
+    evaluator = CrossEncoderNanoBEIREvaluator(
+        dataset_names=["msmarco", "nfcorpus", "nq"],
+        batch_size=BATCH_SIZE,
+    )
+
+    args = CrossEncoderTrainingArguments(
+        output_dir="models/reranker",
+        num_train_epochs=NUM_EPOCHS,
+        per_device_train_batch_size=BATCH_SIZE,
+        per_device_eval_batch_size=BATCH_SIZE,
+        learning_rate=LEARNING_RATE,
+        warmup_ratio=0.1,
+        bf16=True,
+        eval_strategy="steps",
+        eval_steps=200,
+        save_strategy="steps",
+        save_steps=200,
+        save_total_limit=2,
+        logging_steps=25,
+        logging_first_step=True,
+        load_best_model_at_end=True,
+        metric_for_best_model="eval_loss",
+        greater_is_better=False,
+        push_to_hub=True,
+        hub_model_id=HUB_MODEL_ID,
+        hub_strategy="every_save",
+        report_to="trackio",
+        run_name=RUN_NAME,
+    )
+
+    trainer = CrossEncoderTrainer(
+        model=model,
+        args=args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        evaluator=evaluator,
+    )
+
+    logger.info("Starting training...")
+    trainer.train()
+
+    logger.info(f"Pushing final model to {HUB_MODEL_ID}")
+    model.push_to_hub(HUB_MODEL_ID)
+    logger.info("Done!")
+
+
+if __name__ == "__main__":
+    main()