Complete package: 18 experts + full codebase

ae1e90a 11 months ago

11.9 kB

"""
Training script for the Shared Encoder
This script trains the encoder on a mixed corpus from all domains.
The encoder is trained once and then frozen for expert training.
Usage:
    python train.py --data_dir ../../data --output_dir ../../experts --epochs 10
"""

import os
import argparse
import json
from pathlib import Path
from typing import Dict, List
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR
from transformers import AutoTokenizer
from tqdm import tqdm
import wandb

from model import create_encoder, ContrastiveLoss


class MixedCorpusDataset(Dataset):
    """Dataset that combines samples from all expert domains."""
    
    def __init__(
        self,
        data_dir: Path,
        tokenizer,
        max_length: int = 512,
        samples_per_expert: int = 5000
    ):
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.samples = []
        
        print("Loading datasets from all expert domains...")
        
        # Load samples from each expert's data
        expert_dirs = sorted([d for d in data_dir.iterdir() if d.is_dir()])
        
        for expert_dir in tqdm(expert_dirs, desc="Loading experts"):
            metadata_path = expert_dir / "metadata.json"
            
            if not metadata_path.exists():
                continue
            
            with open(metadata_path) as f:
                metadata = json.load(f)
            
            # Skip synthetic datasets
            if metadata.get("synthetic", False):
                continue
            
            # Load raw datasets
            expert_samples = self._load_expert_data(expert_dir, samples_per_expert)
            
            # Add domain label
            for sample in expert_samples:
                sample["domain"] = metadata["expert_id"]
            
            self.samples.extend(expert_samples)
        
        print(f"Loaded {len(self.samples)} total samples from {len(expert_dirs)} domains")
    
    def _load_expert_data(self, expert_dir: Path, max_samples: int) -> List[Dict]:
        """Load samples from an expert's raw datasets."""
        from datasets import load_from_disk
        
        samples = []
        raw_datasets = sorted(expert_dir.glob("raw_*"))
        
        for raw_path in raw_datasets:
            try:
                dataset = load_from_disk(str(raw_path))
                
                # Extract text from various dataset formats
                for item in dataset:
                    text = self._extract_text(item)
                    if text and len(text) > 20:  # Minimum length filter
                        samples.append({"text": text})
                    
                    if len(samples) >= max_samples:
                        break
                
                if len(samples) >= max_samples:
                    break
            except Exception as e:
                print(f"  Warning: Could not load {raw_path}: {e}")
                continue
        
        return samples[:max_samples]
    
    def _extract_text(self, item: Dict) -> str:
        """Extract text from various dataset formats."""
        # Try common text fields
        for field in ["text", "sentence", "question", "document", "article", "code"]:
            if field in item and item[field]:
                return str(item[field])
        
        # For QA datasets
        if "question" in item and "context" in item:
            return f"{item['question']} {item['context']}"
        
        # For summarization
        if "article" in item:
            return item["article"]
        
        # Fallback: convert entire item to string
        return str(item)
    
    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        sample = self.samples[idx]
        
        # Tokenize
        encoding = self.tokenizer(
            sample["text"],
            max_length=self.max_length,
            padding="max_length",
            truncation=True,
            return_tensors="pt"
        )
        
        return {
            "input_ids": encoding["input_ids"].squeeze(0),
            "attention_mask": encoding["attention_mask"].squeeze(0),
            "domain": hash(sample["domain"]) % 1000  # Simple domain label
        }


def train_epoch(
    model: nn.Module,
    dataloader: DataLoader,
    optimizer: torch.optim.Optimizer,
    criterion: nn.Module,
    device: torch.device,
    epoch: int
) -> float:
    """Train for one epoch."""
    model.train()
    total_loss = 0
    
    pbar = tqdm(dataloader, desc=f"Epoch {epoch}")
    
    for batch_idx, batch in enumerate(pbar):
        input_ids = batch["input_ids"].to(device)
        attention_mask = batch["attention_mask"].to(device)
        labels = batch["domain"].to(device)
        
        # Forward pass
        embeddings = model(input_ids, attention_mask)
        
        # Compute loss
        loss = criterion(embeddings, labels)
        
        # Backward pass
        optimizer.zero_grad()
        loss.backward()
        
        # Gradient clipping
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        
        optimizer.step()
        
        # Track metrics
        total_loss += loss.item()
        avg_loss = total_loss / (batch_idx + 1)
        
        pbar.set_postfix({"loss": f"{avg_loss:.4f}"})
        
        # Log to wandb
        if batch_idx % 100 == 0:
            wandb.log({
                "train/loss": loss.item(),
                "train/avg_loss": avg_loss,
                "train/epoch": epoch
            })
    
    return total_loss / len(dataloader)


