GPT2.py

import os 
import torch
from datasets import load_dataset
from transformers import GPT2Tokenizer, GPT2LMHeadModel

from torch.utils.data import DataLoader
from torch.nn import CrossEntropyLoss
from torch.optim import AdamW
from torch.optim.lr_scheduler import StepLR, CosineAnnealingLR, ExponentialLR
from torch.amp import GradScaler, autocast # Using mixed precision training to reduce memory usage and increase computation speed

# For Code Execution
import ast
import sys
from io import StringIO
import contextlib

# Resource management for code execution
from func_timeout import func_timeout, FunctionTimedOut
import resource

import matplotlib.pyplot as plt
from collections import defaultdict

os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128, expandable_segments: True'

def load_data():
    """Load and preprocess the CodeSearchNet dataset."""
    # The dataset is available via the Hugging Face Datasets library
    # Filter for Python functions only
    dataset = load_dataset("claudios/code_search_net", "python", trust_remote_code=True)
    print("Dataset loaded:", dataset)

    # return dataset

    # Use only a small subset of the dataset for experimentation
    small_dataset = dataset["train"].train_test_split(test_size = 0.01, shuffle=True)["test"]
    return small_dataset

def preprocess_data(dataset, tokenizer):
    """Tokenize and prepare the dataset for training."""
    # Save processed data in a format like .jsonl or .pt
    def tokenize_function(examples):
        # Tokenize the code and documentation strings
        max_length = 128 # reduce sequence length during preprocessing to save memory
        code_tokens = tokenizer(examples['func_code_string'], 
                                padding="max_length", 
                                truncation=True,
                                max_length=max_length) # Return PyTorch tensors
        doc_tokens = tokenizer(examples['func_documentation_string'], 
                               padding="max_length", 
                               truncation=True,
                               max_length=max_length) # Return PyTorch tensors

        # Use the code tokens as input_ids and the documentation tokens as labels
        return {
            "input_ids": doc_tokens["input_ids"], # Remove batch dimension
            "attention_mask": doc_tokens["attention_mask"],
            "labels": code_tokens["input_ids"]
        }
    
    tokenized_datasets = dataset.map(tokenize_function, batched=True, batch_size=32)
    tokenized_datasets.set_format("torch")
    # DEBUGGING LINES
    print("Tokenized dataset features:", tokenized_datasets.features)
    print("Sample tokenized data:", tokenized_datasets[0])
    return tokenized_datasets

def build_model():
    """Initialize or modify the GPT-like model and tokenizer."""
    
    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
    tokenizer.add_special_tokens({'pad_token': '[PAD]'})
    # Resize model embeddings to account for new token
    model = GPT2LMHeadModel.from_pretrained("gpt2")
    model.resize_token_embeddings(len(tokenizer))

    # Check if GPU is available and move model to GPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    return model, tokenizer, device

def generate_code(model, tokenizer, prompt, max_length=128):
    """
    Generate python code with structured prompt, with explicit attention masks.
    Args:
        model: GPT-2
        tokenizer: GPT-2 tokenizer
        max_length = Maximum length of generated sequence 
        prompt: String prompt for code generation
    Returns:
        generated_code: String containing the generated Python code
    """
    prompt = prompt.strip()

    # Prepare the prompt with explicit attention mask
    encoded_input = tokenizer(
        prompt,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=max_length,
        add_special_tokens=True
    )

    # Move everything to the correct device
    inputs = {
        "input_ids": encoded_input["input_ids"].to(model.device),
        "attention_mask": encoded_input["attention_mask"].to(model.device)
    }

    # Generate
    with torch.no_grad():
        outputs = model.generate(
            inputs["input_ids"], 
            attention_mask=inputs["attention_mask"],
            max_length=max_length,
            num_return_sequences=1,
            temperature=0.5,
            top_p=0.85, # nucleus sampling
            top_k=20, # top-k sampling
            pad_token_id=tokenizer.eos_token_id,
            bos_token_id=tokenizer.bos_token_id,
            min_length=20,
            do_sample=True,
            no_repeat_ngram_size=4,
            num_beams=5,
            early_stopping=True
        )
    
    # Decode and return the generated code
    generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # Add post-processing to extract only the code portion
    if prompt in generated_code:
        generated_code = generated_code[len(prompt):].strip()

    return generated_code

