Upload train_stage_2.py

train_stage_2.py

@@ -0,0 +1,163 @@
+#!/usr/bin/env python3
+import os
+import csv
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    get_linear_schedule_with_warmup
+)
+from peft import PeftModel
+from torch.cuda.amp import autocast, GradScaler
+from tqdm.auto import tqdm
+from multiprocessing import freeze_support
+
+class TripletDataset(Dataset):
+    def __init__(self, path):
+        self.samples = []
+        with open(path, newline="") as f:
+            reader = csv.DictReader(f)
+            for row in reader:
+                a_ids   = torch.tensor(list(map(int, row["a_ids"].split())), dtype=torch.long)
+                a_mask  = torch.tensor(list(map(int, row["a_mask"].split())), dtype=torch.long)
+                p_ids   = torch.tensor(list(map(int, row["p_ids"].split())), dtype=torch.long)
+                p_mask  = torch.tensor(list(map(int, row["p_mask"].split())), dtype=torch.long)
+                n_ids   = torch.tensor(list(map(int, row["n_ids"].split())), dtype=torch.long)
+                n_mask  = torch.tensor(list(map(int, row["n_mask"].split())), dtype=torch.long)
+                self.samples.append((a_ids, a_mask, p_ids, p_mask, n_ids, n_mask))
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        return self.samples[idx]
+
+class GemmaTripletModel(nn.Module):
+    def __init__(self, peft_model):
+        super().__init__()
+        self.peft = peft_model
+        H = peft_model.config.hidden_size
+        self.proj = nn.Sequential(
+            nn.Linear(H, 512),
+            nn.ReLU(),
+            nn.Linear(512, 256),
+        )
+
+    def forward(self, ids, mask):
+        out = self.peft.base_model(          # PeftModel stores underlying model as .base_model
+            input_ids=ids,
+            attention_mask=mask,
+            output_hidden_states=True,
+            return_dict=True
+        )
+        last = out.hidden_states[-1]         # (B, T, H)
+        pooled = last.mean(dim=1)            # mean pooling
+        z = self.proj(pooled)
+        norm = z.norm(p=2, dim=1, keepdim=True).clamp_min(1e-6)
+        return z / norm
+
+def collate_fn(batch):
+    return tuple(torch.stack(x) for x in zip(*batch))
+
+def main():
+    # --- Config ---
+    MODEL_NAME   = "google/gemma-3-1b-pt"
+    STAGE1_DIR   = "stage1_simcse/final"
+    TRAIN_FILE   = "train.csv"
+    VAL_FILE     = "val.csv"
+    BATCH_SIZE   = 12
+    LR           = 1e-5
+    WEIGHT_DECAY = 0.01
+    NUM_EPOCHS   = 3
+    MARGIN       = 0.2
+    OUTPUT_DIR   = "phase2_triplet_amp"
+    SEED         = 42
+
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    torch.manual_seed(SEED)
+
+    # --- Tokenizer & PEFT Model (load Stage 1) ---
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)
+    base = AutoModelForCausalLM.from_pretrained(MODEL_NAME, attn_implementation="eager")
+    peft_model = PeftModel.from_pretrained(base, STAGE1_DIR)  # loads your Stage 1 LoRA weights
+
+    # --- Embed + Projector ---
+    model = GemmaTripletModel(peft_model).to(device)
+
+    # --- Datasets & Loaders ---
+    train_ds = TripletDataset(TRAIN_FILE)
+    val_ds   = TripletDataset(VAL_FILE)
+    train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_fn)
+    val_loader   = DataLoader(val_ds,   batch_size=BATCH_SIZE, shuffle=False, collate_fn=collate_fn)
+
+    # --- Optimizer, Scheduler, AMP ---
+    optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)
+    total_steps = len(train_loader) * NUM_EPOCHS
+    scheduler = get_linear_schedule_with_warmup(
+        optimizer,
+        num_warmup_steps=int(0.1 * total_steps),
+        num_training_steps=total_steps
+    )
+    scaler = GradScaler()
+
+    # --- Loss ---
+    triplet_loss = nn.TripletMarginLoss(margin=MARGIN, p=2)
+
+    # --- Training Loop ---
+    for epoch in range(1, NUM_EPOCHS + 1):
+        model.train()
+        running_loss = 0.0
+        for a_ids, a_mask, p_ids, p_mask, n_ids, n_mask in tqdm(train_loader, desc=f"Train {epoch}", unit="batch"):
+            a_ids, a_mask = a_ids.to(device), a_mask.to(device)
+            p_ids, p_mask = p_ids.to(device), p_mask.to(device)
+            n_ids, n_mask = n_ids.to(device), n_mask.to(device)
+
+            optimizer.zero_grad()
+            with autocast():
+                emb_a = model(a_ids, a_mask)
+                emb_p = model(p_ids, p_mask)
+                emb_n = model(n_ids, n_mask)
+                loss  = triplet_loss(emb_a, emb_p, emb_n)
+
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+            scheduler.step()
+            running_loss += loss.item()
+
+        print(f"Epoch {epoch} Train Loss: {running_loss/len(train_loader):.6f}")
+
+        # --- Validation ---
+        model.eval()
+        val_loss = 0.0
+        with torch.no_grad():
+            for a_ids, a_mask, p_ids, p_mask, n_ids, n_mask in tqdm(val_loader, desc=f"Val {epoch}", unit="batch"):
+                a_ids, a_mask = a_ids.to(device), a_mask.to(device)
+                p_ids, p_mask = p_ids.to(device), p_mask.to(device)
+                n_ids, n_mask = n_ids.to(device), n_mask.to(device)
+                with autocast():
+                    emb_a = model(a_ids, a_mask)
+                    emb_p = model(p_ids, p_mask)
+                    emb_n = model(n_ids, n_mask)
+                    val_loss += triplet_loss(emb_a, emb_p, emb_n).item()
+
+        print(f"Epoch {epoch} Val   Loss: {val_loss/len(val_loader):.6f}")
+
+        # --- Checkpoint LoRA only ---
+        ckpt_dir = os.path.join(OUTPUT_DIR, f"epoch{epoch}")
+        peft_model.save_pretrained(ckpt_dir)
+        tokenizer.save_pretrained(ckpt_dir)
+
+    # --- Final Save ---
+    final_dir = os.path.join(OUTPUT_DIR, "final")
+    os.makedirs(final_dir, exist_ok=True)
+    peft_model.save_pretrained(final_dir)
+    tokenizer.save_pretrained(final_dir)
+    print("Phase 2 complete. Checkpoints in", OUTPUT_DIR)
+
+if __name__ == "__main__":
+    freeze_support()
+    main()