Upload train.py with huggingface_hub

train.py

@@ -108,7 +108,7 @@ BENCHMARK = [
 
 # -- Hyperparameters --
 BASE_MODEL = "nomic-ai/nomic-embed-text-v1.5"
-EPOCHS = 30
+EPOCHS = 10
 BATCH_SIZE = 64
 LEARNING_RATE = 1e-4
 WARMUP_RATIO = 0.1
@@ -745,36 +745,25 @@ def train():
     if baseline["failures"]:
         print(f"  Failures ({len(baseline['failures'])}): {baseline['failures']}")
 
-    # Build training objectives — TripletLoss only
-    # Convert positive pairs into triplets by sampling random negatives from corpus
-    qp = "search_query: "
-    dp = "search_document: "
-    corpus_docs = [f"{dp}{name}" for name in corpus_names]
-
-    all_triplet_examples = []
-
-    # Hard negatives first (highest value)
-    for t in triplets:
-        all_triplet_examples.append(InputExample(texts=[t["query"], t["positive"], t["negative"]]))
-
-    # Convert positive pairs to triplets with random corpus negatives
-    for p in positive_pairs:
-        pos_subcat = strip_prefix(p["positive"])
-        # Pick a random negative that isn't the positive
-        neg_candidates = [d for d in corpus_docs if strip_prefix(d) != pos_subcat]
-        if neg_candidates:
-            neg = random.choice(neg_candidates)
-            all_triplet_examples.append(InputExample(texts=[p["query"], p["positive"], neg]))
-
-    random.shuffle(all_triplet_examples)
-    print(f"  Total triplets (hard + random neg): {len(all_triplet_examples)}")
+    # Build training objectives — MNRL (no Matryoshka) + hard negative TripletLoss
+    pos_examples = [InputExample(texts=[p["query"], p["positive"]]) for p in positive_pairs]
+    trip_examples = [InputExample(texts=[t["query"], t["positive"], t["negative"]]) for t in triplets]
 
     train_objectives = []
-    trip_loader = DataLoader(all_triplet_examples, shuffle=True, batch_size=BATCH_SIZE)
-    triplet_loss = losses.TripletLoss(model)
-    if MATRYOSHKA_DIMS:
-        triplet_loss = losses.MatryoshkaLoss(model, triplet_loss, matryoshka_dims=MATRYOSHKA_DIMS)
-    train_objectives.append((trip_loader, triplet_loss))
+
+    if pos_examples:
+        pos_loader = DataLoader(pos_examples, shuffle=True, batch_size=BATCH_SIZE)
+        mnrl_loss = losses.MultipleNegativesRankingLoss(model)
+        if MATRYOSHKA_DIMS:
+            mnrl_loss = losses.MatryoshkaLoss(model, mnrl_loss, matryoshka_dims=MATRYOSHKA_DIMS)
+        train_objectives.append((pos_loader, mnrl_loss))
+
+    if trip_examples:
+        trip_loader = DataLoader(trip_examples, shuffle=True, batch_size=BATCH_SIZE)
+        triplet_loss = losses.TripletLoss(model)
+        if MATRYOSHKA_DIMS:
+            triplet_loss = losses.MatryoshkaLoss(model, triplet_loss, matryoshka_dims=MATRYOSHKA_DIMS)
+        train_objectives.append((trip_loader, triplet_loss))
 
     total_steps = sum(len(loader) for loader, _ in train_objectives) * EPOCHS
     warmup_steps = int(total_steps * WARMUP_RATIO)