def execute(code_str, test_input=None, timeout=2):
    """
    Safely execute generated code with test input, and capture its output.
    Args:
        code_str: String containing Python code to execute.
    Returns:
        success: Boolean indicating whether the code executed successfully.
        output: String containing execution output or error message.
    """
    # Add resource limits
    # resource.setrlimit(resource.RLIMIT_AS, (256 * 1024 * 1024, -1))  # GPU
    resource.setrlimit(resource.RLIMIT_CPU, (2, -1))  # 1 second CPU time

    # Create string buffer to capture output
    output_buffer = StringIO()

    try:
        # First, validate Python syntax
        ast.parse(code_str) # ast is the Abstract Syntax Tree module in Python

        # Create execution environment
        global_env = {"__builtins__": __builtins__}
        local_env = {}

        def execute_code():
            # Redirect stdout to capture output
            with contextlib.redirect_stdout(output_buffer):
                # Execute the code in a safe environment
                exec(code_str, global_env, local_env)
                # Execute test input if provided
                if test_input:
                    exec(test_input, global_env, local_env)

        # Execute with timeout
        func_timeout(timeout, execute_code)
        return True, output_buffer.getvalue()
    
    except FunctionTimedOut:
        return False, "Execution timed out"
    except SyntaxError as e:
        return False, f"Syntax Error: {str(e)}"
    except Exception as e:
        return False, f"Runtime Error: {str(e)}"
    finally:
        output_buffer.close()

def evaluate_code(model, tokenizer, test_cases):
    """
    Evaluate the model's code generation capabilities.
    Args:
        model: GPT-2
        tokenizer: GPT-2 tokenizer
        test_cases: List of dictionaries containing:
        - 'func_documentation_string': Documentation for the function
        - 'expected_output': Expected output
        - 'test_input': Test input to execute the generated code
    Returns:
        metrics: Dictionary containing evaluation metrics
    """
    results = {
        'total_cases': len(test_cases),
        'syntax_correct': 0,
        'execution_successful': 0,
        'output_matches': 0
    }

    for test_case in test_cases:
        # Generate code
        generated_code = generate_code(model, tokenizer, test_case['func_documentation_string'])

        # Execute the generated code, using the test input
        success, output = execute(generated_code, test_case['test_input'])

        # Update metrics
        if success:
            results['syntax_correct'] += 1
            results['execution_successful'] += 1

            # Compare output with expected output
            if output.strip() == test_case['expected_output'].strip():
                results['output_matches'] += 1
        
        print(f"\nPrompt: {test_case['func_documentation_string']}")
        print(f"Generated Code:\n{generated_code}")
        print(f"Execution success: {success}")
        print(f"Output: {output}")

    # Calculate percentages
    total = results['total_cases']
    results['syntax_accuracy'] = results['syntax_correct'] / total * 100
    results['execution_accuracy'] = results['execution_successful'] / total * 100
    results['functional_accuracy'] = results['output_matches'] / total * 100

    return results