@torch.no_grad()
def validate(
    model: nn.Module,
    dataloader: DataLoader,
    criterion: nn.Module,
    device: torch.device
) -> float:
    """Validate the model."""
    model.eval()
    total_loss = 0
    
    for batch in tqdm(dataloader, desc="Validating"):
        input_ids = batch["input_ids"].to(device)
        attention_mask = batch["attention_mask"].to(device)
        labels = batch["domain"].to(device)
        
        embeddings = model(input_ids, attention_mask)
        loss = criterion(embeddings, labels)
        
        total_loss += loss.item()
    
    return total_loss / len(dataloader)


def main():
    parser = argparse.ArgumentParser(description="Train shared encoder")
    parser.add_argument("--data_dir", type=str, default="../../data", help="Data directory")
    parser.add_argument("--output_dir", type=str, default="../../experts", help="Output directory")
    parser.add_argument("--epochs", type=int, default=10, help="Number of epochs")
    parser.add_argument("--batch_size", type=int, default=32, help="Batch size")
    parser.add_argument("--lr", type=float, default=1e-4, help="Learning rate")
    parser.add_argument("--samples_per_expert", type=int, default=5000, help="Samples per expert")
    parser.add_argument("--max_length", type=int, default=512, help="Max sequence length")
    parser.add_argument("--wandb_project", type=str, default="lego-moe", help="W&B project name")
    parser.add_argument("--no_wandb", action="store_true", help="Disable W&B logging")
    
    args = parser.parse_args()
    
    # Setup
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")
    
    data_dir = Path(args.data_dir)
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Initialize wandb
    if not args.no_wandb:
        wandb.init(
            project=args.wandb_project,
            name="shared-encoder-training",
            config=vars(args)
        )
    
    # Load tokenizer
    print("Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
    
    # Create datasets
    print("Creating datasets...")
    train_dataset = MixedCorpusDataset(
        data_dir,
        tokenizer,
        max_length=args.max_length,
        samples_per_expert=args.samples_per_expert
    )
    
    # Split into train/val
    train_size = int(0.95 * len(train_dataset))
    val_size = len(train_dataset) - train_size
    train_dataset, val_dataset = torch.utils.data.random_split(
        train_dataset,
        [train_size, val_size]
    )
    
    train_loader = DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=4,
        pin_memory=True
    )
    
    val_loader = DataLoader(
        val_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=4,
        pin_memory=True
    )
    
    # Create model
    print("Creating model...")
    model = create_encoder({
        "vocab_size": tokenizer.vocab_size,
        "embed_dim": 256,
        "num_layers": 4,
        "num_heads": 8,
        "ff_dim": 1024,
        "dropout": 0.1,
        "max_seq_len": args.max_length
    })
    
    model = model.to(device)
    
    num_params, size_mb = model.get_model_size()
    print(f"Model parameters: {num_params:,}")
    print(f"Model size: {size_mb:.2f} MB")
    
    # Create optimizer and scheduler
    optimizer = AdamW(model.parameters(), lr=args.lr, weight_decay=0.01)
    scheduler = CosineAnnealingLR(optimizer, T_max=args.epochs * len(train_loader))
    
    # Create loss function
    criterion = ContrastiveLoss(temperature=0.07)
    
    # Training loop
    print(f"\nStarting training for {args.epochs} epochs...")
    best_val_loss = float("inf")
    
    for epoch in range(1, args.epochs + 1):
        print(f"\n{'='*60}")
        print(f"Epoch {epoch}/{args.epochs}")
        print(f"{'='*60}")
        
        # Train
        train_loss = train_epoch(model, train_loader, optimizer, criterion, device, epoch)
        
        # Validate
        val_loss = validate(model, val_loader, criterion, device)
        
        # Update scheduler
        scheduler.step()
        
        print(f"Train Loss: {train_loss:.4f}")
        print(f"Val Loss: {val_loss:.4f}")
        
        # Log to wandb
        if not args.no_wandb:
            wandb.log({
                "epoch": epoch,
                "train/epoch_loss": train_loss,
                "val/loss": val_loss,
                "lr": optimizer.param_groups[0]["lr"]
            })
        
        # Save best model
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            checkpoint_path = output_dir / "encoder_best.pt"
            torch.save({
                "epoch": epoch,
                "model_state_dict": model.state_dict(),
                "optimizer_state_dict": optimizer.state_dict(),
                "val_loss": val_loss,
                "config": model.__dict__
            }, checkpoint_path)
            print(f"✓ Saved best model to {checkpoint_path}")
        
        # Save checkpoint
        checkpoint_path = output_dir / f"encoder_epoch_{epoch}.pt"
        torch.save({
            "epoch": epoch,
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "val_loss": val_loss
        }, checkpoint_path)
    
    # Export final model
    print("\nExporting model to ONNX...")
    model.eval()
    onnx_path = output_dir / "encoder.onnx"
    model.export_to_onnx(str(onnx_path))
    
    # Save tokenizer
    tokenizer.save_pretrained(str(output_dir / "tokenizer"))
    
    print("\n✅ Training complete!")
    print(f"Best validation loss: {best_val_loss:.4f}")
    print(f"Model saved to: {output_dir}")
    print(f"\nNext steps:")
    print(f"  1. Generate embeddings: python ../datasets/generate_embeddings.py")
    print(f"  2. Train experts: python ../distillation/train_expert.py")
    
    if not args.no_wandb:
        wandb.finish()


