custom training loop — bypass SentenceTransformerTrainer

NVEmbedModel has non-standard forward() incompatible with ST.
Uses model.encode() with torch.enable_grad() for differentiable
embeddings + manual MNRL loss + AdamW with warmup scheduler.
No sentence-transformers dependency needed for training.

Made-with: Cursor

requirements.txt

@@ -1,5 +1,4 @@
 transformers==4.45.2
-sentence-transformers>=3.3.0,<3.4.0
 datasets>=2.14.0
 huggingface-hub>=0.20.0
 accelerate>=0.25.0

run.py

@@ -1,12 +1,13 @@
 """
-NV-Embed-v2: Evaluate base → Fine-tune with LoRA → Evaluate fine-tuned.
-All on same GPU in one run.
+NV-Embed-v2: Evaluate base → Fine-tune with LoRA (custom loop) → Evaluate fine-tuned.
+Custom training loop because NVEmbedModel.forward() is incompatible with SentenceTransformerTrainer.
 """
 
 import json
 import random
 import logging
 import os
+import gc
 import numpy as np
 import pandas as pd
 import faiss
@@ -42,21 +43,22 @@ LORA_R = 16
 LORA_ALPHA = 32
 LORA_DROPOUT = 0.1
 LORA_TARGETS = ["q_proj", "v_proj", "k_proj", "o_proj"]
+TEMPERATURE = 20.0
 
 
 # ═══════════════════════════════════════════════════════════════════════════
-# EVALUATION (uses AutoModel.encode directly — no sentence-transformers)
+# EVALUATION (uses AutoModel.encode — no gradients needed)
 # ═══════════════════════════════════════════════════════════════════════════
 
-def evaluate_with_automodel(model_name, token):
-    """Load model via AutoModel and run ESCO benchmark."""
+def evaluate_with_automodel(model_name_or_path, token, method_label=None):
     from transformers import AutoModel
 
-    logger.info(f"\n{'='*60}\n  EVALUATING: {model_name}\n{'='*60}")
+    label = method_label or model_name_or_path.split("/")[-1]
+    logger.info(f"\n{'='*60}\n  EVALUATING: {label}\n{'='*60}")
 
-    logger.info(f"Loading {model_name}...")
+    logger.info(f"Loading {model_name_or_path}...")
     model = AutoModel.from_pretrained(
-        model_name, trust_remote_code=True, token=token,
+        model_name_or_path, trust_remote_code=True, token=token,
         torch_dtype=torch.float16, device_map="auto",
     )
     model.eval()
@@ -111,7 +113,7 @@ def evaluate_with_automodel(model_name, token):
 
         met = {
             "test_set": test_name,
-            "method": model_name.split("/")[-1],
+            "method": label,
             "matchable_rows": matchable_count,
             "accuracy_top1": round(top1/matchable_count, 6),
             "accuracy_top3": round(top3/matchable_count, 6),
@@ -126,12 +128,13 @@ def evaluate_with_automodel(model_name, token):
         )
 
     del model
+    gc.collect()
     torch.cuda.empty_cache()
     return all_metrics
 
 
 # ═══════════════════════════════════════════════════════════════════════════
-# TRAINING (uses SentenceTransformer + LoRA)
+# TRAINING (custom loop — model.encode() with torch.enable_grad)
 # ═══════════════════════════════════════════════════════════════════════════
 
 def augment_with_context(sentences, prob=0.5):
@@ -144,98 +147,116 @@ def augment_with_context(sentences, prob=0.5):
     return augmented
 
 
+def mnrl_loss(anchor_emb, positive_emb, temperature=TEMPERATURE):
+    """Multiple Negatives Ranking Loss: in-batch contrastive."""
+    scores = torch.mm(anchor_emb, positive_emb.t()) * temperature
+    labels = torch.arange(scores.size(0), device=scores.device)
+    return F.cross_entropy(scores, labels)
+
+
 def train_model(token):
-    from datasets import load_dataset, Dataset
-    from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer
-    from sentence_transformers.losses import MultipleNegativesRankingLoss
-    from sentence_transformers.training_args import SentenceTransformerTrainingArguments
-    from sentence_transformers.evaluation import InformationRetrievalEvaluator
+    from transformers import AutoModel
     from peft import LoraConfig, get_peft_model
+    from datasets import load_dataset
 