def train_model(model, train_dataset, valid_dataset):
    """Train the model on the dataset."""
    train_dataloader = DataLoader(train_dataset, batch_size=8, shuffle=True)
    valid_dataloader = DataLoader(valid_dataset, batch_size=8, shuffle=True)
    history = {
        'train_loss': [],
        'valid_loss': []
    }
    print(f"Using device: {model.device}") # DEBUGGING LINE
    # DEBUGGING LINES
    # try:
    #     sample_batch = next(iter(train_dataloader))
    #     if torch.is_tensor(sample_batch['input_ids']):
    #         print(f"Sample batch device: {sample_batch['input_ids'].device}")
    #     else:
    #         print("Warning: input_ids is not a tensor")
    # except Exception as e:
    #     print(f"Debug info error: {e}")

    optimizer = AdamW(
        model.parameters(), 
        lr=5e-5, 
        weight_decay=0.001, 
        eps=1e-8
        ) #Add L2 regularization/weight decay to prevent overfitting
    gradient_accumulation_steps = 8

    num_epochs = 4
    num_training_steps = len(train_dataloader) * num_epochs # 20 epochs
    num_warmup_steps = num_training_steps // 20

    scheduler = ExponentialLR(
        optimizer,
        gamma=0.7 # Decay factor (after each epoch)
    )
    current_lr = scheduler.get_last_lr()[0]  # Get initial learning rate

    criterion = CrossEntropyLoss() # Initialize the loss function

    scaler = GradScaler()

    print("Training process started...")
    for epoch in range(num_epochs):
        model.train()
        total_train_loss = 0
        num_batches = len(train_dataloader)

        for i, batch in enumerate(train_dataloader):
            inputs = batch['input_ids'].to(model.device) # Move input data to GPU (tokenized code inputs)
            labels = batch['labels'].to(model.device) # Move labels to GPU (expected outputs/targets)

            # outputs = model(inputs, labels=labels) # forward pass
            # loss = criterion(outputs.logits.view(-1, outputs.logits.size(-1)), labels.view(-1)) # Calculate loss from model outputs

            # loss.backward() # Calculate gradients with backpropagation
            # optimizer.step() # Update model weights
            
            with autocast(device_type='cuda'): # starts a context manager for automatic mixed precision training (AMP).
                # Certain operations will use lower precision (float16) to save memory and computation time
                outputs = model(inputs, labels = labels)
                loss = outputs.loss
                scaled_loss = loss / gradient_accumulation_steps # Normalize loss by accumulation steps
            
            scaler.scale(scaled_loss).backward() # Backpropagation with PyTorch AMP. Scale and compute gradients.

            # Only update after accumulating gradients
            if (i+1) % gradient_accumulation_steps == 0:
                scaler.unscale_(optimizer) # Unscale gradients
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0) # Gradient clipping (max L2 norm to 5.0)         
                scaler.step(optimizer) # Performs optimizer step with scaled gradients
                scaler.update()
                optimizer.zero_grad(set_to_none=True) # Reset gradients to zero to avoid accumulation
            
            total_train_loss += loss.item() 

            if i % 100 == 0:
                print(f"Epoch {epoch+1}, Batch {i}, Loss: {loss.item()}, Learning Rate: {current_lr}")
                torch.cuda.empty_cache() # Clear GPU cache every 100 batches

        scheduler.step() # Performs optimizer step with scaled gradients
        # Learning rate logging
        current_lr = scheduler.get_last_lr()[0]

        # Print epoch average
        avg_train_loss = total_train_loss / num_batches
        history['train_loss'].append(avg_train_loss)
        # DEBUGGING LINES
        # print("Training dataset keys:", train_dataset.features)  
        # print("First batch keys:", next(iter(DataLoader(train_dataset, batch_size=1)))) 
        
        # Validation after each epoch
        model.eval()
        total_valid_loss = 0

        with torch.no_grad():
            for batch in valid_dataloader:
                inputs = batch['input_ids'].to(model.device)
                labels = batch['labels'].to(model.device)
                outputs = model(inputs, labels=labels)
                total_valid_loss += outputs.loss.item()
        
        avg_valid_loss = total_valid_loss / len(valid_dataloader)
        history['valid_loss'].append(avg_valid_loss)

        print(f"Epoch {epoch+1}, Train_loss:{avg_train_loss:.4f}, Valid_loss {avg_valid_loss:.4f}, Learning Rate: {current_lr}")    

    return history

