Update README.md

README.md

@@ -87,6 +87,12 @@ Please cite the [Model2Vec repository](https://github.com/MinishLab/model2vec) i
 
 ## Reproducibility
 
+The following script can be used to reproduce this model. All credits go to [Tom Aarsen](https://huggingface.co/tomaarsen) for this fine-tuning approach and code. We make a few modifcations to the original code, namely:
+
+- We start with a pre-trained Model2Vec model ([potion-base-32M](https://huggingface.co/minishlab/potion-base-32M)).
+- We reduce the dataset size by a factor of 10. During experiments we saw that we didn't need the full dataset for the model to converge.
+- We decease the learning rate and train for 3 epochs instead of 1. Using a high learning rate wipes the effects of using a pre-trained model.
+
 ```python
 import random
 import logging
@@ -102,29 +108,29 @@ from sentence_transformers.training_args import BatchSamplers, MultiDatasetBatch
 from sentence_transformers.evaluation import NanoBEIREvaluator
 from sentence_transformers.models.StaticEmbedding import StaticEmbedding
 import wandb
-from transformers import AutoTokenizer
-from sentence_transformers.losses import CachedMultipleNegativesRankingLoss
-from torch.optim import AdamW
-from transformers import get_linear_schedule_with_warmup
-from torch.optim.lr_scheduler import CosineAnnealingLR
 
 logging.basicConfig(
     format="%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO
 )
 random.seed(12)
 
-            
-def load_train_eval_datasets():
+
+def load_train_eval_datasets(factor: int = 1):
     """
-    Either load the train and eval datasets from disk or load them from the datasets library & save them to disk.
+    Loads train and eval datasets from disk if available. Otherwise, downloads 
+    them from Hugging Face, preprocesses, and saves them to disk. If `factor` is 
+    greater than 1, returns a fraction (1/factor) of each dataset subset.
 
-    Upon saving to disk, we quit() to ensure that the datasets are not loaded into memory before training.
+    :param factor: The factor by which the data is reduced. If factor=1, no reduction is performed.
+    :return: (train_dataset: DatasetDict, eval_dataset: DatasetDict)
     """
     try:
+        # Try loading from disk
         train_dataset = DatasetDict.load_from_disk("datasets/train_dataset")
         eval_dataset = DatasetDict.load_from_disk("datasets/eval_dataset")
-        return train_dataset, eval_dataset
     except FileNotFoundError:
+        print("Prebuilt datasets not found on disk. Building from scratch...")
+
         print("Loading gooaq dataset...")
         gooaq_dataset = load_dataset("sentence-transformers/gooaq", split="train")
         gooaq_dataset_dict = gooaq_dataset.train_test_split(test_size=10_000, seed=12)
@@ -133,7 +139,11 @@ def load_train_eval_datasets():
         print("Loaded gooaq dataset.")
 
         print("Loading msmarco dataset...")
-        msmarco_dataset = load_dataset("sentence-transformers/msmarco-co-condenser-margin-mse-sym-mnrl-mean-v1", "triplet", split="train")
+        msmarco_dataset = load_dataset(
+            "sentence-transformers/msmarco-co-condenser-margin-mse-sym-mnrl-mean-v1",
+            "triplet",
+            split="train"
+        )
         msmarco_dataset_dict = msmarco_dataset.train_test_split(test_size=10_000, seed=12)
         msmarco_train_dataset: Dataset = msmarco_dataset_dict["train"]
         msmarco_eval_dataset: Dataset = msmarco_dataset_dict["test"]
@@ -147,15 +157,27 @@ def load_train_eval_datasets():
         print("Loaded squad dataset.")
 
         print("Loading s2orc dataset...")
-        s2orc_dataset = load_dataset("sentence-transformers/s2orc", "title-abstract-pair", split="train[:100000]")
+        s2orc_dataset = load_dataset(
+            "sentence-transformers/s2orc", 
+            "title-abstract-pair", 
+            split="train[:100000]"  # limit to 100k
+        )
         s2orc_dataset_dict = s2orc_dataset.train_test_split(test_size=10_000, seed=12)
         s2orc_train_dataset: Dataset = s2orc_dataset_dict["train"]
         s2orc_eval_dataset: Dataset = s2orc_dataset_dict["test"]
         print("Loaded s2orc dataset.")
 
         print("Loading allnli dataset...")
-        allnli_train_dataset = load_dataset("sentence-transformers/all-nli", "triplet", split="train")
-        allnli_eval_dataset = load_dataset("sentence-transformers/all-nli", "triplet", split="dev")
+        allnli_train_dataset = load_dataset(
+            "sentence-transformers/all-nli", 
+            "triplet", 
+            split="train"
+        )
+        allnli_eval_dataset = load_dataset(
+            "sentence-transformers/all-nli", 
+            "triplet", 
+            split="dev"
+        )
         print("Loaded allnli dataset.")
 
         print("Loading paq dataset...")
@@ -174,21 +196,33 @@ def load_train_eval_datasets():
 
         print("Loading msmarco_10m dataset...")
         msmarco_10m_dataset = load_dataset("bclavie/msmarco-10m-triplets", split="train")
-        msmarco_10m_dataset_dict = msmarco_10m_dataset.train_test_split(test_size=10_000, seed=12)
+        msmarco_10m_dataset_dict = msmarco_10m_dataset.train_test_split(
+            test_size=10_000, seed=12
+        )
         msmarco_10m_train_dataset: Dataset = msmarco_10m_dataset_dict["train"]
         msmarco_10m_eval_dataset: Dataset = msmarco_10m_dataset_dict["test"]
         print("Loaded msmarco_10m dataset.")
 
         print("Loading swim_ir dataset...")
-        swim_ir_dataset = load_dataset("nthakur/swim-ir-monolingual", "en", split="train").select_columns(["query", "text"])
-        swim_ir_dataset_dict = swim_ir_dataset.train_test_split(test_size=10_000, seed=12)
+        swim_ir_dataset = load_dataset(
+            "nthakur/swim-ir-monolingual", 
+            "en", 
+            split="train"
+        ).select_columns(["query", "text"])
+        swim_ir_dataset_dict = swim_ir_dataset.train_test_split(
+            test_size=10_000, seed=12
+        )
         swim_ir_train_dataset: Dataset = swim_ir_dataset_dict["train"]
         swim_ir_eval_dataset: Dataset = swim_ir_dataset_dict["test"]
         print("Loaded swim_ir dataset.")
 
         # NOTE: 20 negatives
         print("Loading pubmedqa dataset...")
-        pubmedqa_dataset = load_dataset("sentence-transformers/pubmedqa", "triplet-20", split="train")
+        pubmedqa_dataset = load_dataset(
+            "sentence-transformers/pubmedqa", 
+            "triplet-20", 
+            split="train"
+        )
         pubmedqa_dataset_dict = pubmedqa_dataset.train_test_split(test_size=100, seed=12)
         pubmedqa_train_dataset: Dataset = pubmedqa_dataset_dict["train"]
         pubmedqa_eval_dataset: Dataset = pubmedqa_dataset_dict["test"]
@@ -196,7 +230,11 @@ def load_train_eval_datasets():
 
         # NOTE: A lot of overlap with anchor/positives
         print("Loading miracl dataset...")
-        miracl_dataset = load_dataset("sentence-transformers/miracl", "en-triplet-all", split="train")
+        miracl_dataset = load_dataset(
+            "sentence-transformers/miracl", 
+            "en-triplet-all", 
+            split="train"
+        )
         miracl_dataset_dict = miracl_dataset.train_test_split(test_size=10_000, seed=12)
         miracl_train_dataset: Dataset = miracl_dataset_dict["train"]
         miracl_eval_dataset: Dataset = miracl_dataset_dict["test"]
@@ -204,7 +242,11 @@ def load_train_eval_datasets():
 
         # NOTE: A lot of overlap with anchor/positives
         print("Loading mldr dataset...")
-        mldr_dataset = load_dataset("sentence-transformers/mldr", "en-triplet-all", split="train")
+        mldr_dataset = load_dataset(
+            "sentence-transformers/mldr", 
+            "en-triplet-all", 
+            split="train"
+        )
         mldr_dataset_dict = mldr_dataset.train_test_split(test_size=10_000, seed=12)
         mldr_train_dataset: Dataset = mldr_dataset_dict["train"]
         mldr_eval_dataset: Dataset = mldr_dataset_dict["test"]
@@ -212,7 +254,11 @@ def load_train_eval_datasets():
 
         # NOTE: A lot of overlap with anchor/positives
         print("Loading mr_tydi dataset...")
-        mr_tydi_dataset = load_dataset("sentence-transformers/mr-tydi", "en-triplet-all", split="train")
+        mr_tydi_dataset = load_dataset(
+            "sentence-transformers/mr-tydi", 
+            "en-triplet-all", 
+            split="train"
+        )
         mr_tydi_dataset_dict = mr_tydi_dataset.train_test_split(test_size=10_000, seed=12)
         mr_tydi_train_dataset: Dataset = mr_tydi_dataset_dict["train"]
         mr_tydi_eval_dataset: Dataset = mr_tydi_dataset_dict["test"]
@@ -249,43 +295,35 @@ def load_train_eval_datasets():
             "mr_tydi": mr_tydi_eval_dataset,
         })
 
+        # Save to disk for next time
         train_dataset.save_to_disk("datasets/train_dataset")
         eval_dataset.save_to_disk("datasets/eval_dataset")
-        
-        # The `train_test_split` calls have put a lot of the datasets in memory, while we want it to just be on disk
+
+        # Quit to avoid memory overhead on large datasets
         quit()
-    
-
-    
-def load_train_eval_datasets_reduced():
-    # 1. Load the full datasets from disk
-    train_dataset = DatasetDict.load_from_disk("datasets/train_dataset")
-    eval_dataset = DatasetDict.load_from_disk("datasets/eval_dataset")
-    
-    factor = 10
-    for subset_name in train_dataset:
-        ds = train_dataset[subset_name]
-        ds = ds.shuffle(seed=42)  # Shuffle for a random subset
-        new_len = len(ds) // factor   # Keep 1/10th
-        ds = ds.select(range(new_len))
-        train_dataset[subset_name] = ds
-
-    for subset_name in eval_dataset:
-        ds = eval_dataset[subset_name]
-        ds = ds.shuffle(seed=42)
-        new_len = len(ds) // factor
-        ds = ds.select(range(new_len))
-        eval_dataset[subset_name] = ds
-    
+
+    # Reduce the dataset if factor > 1
+    if factor > 1:
+        for subset_name in train_dataset:
+            ds = train_dataset[subset_name].shuffle(seed=42)
+            new_len = len(ds) // factor
+            train_dataset[subset_name] = ds.select(range(new_len))
+
+        for subset_name in eval_dataset:
+            ds = eval_dataset[subset_name].shuffle(seed=42)
+            new_len = len(ds) // factor
+            eval_dataset[subset_name] = ds.select(range(new_len))
+
     return train_dataset, eval_dataset
 
+
 def main():
     wandb.init(entity="minishlab", project="minishlab")
-    # 1. Load a model to finetune with 2. (Optional) model card data
-    static_embedding = StaticEmbedding.from_model2vec("minishlab/potion-base-512dim-60kvocab")
-    # 2. Initialize the SentenceTransformer model as usual
+    # 1. Load a model to finetune
+    static_embedding = StaticEmbedding.from_model2vec("minishlab/potion-base-32M")
 
-    model_name = "potion-retrieval-512dim-60kvocab-v1"
+    # 2. Initialize the SentenceTransformer model
+    model_name = "potion-retrieval-32M"
     model = SentenceTransformer(
         modules=[static_embedding],
         model_card_data=SentenceTransformerModelCardData(
@@ -294,33 +332,31 @@ def main():
             model_name=model_name,
         ),
     )
-    # 3. Set up training & evaluation datasets - each dataset is trained with MNRL (with MRL)
-    train_dataset, eval_dataset = load_train_eval_datasets_reduced()
+
+    # 3. Load training & evaluation datasets
+    # NOTE: we reduce the total dataset size by a factor of 10 
+    train_dataset, eval_dataset = load_train_eval_datasets(factor=10)
     print(train_dataset)
 
     # 4. Define a loss function
     loss = MultipleNegativesRankingLoss(model)
     loss = MatryoshkaLoss(model, loss, matryoshka_dims=[32, 64, 128, 256, 512])
 
-
-    # 5. (Optional) Specify training arguments
+    # 5. Specify training arguments
     run_name = model_name
     epochs = 3
-    lr = 0.05 
+    lr = 0.05
     args = SentenceTransformerTrainingArguments(
-        # Required parameter:
         output_dir=f"models/{run_name}",
-        # Optional training parameters:
         num_train_epochs=epochs,
         per_device_train_batch_size=2048,
         per_device_eval_batch_size=2048,
         learning_rate=lr,
         warmup_ratio=0.1,
-        fp16=False,  # Set to False if you get an error that your GPU can't run on FP16
-        bf16=True,  # Set to True if you have a GPU that supports BF16
-        batch_sampler=BatchSamplers.NO_DUPLICATES,  # MultipleNegativesRankingLoss benefits from no duplicate samples in a batch
+        fp16=False,
+        bf16=True,
+        batch_sampler=BatchSamplers.NO_DUPLICATES,
         multi_dataset_batch_sampler=MultiDatasetBatchSamplers.PROPORTIONAL,
-        # Optional tracking/debugging parameters:
         eval_strategy="steps",
         eval_steps=250,
         save_strategy="steps",
@@ -328,17 +364,15 @@ def main():
         save_total_limit=2,
         logging_steps=250,
         logging_first_step=True,
-        run_name=run_name,  # Will be used in W&B if `wandb` is installed
+        run_name=run_name,
         report_to=["wandb"],
         load_best_model_at_end=True,
         metric_for_best_model="eval_NanoBEIR_mean_cosine_ndcg@10",
         greater_is_better=True,
-        #
     )
 
-    # 6. (Optional) Create an evaluator & evaluate the base model
+    # 6. Create an evaluator & evaluate the base model
     evaluator = NanoBEIREvaluator()
-    
     evaluator(model)
 
     # 7. Create a trainer & train
@@ -352,12 +386,11 @@ def main():
     )
     trainer.train()
 
-    # (Optional) Evaluate the trained model on the evaluator after training
+    # 8. Evaluate the trained model and save
     evaluator(model)
-
-    # 8. Save the trained model
     model.save_pretrained(f"models/{run_name}/final")
 
+
 if __name__ == "__main__":
     main()
 ```