gemma-embed / train_stage_2.py

Upload 54 files

464767c verified over 1 year ago

6.11 kB

#!/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()

	#!/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()