ComputerAssignments/CA8_neural_program_synthesis/scripts/train.py

#!/usr/bin/env python3

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
import numpy as np
import json
import os
import time
from typing import Dict, List, Tuple, Optional
from tqdm import tqdm

from neural_program_synthesis import (
    ProgramSynthesisTrainer,
    Seq2SeqTrainer,
    ProgramEvaluator,
)

class ProgramDataset(Dataset):

    def __init__(self, data: List[Dict], max_length: int = 100):
        self.data = data
        self.max_length = max_length

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

    def __getitem__(self, idx):
        item = self.data[idx]


        input_seq = self._pad_sequence(item["input"], self.max_length)
        target_seq = self._pad_sequence(item["target"], self.max_length)

        return {
            "input_ids": torch.tensor(input_seq, dtype=torch.long),
            "target_ids": torch.tensor(target_seq, dtype=torch.long),
            "input_length": item["input_length"],
            "target_length": item["target_length"],
        }

    def _pad_sequence(self, sequence: List[int], max_length: int) -> List[int]:
        if len(sequence) > max_length:
            return sequence[:max_length]
        else:
            return sequence + [0] * (max_length - len(sequence))

def save_checkpoint(
    model: nn.Module,
    optimizer: optim.Optimizer,
    epoch: int,
    loss: float,
    save_path: str,
    additional_info: Dict = None,
):
    os.makedirs(os.path.dirname(save_path), exist_ok=True)

    checkpoint = {
        "epoch": epoch,
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "loss": loss,
    }

    if additional_info:
        checkpoint.update(additional_info)

    torch.save(checkpoint, save_path)
    print(f"Checkpoint saved to {save_path}")

def load_checkpoint(checkpoint_path: str) -> Dict:
    if not os.path.exists(checkpoint_path):
        raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")

    checkpoint = torch.load(checkpoint_path, map_location="cpu")
    print(f"Checkpoint loaded from {checkpoint_path}")
    return checkpoint

def create_data_loader(
    dataset: Dataset, batch_size: int, shuffle: bool = True
) -> DataLoader:

    def collate_fn(batch):
        input_ids = torch.stack([item["input_ids"] for item in batch])
        target_ids = torch.stack([item["target_ids"] for item in batch])
        input_lengths = torch.tensor([item["input_length"] for item in batch])
        target_lengths = torch.tensor([item["target_length"] for item in batch])

        return {
            "input_ids": input_ids,
            "target_ids": target_ids,
            "input_lengths": input_lengths,
            "target_lengths": target_lengths,
        }

    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        collate_fn=collate_fn,
        num_workers=0,
    )

def train_epoch(
    model: nn.Module,
    trainer: ProgramSynthesisTrainer,
    data_loader: DataLoader,
    device: torch.device,
) -> Dict[str, float]:
    model.train()

    total_loss = 0.0
    num_batches = 0

    progress_bar = tqdm(data_loader, desc="Training")
    for batch in progress_bar:

        batch = {
            k: v.to(device) if isinstance(v, torch.Tensor) else v
            for k, v in batch.items()
        }


        metrics = trainer.train_step(batch)

        total_loss += metrics["loss"]
        num_batches += 1


        progress_bar.set_postfix(loss=f"{metrics['loss']:.4f}")

    avg_loss = total_loss / num_batches if num_batches > 0 else 0.0

    return {"loss": avg_loss, "num_batches": num_batches}

def validate_epoch(
    model: nn.Module,
    trainer: ProgramSynthesisTrainer,
    data_loader: DataLoader,
    device: torch.device,
) -> Dict[str, float]:
    model.eval()

    total_loss = 0.0
    num_batches = 0

    with torch.no_grad():
        progress_bar = tqdm(data_loader, desc="Validation")
        for batch in progress_bar:

            batch = {
                k: v.to(device) if isinstance(v, torch.Tensor) else v
                for k, v in batch.items()
            }


            metrics = trainer.evaluate(batch)

            total_loss += metrics["loss"]
            num_batches += 1


            progress_bar.set_postfix(loss=f"{metrics['loss']:.4f}")

    avg_loss = total_loss / num_batches if num_batches > 0 else 0.0

    return {"loss": avg_loss, "num_batches": num_batches}