# Evaluation
def test_model(model, test_dataset):
    """Evaluate on test data."""
    test_dataloader = DataLoader(test_dataset, batch_size=8)
    total_test_loss = 0
    num_batches = len(test_dataloader)

    model.eval()
    print("Testing process started...")
    with torch.no_grad():
        for i, batch in enumerate(test_dataloader):
            inputs = batch["input_ids"].to(model.device)
            labels = batch["labels"].to(model.device)
            outputs = model(inputs, labels=labels)
            total_test_loss += outputs.loss.item()

            if i % 100 == 0:
                print(f"Batch {i}, Loss: {outputs.loss.item()}")

    avg_test_loss = total_test_loss / num_batches
    print(f"Average Test Loss: {avg_test_loss:.4f}")

    return avg_test_loss

def plot_learning_curves(history):
    """Plot training vs. validation learning curves"""
    plt.figure(figsize=(10, 6))
    epochs = range(1, len(history['train_loss']) + 1) # integer x-axis values
    plt.plot(epochs, history['train_loss'], label='Training Loss')
    plt.plot(epochs, history['valid_loss'], label='Validation Loss')
    plt.title('Learning Curves')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    plt.grid(True)
    plt.xticks(epochs) # Force integer ticks on x-axis
    plt.savefig('Learning_Curves_GPT2.png')
    plt.close()

def save_model(model, path="./model_checkpoint_gpt2"):
    """Save the trained model."""
    model.save_pretrained(path)
    print(f"Model saved to {path}")

def load_model(path="./model_checkpoint_gpt2"):
    """Load a saved model."""
    model = GPT2LMHeadModel.from_pretrained(path)
    print(f"Model loaded from {path}")
    return model

def main():
    dataset = load_data()
    print("Original dataset features:", dataset.features) # DEBUGGING LINE
    model, tokenizer, device = build_model()
    tokenized_dataset = preprocess_data(dataset, tokenizer)
    print("Tokenized dataset features:", tokenized_dataset.features) # DEBUGGING LINE

    # split_dataset = dataset["train"].train_test_split(test_size=0.2, shuffle=True)
    # Using the following line for small_dataset for experimentation
    
    split_dataset = tokenized_dataset.train_test_split(test_size=0.2, shuffle=True)
    split_temp = split_dataset["train"].train_test_split(test_size=0.5, shuffle=True)
    train_dataset = split_dataset["train"]
    valid_dataset = split_temp["train"]
    test_dataset = split_temp["test"]

    print(torch.cuda.is_available())  # Should return True if GPU is available
    print(torch.cuda.device_count())  # Number of available GPUs
    print(torch.cuda.get_device_name(0))  # GPU name

    history = train_model(model, train_dataset, valid_dataset)
    test_model(model, test_dataset)
    save_model(model)

    # Generate Learning curves
    plot_learning_curves(history)

    # The following are temporary & for experimentation. Need to be replaced.
    # Define test cases for evaluation
    test_cases = [
        {
            'func_documentation_string': 'Add two numbers a and b',
            'expected_output': '12',
            'test_input': 'print(add(5,7))'
        },
        {
            'func_documentation_string': 'Calculate factorial of a number n',
            'expected_output': '120',
            'test_input': 'print(factorial(5))'
        }
    ]
    
    # Evaluate the model
    print("\nEvaluating model...")
    evaluation_results = evaluate_code(model, tokenizer, test_cases)
    
    # Print evaluation results
    print("\nEvaluation Results:")
    print(f"Syntax Accuracy: {evaluation_results['syntax_accuracy']:.2f}%")
    print(f"Execution Accuracy: {evaluation_results['execution_accuracy']:.2f}%")
    print(f"Functional Accuracy: {evaluation_results['functional_accuracy']:.2f}%")


if __name__ == "__main__":
    main()