-    logger.info(f"\n{'='*60}\n  FINE-TUNING WITH LoRA\n{'='*60}")
+    logger.info(f"\n{'='*60}\n  FINE-TUNING WITH LoRA (custom loop)\n{'='*60}")
 
-    lora_config = LoraConfig(
-        r=LORA_R, lora_alpha=LORA_ALPHA, target_modules=LORA_TARGETS,
-        lora_dropout=LORA_DROPOUT, bias="none", task_type="FEATURE_EXTRACTION",
+    # Load model
+    logger.info(f"Loading {MODEL_ID}...")
+    model = AutoModel.from_pretrained(
+        MODEL_ID, trust_remote_code=True, token=token,
+        torch_dtype=torch.bfloat16,
     )
 
-    logger.info(f"Loading {MODEL_ID} via SentenceTransformer...")
-    model = SentenceTransformer(
-        MODEL_ID, trust_remote_code=True,
-        model_kwargs={"torch_dtype": torch.bfloat16},
+    # Apply LoRA
+    lora_config = LoraConfig(
+        r=LORA_R, lora_alpha=LORA_ALPHA, target_modules=LORA_TARGETS,
+        lora_dropout=LORA_DROPOUT, bias="none",
     )
-    model.max_seq_length = 512
-
-    logger.info("Applying LoRA adapter...")
-    model[0].auto_model = get_peft_model(model[0].auto_model, lora_config)
-    model[0].auto_model.print_trainable_parameters()
-
-    model.prompts = {"anchor": QUERY_INSTRUCTION, "positive": ""}
+    model = get_peft_model(model, lora_config)
+    model.print_trainable_parameters()
+    model.cuda()
+    model.train()
 
     # Dataset
     logger.info(f"Loading dataset: {DATASET_ID}")
     raw = load_dataset(DATASET_ID, split="train").shuffle(seed=42)
-    split = raw.train_test_split(test_size=0.05, seed=42)
-
-    sentences = split["train"]["sentence"]
-    skills = split["train"]["skill"]
-    aug = augment_with_context(sentences, prob=AUGMENT_PROB)
-    train_dataset = Dataset.from_dict({"anchor": aug, "positive": skills})
-    eval_raw = split["test"]
-    logger.info(f"Train: {len(train_dataset)}, Eval: {len(eval_raw)}")
-
-    loss = MultipleNegativesRankingLoss(model)
-
-    sample = eval_raw.select(range(min(500, len(eval_raw))))
-    evaluator = InformationRetrievalEvaluator(
-        queries={str(i): row["sentence"] for i, row in enumerate(sample)},
-        corpus={s: s for s in set(sample["skill"])},
-        relevant_docs={str(i): {row["skill"]} for i, row in enumerate(sample)},
-        name="esco-eval",
-        score_functions={"cosine": lambda a, b: (a @ b.T)},
+    anchors_raw = raw["sentence"]
+    positives_raw = raw["skill"]
+    logger.info(f"Dataset: {len(anchors_raw)} pairs")
+
+    logger.info("Augmenting anchors...")
+    anchors = augment_with_context(anchors_raw, prob=AUGMENT_PROB)
+    positives = positives_raw
+
+    # Optimizer (only LoRA params)
+    trainable_params = [p for p in model.parameters() if p.requires_grad]
+    optimizer = torch.optim.AdamW(trainable_params, lr=LR, weight_decay=0.01)
+
+    total_micro_steps = len(anchors) // TRAIN_BATCH
+    total_optim_steps = total_micro_steps // GRAD_ACCUM
+    logger.info(f"Micro-steps: {total_micro_steps}, Optimizer steps: {total_optim_steps}")
+    logger.info(f"Warmup: {WARMUP_STEPS} steps, LR: {LR}")
+
+    # LR scheduler with warmup
+    from transformers import get_linear_schedule_with_warmup
+    scheduler = get_linear_schedule_with_warmup(
+        optimizer, num_warmup_steps=WARMUP_STEPS, num_training_steps=total_optim_steps
     )
 
-    push = bool(HUB_MODEL_ID)
-    args_kwargs = dict(
-        output_dir=OUTPUT_DIR,
-        num_train_epochs=EPOCHS,
-        per_device_train_batch_size=TRAIN_BATCH,
-        per_device_eval_batch_size=TRAIN_BATCH,
-        gradient_accumulation_steps=GRAD_ACCUM,
-        learning_rate=LR,
-        warmup_steps=WARMUP_STEPS,
-        bf16=True,
-        eval_strategy="steps", eval_steps=500,
-        save_strategy="steps", save_steps=500,
-        save_total_limit=2,
-        load_best_model_at_end=True,
-        metric_for_best_model="esco-eval_cosine_ndcg@10",
-        logging_steps=50,
-        gradient_checkpointing=False,
-        dataloader_pin_memory=False,
-        push_to_hub=push,
-    )
-    if push:
-        args_kwargs["hub_model_id"] = HUB_MODEL_ID
-        args_kwargs["hub_strategy"] = "every_save"
-
-    trainer = SentenceTransformerTrainer(
-        model=model,
-        args=SentenceTransformerTrainingArguments(**args_kwargs),
-        train_dataset=train_dataset,
-        eval_dataset=train_dataset.select(range(500)),
-        loss=loss, evaluator=evaluator,
-    )
+    # Training loop
+    indices = list(range(len(anchors)))
+    random.shuffle(indices)
+
+    optimizer.zero_grad()
+    running_loss = 0.0
+    micro_step = 0
+
+    for i in range(0, len(indices) - TRAIN_BATCH + 1, TRAIN_BATCH):
+        batch_idx = indices[i:i+TRAIN_BATCH]
+        batch_anchors = [anchors[j] for j in batch_idx]
+        batch_positives = [positives[j] for j in batch_idx]
+
+        # Get embeddings WITH gradients via torch.enable_grad()
+        with torch.enable_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16):
+            anchor_emb = model.encode(batch_anchors, instruction=QUERY_INSTRUCTION, max_length=512)
+            positive_emb = model.encode(batch_positives, instruction="", max_length=512)
+
+            anchor_emb = F.normalize(anchor_emb, p=2, dim=1)
+            positive_emb = F.normalize(positive_emb, p=2, dim=1)
+
+            loss = mnrl_loss(anchor_emb, positive_emb) / GRAD_ACCUM
+
+        loss.backward()
+        running_loss += loss.item()
+        micro_step += 1
+
+        if micro_step % GRAD_ACCUM == 0:
+            torch.nn.utils.clip_grad_norm_(trainable_params, 1.0)
+            optimizer.step()
+            scheduler.step()
+            optimizer.zero_grad()
+
+            optim_step = micro_step // GRAD_ACCUM
+            avg_loss = running_loss
+            running_loss = 0.0
+
+            if optim_step % 50 == 0:
+                lr_now = scheduler.get_last_lr()[0]
+                logger.info(f"  step {optim_step}/{total_optim_steps}  loss={avg_loss:.4f}  lr={lr_now:.2e}")
 
-    logger.info("Starting training...")
-    trainer.train()
+    # Save
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
+    logger.info(f"Saving to {OUTPUT_DIR}...")
+    model.save_pretrained(OUTPUT_DIR)
 
-    model.save_pretrained(f"{OUTPUT_DIR}/final")
-    if push:
+    if HUB_MODEL_ID:
         logger.info(f"Pushing to Hub: {HUB_MODEL_ID}")
-        model.push_to_hub(HUB_MODEL_ID, exist_ok=True)
+        model.push_to_hub(HUB_MODEL_ID, token=token)
+        logger.info("Pushed.")
 
-    del model, trainer
+    del model, optimizer, trainable_params
+    gc.collect()
     torch.cuda.empty_cache()
     logger.info("Training complete.")
 
@@ -256,14 +277,14 @@ def main():
         logger.info(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
 
     # Phase 1: Evaluate base model
-    base_metrics = evaluate_with_automodel(MODEL_ID, token)
+    base_metrics = evaluate_with_automodel(MODEL_ID, token, "nv-embed-v2")
 
     # Phase 2: Fine-tune with LoRA
     train_model(token)
 
     # Phase 3: Evaluate fine-tuned model
-    ft_model_id = HUB_MODEL_ID if HUB_MODEL_ID else f"{OUTPUT_DIR}/final"
-    ft_metrics = evaluate_with_automodel(ft_model_id, token)
+    ft_source = HUB_MODEL_ID if HUB_MODEL_ID else OUTPUT_DIR
+    ft_metrics = evaluate_with_automodel(ft_source, token, "nv-embed-v2-ft")
 
     # Summary
     all_metrics = base_metrics + ft_metrics