Update README.md

README.md

@@ -34,11 +34,11 @@ You can train a better model if you have access to adequate compute (can finetun
 from datasets import load_dataset
 from ragatouille import RAGTrainer
 sample_size = 100000
-ds = load_dataset('akhooli/arabic-triplets-1m-curated-sims-len', 'arabic', split="train", trust_remote_code=True, streaming=True)
+ds = load_dataset('akhooli/arabic-triplets-1m-curated-sims-len', split="train", trust_remote_code=True, streaming=True)
 
-# some data processing not in this script
+# some data processing not in this script (data filtered based on similarity scores) and 100K selected at random
 sds = ds.shuffle(seed=42, buffer_size=10_000)
-dsf = sds.take(sample_size)
+dsf = sds
 triplets = []
 for item in iter(dsf):
   triplets.append((item["query"], item["positive"], item["negative"]))
@@ -46,11 +46,11 @@ trainer = RAGTrainer(model_name="Arabic-ColBERT-100k", pretrained_model_name="au
 trainer.prepare_training_data(raw_data=triplets, mine_hard_negatives=False)
 
 trainer.train(batch_size=32,
-    nbits=2, # How many bits will the trained model use when compressing indexes
-    maxsteps=100000, # Maximum steps hard stop
+    nbits=4, # How many bits will the trained model use when compressing indexes
+    maxsteps=3125, # Maximum steps hard stop
     use_ib_negatives=True, # Use in-batch negative to calculate loss
     dim=128, # How many dimensions per embedding. 128 is the default and works well.
-    learning_rate=5e-6, # Learning rate, small values ([3e-6,3e-5] work best if the base model is BERT-like, 5e-6 is often the sweet spot)
+    learning_rate=1e-5, # Learning rate, small values ([3e-6,3e-5] work best if the base model is BERT-like, 5e-6 is often the sweet spot)
     doc_maxlen=256, # Maximum document length. Because of how ColBERT works, smaller chunks (128-256) work very well.
     use_relu=False, # Disable ReLU -- doesn't improve performance
     warmup_steps="auto", # Defaults to 10%