Upload train.py with huggingface_hub

train.py

@@ -110,9 +110,9 @@ BENCHMARK = [
 BASE_MODEL = "nomic-ai/nomic-embed-text-v1.5"
 EPOCHS = 30
 BATCH_SIZE = 64
-LEARNING_RATE = 5e-5
+LEARNING_RATE = 1e-4
 WARMUP_RATIO = 0.1
-MATRYOSHKA_DIMS = [768, 384, 256, 128]
+MATRYOSHKA_DIMS = []  # Disabled — dilutes gradient across dims
 GARBAGE_THRESHOLD = 0.3
 
 # -- Training Data --
@@ -745,25 +745,36 @@ def train():
     if baseline["failures"]:
         print(f"  Failures ({len(baseline['failures'])}): {baseline['failures']}")
 
-    # Build training objectives
-    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]
+    # 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]
 
-    train_objectives = []
+    all_triplet_examples = []
+
+    # Hard negatives first (highest value)
+    for t in triplets:
+        all_triplet_examples.append(InputExample(texts=[t["query"], t["positive"], t["negative"]]))
 
-    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))
+    # 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)}")
+
+    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))
 
     total_steps = sum(len(loader) for loader, _ in train_objectives) * EPOCHS
     warmup_steps = int(total_steps * WARMUP_RATIO)