StarCoder.py

# StarCoderBase1B
# https://huggingface.co/bigcode/starcoderbase-1b

import os 
import json
import torch
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig    

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 pandas as pd
import matplotlib.pyplot as plt
from collections import defaultdict

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

# Add after imports
class Config:
    """Configuration class for model parameters"""
    def __init__(self):
        self.test_size = 0.01
        self.max_length = 128
        self.batch_size = 32
        self.num_epochs = 4
        self.learning_rate = 5e-5

# Create global config
config = Config()

test_size = 0.01

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 = test_size, 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 StarCoderBase model and tokenizer."""
    model_name = "bigcode/starcoderbase-1b"

    # Use 4-bit quantization for lower memory usage
    quantization_config =BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True
    )

    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(
        model_name,
        trust_remote_code=True,
        padding_side="left" # Note: StarCoder uses left padding
    )
    tokenizer.pad_token = tokenizer.eos_token # Set pad token

    # Load model with quantization config
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        trust_remote_code=True,
        device_map="auto",
        quantization_config=quantization_config
    )

    # Check if GPU is available and move model to GPU
    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # model.to(device)
    device = next(model.parameters()).device
    print(f"Model loaded on device: {device}")
    return model, tokenizer, device

def generate_code(model, tokenizer, prompt, max_length=128, test_input=None):
    """
    Generate python code with structured prompt, with explicit attention masks.
    Use function name extrated from test input. (This previously posed an issue.)
    Args:
        model: StarCoderBase1B 
        tokenizer: StarCoderBase1B tokenizer
        max_length = Maximum length of generated sequence 
        prompt: String prompt for code generation
    Returns:
        generated_code: String containing the generated Python code
    """
    # Extract function name from test input
    function_name = None
    if test_input:
        import re
        match = re.search(r'print\((\w+)\(', test_input)
        if match:
            function_name = match.group(1)
    
    # Format prompt to include function name
    if function_name:
        formatted_prompt = f"Write a Python function named '{function_name}' that {prompt}\n\n```python\ndef {function_name}"
    else:
        formatted_prompt = f"Write Python code: {prompt}\n\n```python\n"
     
    # Tokenize the prompt
    encoded_input = tokenizer(
        formatted_prompt,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=max_length
    ).to(model.device)
    
    # Generate with appropriate parameters for code
    with torch.no_grad():
        outputs = model.generate(
            **encoded_input,
            max_new_tokens=max_length,
            temperature=0.2,
            top_p=0.95,
            top_k=40,
            repetition_penalty=1.1,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
            num_return_sequences=1
        )
    
    # Decode and improve code extraction (issue: older output includes instruction text)
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Clean up the generated code
    # TODO: Review this code block and correct or comment as necessary.
    # improvements:
        # First looks for Python code blocks (```python) as StarCoder often uses these
        # Better handling of indentation standardization
        # Preserves empty lines for better code readability
        # More robust fallback mechanism if code block markers aren't found
        # Handles cases where the function definition might not start with "def"
        # Standardizes all indentation to 4 spaces while preserving relative indentation levels
    if "```python" in generated_text:
        # Extract code between Python code block markers
        code_start = generated_text.find("```python") + 9
        code_end = generated_text.find("```", code_start)
        if code_end == -1:  # If no closing backticks, take rest of text
            generated_code = generated_text[code_start:].strip()
        else:
            generated_code = generated_text[code_start:code_end].strip()
    else:
        # Fallback to looking for function definition
        if "def " in generated_text:
            code_start = generated_text.find("def ")
            generated_code = generated_text[code_start:].strip()
        else:
            generated_code = generated_text.strip()

    # Process the extracted code
    if generated_code:
        lines = generated_code.split('\n')
        processed_lines = []
        indent_level = None

        for line in lines:
            if line.strip():  # Non-empty line
                if indent_level is None and not line.startswith('def '):
                    # Set indent level based on first non-empty, non-def line
                    indent_level = len(line) - len(line.lstrip())
                
                if indent_level is not None and not line.startswith('def '):
                    # Standardize indentation to 4 spaces
                    processed_lines.append(' ' * 4 + line[indent_level:])
                else:
                    processed_lines.append(line)
            else:
                processed_lines.append('')  # Preserve empty lines

        generated_code = '\n'.join(processed_lines)

        # Ensure proper indentation
        lines = generated_code.split('\n')
        if len(lines) > 1:
            base_indent = len(lines[1]) - len(lines[1].lstrip())
            formatted_lines = [lines[0]]  # Keep first line as is
            for line in lines[1:]:
                if line.strip():  # Skip empty lines
                    formatted_lines.append(" " * 4 + line[base_indent:])
            generated_code = '\n'.join(formatted_lines)
    else:
        generated_code = "def " + generated_text.split("def", 1)[1] if "def" in generated_text else generated_text

        
    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: CodeQwen1.5-1.8B-Chat
        tokenizer: CodeQwen1.5-1.8B-Chat 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,
        'test_results': []
    }

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

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

        # Create individual test result
        test_result = {
            'generated_code': generated_code,
            'success': success,
            'output': output,
            'output_matches': False
        }

        # 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

        # Add to test results
        results['test_results'].append(test_result)
        
        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=2, shuffle=True)
    valid_dataloader = DataLoader(valid_dataset, batch_size=2, 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.01, 
        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

    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
            
            outputs = model(inputs, labels = labels)
            loss = outputs.loss

            # Only update after accumulating gradients
            if (i+1) % gradient_accumulation_steps == 0:
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0) # Gradient clipping (max L2 norm to 5.0)         
                optimizer.step()
                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_StarCoderBase.png')
    plt.close()

def save_model(model, tokenizer, history=None, path="./model_checkpoint_starcoderbase"):
    """Save the trained StarCoderBase model."""
    os.makedirs(path, exist_ok=True)
    # Save the quantized model
    model.save_pretrained(
        path,
        safe_serialization=True # Use safetensors format
    )

    # Save the tokenizer
    tokenizer.save_pretrained(path)
    print(f"Model saved to {path}")

    # Save training history
    if history:
        with open(os.path.join(path, "training_history.json"), "w") as f:
            json.dump(history, f)
        print(f"Training history saved to {path}")
        
    print(f"Model saved to {path}")

def load_model(path="./model_checkpoint_starcoderbase"):
    """Load the saved StarCoderBase model and tokenizer."""
    model = AutoModelForCausalLM.from_pretrained(
        path,
        trust_remote_code=True,
        device_map="auto",
        quantization_config=BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4big_quant_type="nf4",
            bnb_4bit_use_double_quant=True
        )
    )

    # Load the tokenizer
    tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)

    print(f"Model and tokenizer loaded from {path}")
    return model, tokenizer

def main():
    # Add this test code temporarily
    from huggingface_hub import HfApi
    api = HfApi()
    try:
        api.list_models(author="bigcode")
        print("Successfully connected to Hugging Face!")
    except Exception as e:
        print(f"Hugging Face connection failed: {e}")

    if test_size > 0.0:
        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, tokenizer,history)

        # Generate Learning curves
        plot_learning_curves(history)
    else:
        # Load pretrained model only
        model, tokenizer, device = build_model()
        evaluate_on_humaneval(model, tokenizer)

    # 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))'
        },
        {
            'func_documentation_string': 'Check if a number is prime',
            'expected_output': 'True',
            'test_input': 'print(is_prime(17))'
        },
        {
            'func_documentation_string': 'Reverse a string',
            'expected_output': 'olleh',
            'test_input': 'print(reverse_string("hello"))'
        },
        {
            'func_documentation_string': 'Find the maximum value in a list of numbers',
            'expected_output': '42',
            'test_input': 'print(find_max([5, 42, 17, 8, 1]))'
        },
        {
            'func_documentation_string': 'Count the number of vowels in a string',
            'expected_output': '2',
            'test_input': 'print(count_vowels("hello"))'
        },
        {
            'func_documentation_string': 'Generate Fibonacci sequence up to n terms',
            'expected_output': '[0, 1, 1, 2, 3]',
            'test_input': 'print(fibonacci(5))'
        },
        {
            'func_documentation_string': 'Convert temperature from Celsius to Fahrenheit',
            'expected_output': '98.6',
            'test_input': 'print(celsius_to_fahrenheit(37))'
        }
    ]
    
    # # 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()