Clean up model naming and remove training comparison

README.md

@@ -10,7 +10,7 @@ tags:
 - generated_from_trainer
 - dataset_size:501907
 - loss:MultiVectorMultipleNegativesRankingLoss
-base_model: multi-vector-encoder-testing/bert-tiny-msmarco
+base_model: prajjwal1/bert-tiny
 widget:
 - text: 'Kroger Pharmacy - Keller 976 Keller Pkwy, Keller TX 76248 Phone Number: (817)
     431-5178'
@@ -954,16 +954,15 @@ model-index:
       name: Maxsim Map@100
 ---
 
-See [the training comparison](RUN_SUMMARY.md) for the longer run settings and all before-and-after results.
 
 # BERT tiny multi-vector encoder trained on MS MARCO
 
-This is a [Multi-Vector Encoder](https://www.sbert.net/docs/multi_vector_encoder/usage/usage.html) model finetuned from [multi-vector-encoder-testing/bert-tiny-msmarco](https://huggingface.co/multi-vector-encoder-testing/bert-tiny-msmarco) on the [msmarco-bm25](https://huggingface.co/datasets/sentence-transformers/msmarco-bm25) dataset using the [sentence-transformers](https://www.SBERT.net) library. It maps inputs to sequences of 128-dimensional token-level vectors and scores them with late interaction (MaxSim), useful for semantic search with late interaction.
+This is a [Multi-Vector Encoder](https://www.sbert.net/docs/multi_vector_encoder/usage/usage.html) model finetuned in two stages from [prajjwal1/bert-tiny](https://huggingface.co/prajjwal1/bert-tiny) on the [msmarco-bm25](https://huggingface.co/datasets/sentence-transformers/msmarco-bm25) dataset using the [sentence-transformers](https://www.SBERT.net) library. It maps inputs to sequences of 128-dimensional token-level vectors and scores them with late interaction (MaxSim), useful for semantic search with late interaction.
 ## Model Details
 
 ### Model Description
 - **Model Type:** Multi-Vector Encoder
-- **Base model:** [multi-vector-encoder-testing/bert-tiny-msmarco](https://huggingface.co/multi-vector-encoder-testing/bert-tiny-msmarco) <!-- at revision 81c5b4e78ac3bdbb01606e60e82bc34d86ed897b -->
+- **Base model:** [prajjwal1/bert-tiny](https://huggingface.co/prajjwal1/bert-tiny)
 - **Maximum Sequence Length:** 512 tokens
     - **Maximum Query Length:** 32 tokens
     - **Maximum Document Length:** 256 tokens
@@ -1007,7 +1006,7 @@ Then you can load this model and run inference.
 from sentence_transformers import MultiVectorEncoder
 
 # Download from the 🤗 Hub
-model = MultiVectorEncoder("multi-vector-encoder-testing/bert-tiny-msmarco-long")
+model = MultiVectorEncoder("multi-vector-encoder-testing/bert-tiny-multi-vector")
 # Run inference: each input becomes a sequence of per-token vectors (variable length).
 queries = [
     'calories in kirkland ravioli',

RUN_SUMMARY.md

@@ -1,45 +0,0 @@
-# Longer BERT tiny training run
-
-Continued multi-vector-encoder-testing/bert-tiny-msmarco at revision 81c5b4e78ac3bdbb01606e60e82bc34d86ed897b.
-
-Trained for 10,000 additional steps with batch size 128 on 501,907 MS MARCO BM25 triplets, with 1,024 held-out rows for evaluation loss. This processed 1.28 million triplets over about 2.55 epochs. Learning rate was 1e-5, with 5% warmup and linear decay. Training used bf16 on one RTX 3090 and took 15.6 minutes including periodic evaluation.
-
-Selected step 6,000 using mean nDCG@10 on NanoMSMARCO, NanoNQ, and NanoFiQA2018, evaluated every 2,000 steps. The remaining ten datasets were evaluated only before and after training.
-
-## Mean nDCG@10
-
-| Evaluation group | Initial model | Longer run |
-| --- | ---: | ---: |
-| Three selection datasets | 0.3326 | 0.3729 |
-| All 13 datasets | 0.3999 | 0.4468 |
-| Ten additional datasets | 0.4201 | 0.4690 |
-
-## Per-dataset nDCG@10
-
-| Dataset | Initial model | Longer run | Change |
-| --- | ---: | ---: | ---: |
-| NanoClimateFEVER | 0.1431 | 0.1919 | +0.0489 |
-| NanoDBPedia | 0.3914 | 0.4820 | +0.0907 |
-| NanoFEVER | 0.6310 | 0.6793 | +0.0483 |
-| NanoFiQA2018 | 0.2658 | 0.3027 | +0.0369 |
-| NanoHotpotQA | 0.5766 | 0.6840 | +0.1073 |
-| NanoMSMARCO | 0.4328 | 0.3859 | -0.0469 |
-| NanoNFCorpus | 0.2569 | 0.2852 | +0.0283 |
-| NanoNQ | 0.2992 | 0.4300 | +0.1307 |
-| NanoQuoraRetrieval | 0.7976 | 0.8355 | +0.0379 |
-| NanoSCIDOCS | 0.2029 | 0.2217 | +0.0188 |
-| NanoArguAna | 0.2863 | 0.3539 | +0.0676 |
-| NanoSciFact | 0.5060 | 0.5650 | +0.0590 |
-| NanoTouche2020 | 0.4096 | 0.3920 | -0.0176 |
-
-Improved on 11 of 13 datasets. NanoMSMARCO and NanoTouche2020 declined.
-
-## Reproduce
-
-Run with a compatible Sentence Transformers checkout and its training dependencies:
-
-```bash
-python train.py --long-run
-```
-
-The script pins the initial model revision. Training arguments, every evaluation, and per-dataset before-and-after results are included in results.json.

train.py

@@ -4,6 +4,7 @@ Adapted from the multi-vector training skill template and training_contrastive.p
 Run from the repository root with Python and the training dependencies installed.
 Use --smoke-test for one step, or --push-to-hub to upload the best checkpoint.
 Use --long-run to continue the initial model for 10,000 steps on the full dataset.
+Run without --long-run first to create the local initial model.
 """
 
 import argparse
@@ -30,8 +31,7 @@ from sentence_transformers.multi_vector_encoder.losses import MultiVectorMultipl
 from sentence_transformers.multi_vector_encoder.modules import MultiVectorMask
 
 RUN_NAME = "bert-tiny-msmarco"
-REPO_ID = f"multi-vector-encoder-testing/{RUN_NAME}"
-INITIAL_REVISION = "81c5b4e78ac3bdbb01606e60e82bc34d86ed897b"
+REPO_ID = "multi-vector-encoder-testing/bert-tiny-multi-vector"
 
 
 class LogProgress(TrainerCallback):
@@ -58,8 +58,8 @@ def main():
     parser.add_argument("--push-to-hub", action="store_true")
     parser.add_argument("--long-run", action="store_true")
     cli = parser.parse_args()
-    repo_id = REPO_ID + ("-long" if cli.long_run else "")
-    run_name = repo_id.split("/")[-1] + ("-smoke" if cli.smoke_test else "")
+    repo_id = REPO_ID
+    run_name = RUN_NAME + ("-long" if cli.long_run else "") + ("-smoke" if cli.smoke_test else "")
     output_dir = Path("models") / run_name
     output_dir.mkdir(parents=True, exist_ok=True)
     Path("logs").mkdir(exist_ok=True)
@@ -82,7 +82,11 @@ def main():
         model_id=repo_id,
     )
     if cli.long_run:
-        model = MultiVectorEncoder(REPO_ID, revision=INITIAL_REVISION, model_card_data=card)
+        initial_model = Path("models") / RUN_NAME / "final"
+        if not initial_model.is_dir():
+            parser.error("Run without --long-run first to create the local initial model.")
+        model = MultiVectorEncoder(str(initial_model), model_card_data=card)
+        model.model_card_data.set_base_model("prajjwal1/bert-tiny")
     else:
         # The original checkpoint lacks model_type, which recent AutoConfig versions require.
         base_dir = output_dir / "base"