def train_model(
    model: nn.Module,
    trainer: ProgramSynthesisTrainer,
    train_dataset: Dataset,
    val_dataset: Optional[Dataset] = None,
    epochs: int = 100,
    batch_size: int = 32,
    device: torch.device = torch.device("cpu"),
    save_dir: str = "./checkpoints",
    patience: int = 10,
    checkpoints_every: int = 10,
) -> Dict:
    print(f"Starting training for {epochs} epochs...")
    print(f"Training samples: {len(train_dataset)}")
    if val_dataset:
        print(f"Validation samples: {len(val_dataset)}")
    print(f"Batch size: {batch_size}")
    print(f"Device: {device}")


    train_loader = create_data_loader(train_dataset, batch_size, shuffle=True)
    val_loader = None
    if val_dataset:
        val_loader = create_data_loader(val_dataset, batch_size, shuffle=False)


    history = {
        "train_loss": [],
        "val_loss": [],
        "epochs": [],
        "best_val_loss": float("inf"),
        "best_epoch": 0,
    }


    no_improvement_count = 0
    start_time = time.time()

    for epoch in range(epochs):
        epoch_start_time = time.time()

        print(f"\nEpoch {epoch + 1}/{epochs}")
        print("-" * 50)


        train_metrics = train_epoch(model, trainer, train_loader, device)
        train_loss = train_metrics["loss"]


        val_loss = None
        if val_loader:
            val_metrics = validate_epoch(model, trainer, val_loader, device)
            val_loss = val_metrics["loss"]


        history["train_loss"].append(train_loss)
        history["epochs"].append(epoch + 1)

        if val_loss is not None:
            history["val_loss"].append(val_loss)


            if val_loss < history["best_val_loss"]:
                history["best_val_loss"] = val_loss
                history["best_epoch"] = epoch + 1
                no_improvement_count = 0


                best_model_path = os.path.join(save_dir, "best_model.pt")
                save_checkpoint(
                    model, trainer.optimizer, epoch + 1, val_loss, best_model_path
                )
            else:
                no_improvement_count += 1
        else:

            no_improvement_count = 0


        epoch_time = time.time() - epoch_start_time
        print(f"Epoch {epoch + 1} Summary:")
        print(f"  Train Loss: {train_loss:.4f}")
        if val_loss is not None:
            print(f"  Val Loss:   {val_loss:.4f}")
            print(
                f"  Best Val Loss: {history['best_val_loss']:.4f} (Epoch {history['best_epoch']})"
            )
        print(f"  Time:       {epoch_time:.2f}s")


        if (epoch + 1) % checkpoints_every == 0:
            checkpoint_path = os.path.join(save_dir, f"checkpoint_epoch_{epoch + 1}.pt")
            additional_info = {"val_loss": val_loss, "timestamp": time.time()}
            save_checkpoint(
                model,
                trainer.optimizer,
                epoch + 1,
                train_loss,
                checkpoint_path,
                additional_info,
            )


        if patience > 0 and no_improvement_count >= patience:
            print(
                f"\nEarly stopping triggered after {patience} epochs without improvement"
            )
            break


    final_model_path = os.path.join(save_dir, "final_model.pt")
    additional_info = {
        "val_loss": val_loss,
        "total_epochs": len(history["epochs"]),
        "final_train_loss": train_loss,
    }
    save_checkpoint(
        model,
        trainer.optimizer,
        len(history["epochs"]),
        train_loss,
        final_model_path,
        additional_info,
    )


    total_time = time.time() - start_time
    print(f"\n" + "=" * 60)
    print("TRAINING SUMMARY")
    print("=" * 60)
    print(f"Total epochs completed: {len(history['epochs'])}")
    print(f"Total training time: {total_time:.2f}s ({total_time/60:.2f} minutes)")
    print(f"Final train loss: {train_loss:.4f}")
    if val_loss is not None:
        print(f"Final val loss: {val_loss:.4f}")
        print(
            f"Best val loss: {history['best_val_loss']:.4f} (Epoch {history['best_epoch']})"
        )
    print(f"Models saved to: {save_dir}")
    print("=" * 60)

    return history

def fine_tune_model(
    model: nn.Module,
    trainer: ProgramSynthesisTrainer,
    fine_tune_dataset: Dataset,
    epochs: int = 10,
    batch_size: int = 32,
    device: torch.device = torch.device("cpu"),
    learning_rate: float = 0.0001,
    save_dir: str = "./fine_tuned_checkpoints",
) -> Dict:
    print(f"Fine-tuning model for {epochs} epochs...")
    print(f"Fine-tuning samples: {len(fine_tune_dataset)}")
    print(f"Learning rate: {learning_rate}")


    for param_group in trainer.optimizer.param_groups:
        param_group["lr"] = learning_rate


    history = train_model(
        model=model,
        trainer=trainer,
        train_dataset=fine_tune_dataset,
        val_dataset=None,
        epochs=epochs,
        batch_size=batch_size,
        device=device,
        save_dir=save_dir,
        patience=0,
        checkpoints_every=5,
    )

    print("Fine-tuning completed!")
    return history

if __name__ == "__main__":
    print("Training module for Neural Program Synthesis")
    print("This module should be imported and used by main.py")