train.py

#!/usr/bin/env python
"""
Quick training script for VulnAI Vulnerability Detection Model

Usage:
    python train.py
    python train.py --epochs 20 --batch-size 32
    python train.py --bigvul --epochs 5
"""

import argparse
import os
from pathlib import Path

from vulnai.models.trainer import ModelTrainer, TrainingConfig
from vulnai.data.loader import load_training_data, load_bigvul_training_data
from vulnai.core.logger import setup_logger


logger = setup_logger()


def main():
    parser = argparse.ArgumentParser(description="Train VulnAI model")
    parser.add_argument("--epochs", type=int, default=10, help="Number of training epochs")
    parser.add_argument("--batch-size", type=int, default=16, help="Batch size")
    parser.add_argument("--lr", type=float, default=2e-5, help="Learning rate")
    parser.add_argument("--output", type=str, default="models/trained/vulnai_classifier.pt", 
                        help="Output model path")
    parser.add_argument("--data", type=str, default="data/processed", help="Data directory")
    parser.add_argument("--bigvul", action="store_true", help="Use BigVul dataset")
    parser.add_argument("--bigvul-file", type=str, default="bigvul_20k.json", 
                        help="Path to BigVul JSON file")
    parser.add_argument("--no-balance", action="store_true", help="Disable class balancing")
    
    args = parser.parse_args()
    
    logger.info("Loading training data...")
    
    # Determine which dataset to use
    use_bigvul = args.bigvul or os.path.exists(args.bigvul_file)
    
    # Load training data
    try:
        if use_bigvul:
            # Use BigVul dataset
            bigvul_path = args.bigvul_file if os.path.exists(args.bigvul_file) else args.bigvul_file
            logger.info(f"Loading BigVul dataset from {bigvul_path}")
            train_split, val_split, test_split = load_bigvul_training_data(
                bigvul_file=bigvul_path,
                val_size=0.1,
                test_size=0.1,
                balance=not args.no_balance,
                include_safe=True
            )
            if train_split is None:
                raise ValueError("Failed to load BigVul data")
            if len(train_split.codes) == 0:
                raise ValueError("BigVul data loaded but has no samples")
        else:
            # Use default data directory
            train_split, val_split, test_split = load_training_data(
                data_path=args.data,
                val_size=0.1,
                test_size=0.1,
                balance=not args.no_balance
            )
    except Exception as e:
        logger.warning(f"Could not load data: {e}")
        logger.info("Using sample data for demonstration...")
        from vulnai.data.loader import SAMPLE_CODES, SAMPLE_LABELS, SAMPLE_LANGUAGES
        
        from sklearn.model_selection import train_test_split
        train_codes, temp_codes, train_labels, temp_labels, train_langs, temp_langs = train_test_split(
            SAMPLE_CODES, SAMPLE_LABELS, SAMPLE_LANGUAGES,
            test_size=0.2, random_state=42
        )
        val_codes, test_codes, val_labels, test_labels, val_langs, test_langs = train_test_split(
            temp_codes, temp_labels, temp_langs,
            test_size=0.5, random_state=42
        )
        
        from vulnai.data.loader import DatasetSplit
        train_split = DatasetSplit(train_codes, train_labels, train_langs)
        val_split = DatasetSplit(val_codes, val_labels, val_langs)
        test_split = DatasetSplit(test_codes, test_labels, test_langs)
    
    logger.info(f"Training samples: {len(train_split.codes)}")
    logger.info(f"Validation samples: {len(val_split.codes)}")
    logger.info(f"Test samples: {len(test_split.codes)}")
    
    # Configure training
    config = TrainingConfig(
        model_name="microsoft/codebert-base",
        num_epochs=args.epochs,
        batch_size=args.batch_size,
        learning_rate=args.lr,
        num_classes=11,
        device="cuda" if os.environ.get("CUDA_VISIBLE_DEVICES") else "cpu"
    )
    
    logger.info(f"Training configuration: {config}")
    
    # Create trainer
    trainer = ModelTrainer(config)
    
    # Create data loaders
    from vulnai.models.trainer import VulnerabilityDataset
    
    train_dataset = VulnerabilityDataset(
        train_split.codes,
        [0] * len(train_split.codes),  # Placeholder labels
        trainer.tokenizer,
        config.max_seq_length
    )
    
    val_dataset = VulnerabilityDataset(
        val_split.codes,
        [0] * len(val_split.codes),
        trainer.tokenizer,
        config.max_seq_length
    )
    
    # Use simple labels for demo
    from vulnai.data.loader import DatasetLoader
    loader = DatasetLoader()
    train_labels_idx = loader.convert_labels_to_indices(train_split.labels)
    val_labels_idx = loader.convert_labels_to_indices(val_split.labels)
    
    train_dataset.labels = train_labels_idx
    val_dataset.labels = val_labels_idx
    
    from torch.utils.data import DataLoader
    train_loader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=config.batch_size, shuffle=False)
    
    # Train
    logger.info("Starting training...")
    history = trainer.train(
        train_loader,
        val_loader,
        num_epochs=args.epochs,
        save_best=True,
        save_path=args.output
    )
    
    logger.info(f"Training complete! Model saved to {args.output}")
    
    # Print final metrics
    logger.info(f"Final train loss: {history['train_loss'][-1]:.4f}")
    logger.info(f"Final val loss: {history['val_loss'][-1]:.4f}")


if __name__ == "__main__":
    main()