if __name__ == "__main__":
    main()

	"""
	Training script for the Shared Encoder

	This script trains the encoder on a mixed corpus from all domains.
	The encoder is trained once and then frozen for expert training.

	Usage:
	python train.py --data_dir ../../data --output_dir ../../experts --epochs 10
	"""

	import os
	import argparse
	import json
	from pathlib import Path
	from typing import Dict, List
	import torch
	import torch.nn as nn
	from torch.utils.data import Dataset, DataLoader
	from torch.optim import AdamW
	from torch.optim.lr_scheduler import CosineAnnealingLR
	from transformers import AutoTokenizer
	from tqdm import tqdm
	import wandb

	from model import create_encoder, ContrastiveLoss


	class MixedCorpusDataset(Dataset):
	"""Dataset that combines samples from all expert domains."""

	def __init__(
	self,
	data_dir: Path,
	tokenizer,
	max_length: int = 512,
	samples_per_expert: int = 5000
	):
	self.tokenizer = tokenizer
	self.max_length = max_length
	self.samples = []

	print("Loading datasets from all expert domains...")

	# Load samples from each expert's data
	expert_dirs = sorted([d for d in data_dir.iterdir() if d.is_dir()])

	for expert_dir in tqdm(expert_dirs, desc="Loading experts"):
	metadata_path = expert_dir / "metadata.json"

	if not metadata_path.exists():
	continue

	with open(metadata_path) as f:
	metadata = json.load(f)

	# Skip synthetic datasets
	if metadata.get("synthetic", False):
	continue

	# Load raw datasets
	expert_samples = self._load_expert_data(expert_dir, samples_per_expert)

	# Add domain label
	for sample in expert_samples:
	sample["domain"] = metadata["expert_id"]

	self.samples.extend(expert_samples)

	print(f"Loaded {len(self.samples)} total samples from {len(expert_dirs)} domains")

	def _load_expert_data(self, expert_dir: Path, max_samples: int) -> List[Dict]:
	"""Load samples from an expert's raw datasets."""
	from datasets import load_from_disk

	samples = []
	raw_datasets = sorted(expert_dir.glob("raw_*"))

	for raw_path in raw_datasets:
	try:
	dataset = load_from_disk(str(raw_path))

	# Extract text from various dataset formats
	for item in dataset:
	text = self._extract_text(item)
	if text and len(text) > 20: # Minimum length filter
	samples.append({"text": text})

	if len(samples) >= max_samples:
	break

	if len(samples) >= max_samples:
	break
	except Exception as e:
	print(f" Warning: Could not load {raw_path}: {e}")
	continue

	return samples[:max_samples]

	def _extract_text(self, item: Dict) -> str:
	"""Extract text from various dataset formats."""
	# Try common text fields
	for field in ["text", "sentence", "question", "document", "article", "code"]:
	if field in item and item[field]:
	return str(item[field])

	# For QA datasets
	if "question" in item and "context" in item:
	return f"{item['question']} {item['context']}"

	# For summarization
	if "article" in item:
	return item["article"]

	# Fallback: convert entire item to string
	return str(item)

	def __len__(self):
	return len(self.samples)

	def __getitem__(self, idx):
	sample = self.samples[idx]

	# Tokenize
	encoding = self.tokenizer(
	sample["text"],
	max_length=self.max_length,
	padding="max_length",
	truncation=True,
	return_tensors="pt"
	)

	return {
	"input_ids": encoding["input_ids"].squeeze(0),
	"attention_mask": encoding["attention_mask"].squeeze(0),
	"domain": hash(sample["domain"]) % 1000 # Simple domain label
	}


	def train_epoch(
	model: nn.Module,
	dataloader: DataLoader,
	optimizer: torch.optim.Optimizer,
	criterion: nn.Module,
	device: torch.device,
	epoch: int
	) -> float:
	"""Train for one epoch."""
	model.train()
	total_loss = 0

	pbar = tqdm(dataloader, desc=f"Epoch {epoch}")

	for batch_idx, batch in enumerate(pbar):
	input_ids = batch["input_ids"].to(device)
	attention_mask = batch["attention_mask"].to(device)
	labels = batch["domain"].to(device)

	# Forward pass
	embeddings = model(input_ids, attention_mask)

	# Compute loss
	loss = criterion(embeddings, labels)

	# Backward pass
	optimizer.zero_grad()
	loss.backward()

	# Gradient clipping
	torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

	optimizer.step()

	# Track metrics
	total_loss += loss.item()
	avg_loss = total_loss / (batch_idx + 1)

	pbar.set_postfix({"loss": f"{avg_loss:.4f}"})

	# Log to wandb
	if batch_idx % 100 == 0:
	wandb.log({
	"train/loss": loss.item(),
	"train/avg_loss": avg_loss,
	"train/epoch": epoch
	})

	return total_loss / len(dataloader)


	@torch.no_grad()
	def validate(
	model: nn.Module,
	dataloader: DataLoader,
	criterion: nn.Module,
	device: torch.device
	) -> float:
	"""Validate the model."""
	model.eval()
	total_loss = 0

	for batch in tqdm(dataloader, desc="Validating"):
	input_ids = batch["input_ids"].to(device)
	attention_mask = batch["attention_mask"].to(device)
	labels = batch["domain"].to(device)

	embeddings = model(input_ids, attention_mask)
	loss = criterion(embeddings, labels)

	total_loss += loss.item()

	return total_loss / len(dataloader)


	def main():
	parser = argparse.ArgumentParser(description="Train shared encoder")
	parser.add_argument("--data_dir", type=str, default="../../data", help="Data directory")
	parser.add_argument("--output_dir", type=str, default="../../experts", help="Output directory")
	parser.add_argument("--epochs", type=int, default=10, help="Number of epochs")
	parser.add_argument("--batch_size", type=int, default=32, help="Batch size")
	parser.add_argument("--lr", type=float, default=1e-4, help="Learning rate")
	parser.add_argument("--samples_per_expert", type=int, default=5000, help="Samples per expert")
	parser.add_argument("--max_length", type=int, default=512, help="Max sequence length")
	parser.add_argument("--wandb_project", type=str, default="lego-moe", help="W&B project name")
	parser.add_argument("--no_wandb", action="store_true", help="Disable W&B logging")

	args = parser.parse_args()

	# Setup
	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
	print(f"Using device: {device}")

	data_dir = Path(args.data_dir)
	output_dir = Path(args.output_dir)
	output_dir.mkdir(parents=True, exist_ok=True)

	# Initialize wandb
	if not args.no_wandb:
	wandb.init(
	project=args.wandb_project,
	name="shared-encoder-training",
	config=vars(args)
	)

	# Load tokenizer
	print("Loading tokenizer...")
	tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

	# Create datasets
	print("Creating datasets...")
	train_dataset = MixedCorpusDataset(
	data_dir,
	tokenizer,
	max_length=args.max_length,
	samples_per_expert=args.samples_per_expert
	)

	# Split into train/val
	train_size = int(0.95 * len(train_dataset))
	val_size = len(train_dataset) - train_size
	train_dataset, val_dataset = torch.utils.data.random_split(
	train_dataset,
	[train_size, val_size]
	)

	train_loader = DataLoader(
	train_dataset,
	batch_size=args.batch_size,
	shuffle=True,
	num_workers=4,
	pin_memory=True
	)

	val_loader = DataLoader(
	val_dataset,
	batch_size=args.batch_size,
	shuffle=False,
	num_workers=4,
	pin_memory=True
	)

	# Create model
	print("Creating model...")
	model = create_encoder({
	"vocab_size": tokenizer.vocab_size,
	"embed_dim": 256,
	"num_layers": 4,
	"num_heads": 8,
	"ff_dim": 1024,
	"dropout": 0.1,
	"max_seq_len": args.max_length
	})

	model = model.to(device)

	num_params, size_mb = model.get_model_size()
	print(f"Model parameters: {num_params:,}")
	print(f"Model size: {size_mb:.2f} MB")

	# Create optimizer and scheduler
	optimizer = AdamW(model.parameters(), lr=args.lr, weight_decay=0.01)
	scheduler = CosineAnnealingLR(optimizer, T_max=args.epochs * len(train_loader))

	# Create loss function
	criterion = ContrastiveLoss(temperature=0.07)

	# Training loop
	print(f"\nStarting training for {args.epochs} epochs...")
	best_val_loss = float("inf")

	for epoch in range(1, args.epochs + 1):
	print(f"\n{'='*60}")
	print(f"Epoch {epoch}/{args.epochs}")
	print(f"{'='*60}")

	# Train
	train_loss = train_epoch(model, train_loader, optimizer, criterion, device, epoch)

	# Validate
	val_loss = validate(model, val_loader, criterion, device)

	# Update scheduler
	scheduler.step()

	print(f"Train Loss: {train_loss:.4f}")
	print(f"Val Loss: {val_loss:.4f}")

	# Log to wandb
	if not args.no_wandb:
	wandb.log({
	"epoch": epoch,
	"train/epoch_loss": train_loss,
	"val/loss": val_loss,
	"lr": optimizer.param_groups[0]["lr"]
	})

	# Save best model
	if val_loss < best_val_loss:
	best_val_loss = val_loss
	checkpoint_path = output_dir / "encoder_best.pt"
	torch.save({
	"epoch": epoch,
	"model_state_dict": model.state_dict(),
	"optimizer_state_dict": optimizer.state_dict(),
	"val_loss": val_loss,
	"config": model.__dict__
	}, checkpoint_path)
	print(f"✓ Saved best model to {checkpoint_path}")

	# Save checkpoint
	checkpoint_path = output_dir / f"encoder_epoch_{epoch}.pt"
	torch.save({
	"epoch": epoch,
	"model_state_dict": model.state_dict(),
	"optimizer_state_dict": optimizer.state_dict(),
	"val_loss": val_loss
	}, checkpoint_path)

	# Export final model
	print("\nExporting model to ONNX...")
	model.eval()
	onnx_path = output_dir / "encoder.onnx"
	model.export_to_onnx(str(onnx_path))

	# Save tokenizer
	tokenizer.save_pretrained(str(output_dir / "tokenizer"))

	print("\n✅ Training complete!")
	print(f"Best validation loss: {best_val_loss:.4f}")
	print(f"Model saved to: {output_dir}")
	print(f"\nNext steps:")
	print(f" 1. Generate embeddings: python ../datasets/generate_embeddings.py")
	print(f" 2. Train experts: python ../distillation/train_expert.py")

	if not args.no_wandb:
	wandb.finish()


	if __name__ == "__main__":
	main()