CodeParrot Dataset

The CodeParrot Dataset, created by the authors of the "Natural Language Processing with Transformers" book (Lewis Tunstall, Leandro von Werra, and Thomas Wolf), is a comprehensive collection of Python source code designed specifically for training code generation language models. This dataset contains approximately 180GB of deduplicated Python code from public GitHub repositories, making it an ideal resource for learning how to train domain-specific language models and build Python programming assistants.

Available on Hugging Face, this dataset is excellent for training GPT-style code models, building Python code completion tools, learning the end-to-end process of dataset curation and model training, and developing AI-powered programming assistants - all while following best practices from the transformers library creators.

Key Features

• Records: Millions of Python files from GitHub repositories.
• Size: Approximately 180GB of Python source code (after deduplication and filtering).
• Variables/Content:
  • Content: Complete Python file contents
  • Repository Name: Source repository identifier
  • File Path: Original file location within repository
  • License: Repository license information (permissive licenses only)
  • Size: File size in bytes
  • Language: Python (exclusively)
  • Extension: .py files
• Data Type: Text data (Python source code).
• Format: Parquet files, accessible via Hugging Face datasets library.
• Filtering Criteria:
  • Python files only (.py extension)
  • Permissive licenses (MIT, Apache, BSD, etc.)
  • Near-deduplication applied
  • Quality filtering (minimum file size, valid Python syntax)
  • Repository stars threshold (quality proxy)
• Training Splits: Pre-split into train and validation sets.
• Educational Focus: Designed as a teaching dataset with documented curation process.

Why This Dataset

CodeParrot was specifically created to demonstrate the complete process of building a domain-specific language model, from data collection to model deployment. Unlike larger multi-language datasets, its Python-only focus makes it more manageable for learning and experimentation. It's ideal for projects that aim to:

1. Train GPT-style causal language models for Python code generation.
2. Build Python-specific code completion and suggestion tools.
3. Learn the end-to-end pipeline of dataset curation, preprocessing, and model training.
4. Develop programming assistants focused on Python and data science workflows.
5. Create code search and retrieval systems for Python repositories.
6. Study Python coding patterns, best practices, and common idioms.
7. Fine-tune pre-trained models for specific Python domains (web dev, ML, data science).
8. Experiment with code generation techniques on a manageable dataset size.

How to Use the Dataset

1. Load via Hugging Face:

   from datasets import load_dataset
   
   # Load full dataset
   ds = load_dataset("transformersbook/codeparrot")
   
   # Access splits
   train_ds = ds['train']
   valid_ds = ds['valid']
   
   print(f"Training samples: {len(train_ds)}")
   print(f"Validation samples: {len(valid_ds)}")

2. Stream for memory efficiency:

   # Stream without downloading entire dataset
   ds = load_dataset(
       "transformersbook/codeparrot",
       split="train",
       streaming=True
   )
   
   # Iterate through examples
   for i, example in enumerate(ds):
       if i >= 10:  # First 10 examples
           break
       print(f"File: {example['repo_name']}/{example['path']}")
       print(f"Code preview: {example['content'][:200]}...")

3. Explore data structure:

   # Examine a sample
   sample = train_ds[0]
   print("Keys:", sample.keys())
   print("Content length:", len(sample['content']))
   print("Repository:", sample['repo_name'])
   print("File path:", sample['path'])
   print("
Code sample:")
   print(sample['content'][:500])

4. Preprocess for training:

   from transformers import AutoTokenizer
   
   # Use GPT-2 tokenizer (or train custom tokenizer)
   tokenizer = AutoTokenizer.from_pretrained("gpt2")
   
   # Add special tokens for code
   tokenizer.add_special_tokens({
       'additional_special_tokens': ['<|endoftext|>']
   })
   
   def tokenize_function(examples):
       # Tokenize code samples
       outputs = tokenizer(
           examples['content'],
           truncation=True,
           max_length=1024,
           return_overflowing_tokens=True,
           return_length=True
       )
       return outputs
   
   # Apply tokenization
   tokenized_ds = train_ds.map(
       tokenize_function,
       batched=True,
       remove_columns=train_ds.column_names,
       desc="Tokenizing dataset"
   )

5. Filter by quality metrics:

   def filter_short_files(example):
       # Filter out very short files
       return len(example['content']) > 200
   
   def filter_large_files(example):
       # Filter out extremely large files
       return len(example['content']) < 100000
   
   filtered_ds = train_ds.filter(filter_short_files)
   filtered_ds = filtered_ds.filter(filter_large_files)

6. Create custom train/validation splits:

   # If you want different split ratios
   full_ds = load_dataset("transformersbook/codeparrot", split="train")
   
   # Create 95-5 split
   split_ds = full_ds.train_test_split(test_size=0.05, seed=42)
   train_ds = split_ds['train']
   val_ds = split_ds['test']

7. Prepare for causal language modeling:

   from transformers import DataCollatorForLanguageModeling
   
   # Data collator for CLM (auto-generates labels)
   data_collator = DataCollatorForLanguageModeling(
       tokenizer=tokenizer,
       mlm=False  # Causal LM, not masked LM
   )
   
   # Group texts into chunks of block_size
   block_size = 1024
   
   def group_texts(examples):
       # Concatenate all texts
       concatenated = {k: sum(examples[k], []) for k in examples.keys()}
       total_length = len(concatenated[list(examples.keys())[0]])
       
       # Drop last chunk if smaller than block_size
       total_length = (total_length // block_size) * block_size
       
       # Split into chunks
       result = {
           k: [t[i:i + block_size] for i in range(0, total_length, block_size)]
           for k, t in concatenated.items()
       }
       
       # Create labels (copy of input_ids for CLM)
       result["labels"] = result["input_ids"].copy()
       return result
   
   lm_dataset = tokenized_ds.map(
       group_texts,
       batched=True,
       desc="Grouping texts"
   )

8. Sample for experimentation:

   # Take subset for rapid prototyping
   small_train = train_ds.shuffle(seed=42).select(range(10000))
   small_val = valid_ds.shuffle(seed=42).select(range(1000))

9. Analyze dataset statistics:

   import numpy as np
   
   # File length distribution
   lengths = [len(example['content']) for example in train_ds.select(range(10000))]
   print(f"Mean length: {np.mean(lengths):.0f} chars")
   print(f"Median length: {np.median(lengths):.0f} chars")
   print(f"Max length: {max(lengths)} chars")
   
   # Token count estimation
   token_counts = [
       len(tokenizer.encode(example['content'])) 
       for example in train_ds.select(range(1000))
   ]
   print(f"Mean tokens: {np.mean(token_counts):.0f}")
   print(f"Median tokens: {np.median(token_counts):.0f}")

10. Custom preprocessing pipeline:

    import re
    
    def preprocess_code(example):
        code = example['content']
        
        # Remove excessive blank lines
        code = re.sub(r'
s*
s*
', '

', code)
        
        # Remove trailing whitespace
        code = '
'.join(line.rstrip() for line in code.split('
'))
        
        # Add file context (optional)
        code = f"# File: {example['path']}
{code}"
        
        example['content'] = code
        return example
    
    processed_ds = train_ds.map(preprocess_code)

Possible Project Ideas

• Python code completion model training a GPT-style model for intelligent code suggestions.
• Domain-specific fine-tuning adapting the model for web development, data science, or ML engineering.
• Code documentation generator creating docstrings and comments from Python functions.
• Python code search engine building semantic search using code embeddings.
• Bug detection for Python training models to identify common Python errors and anti-patterns.
• Code style transfer converting between different Python coding styles (PEP 8 compliance).
• Interactive Python tutor creating a chatbot that explains and generates Python code.
• Function name prediction learning to suggest appropriate function names from implementations.
• Import statement generator automatically suggesting required imports based on code.
• Code summarization tool generating natural language descriptions of Python functions.
• Python idiom learner identifying and suggesting Pythonic coding patterns.
• Test generation assistant creating unit tests from function signatures and docstrings.
• Code refactoring suggestions identifying opportunities for code improvements.
• Data science code assistant specialized model for pandas, numpy, sklearn workflows.
• Educational coding companion helping beginners write better Python code.

Dataset Challenges and Considerations

• Python-Only Limitation: Only covers Python; not suitable for multi-language applications.
• Size vs. Diversity: 180GB is substantial but smaller than multi-language datasets like The Stack.
• Quality Variation: Code quality ranges from production to educational/experimental projects.
• Duplication: Despite deduplication, similar patterns exist across repositories.
• Temporal Relevance: Code may use older Python versions or deprecated libraries.
• Domain Bias: Over-representation of certain domains (web dev, data science) over others.
• License Verification: Permissive licenses assumed but ongoing verification needed.
• Computational Requirements: Still requires significant compute for full model training.
• Library Evolution: Popular libraries change APIs; code may be outdated.
• Privacy Concerns: May contain API keys, credentials, or sensitive information despite filtering.

Training a CodeParrot Model

Architecture Choices:

from transformers import GPT2Config, GPT2LMHeadModel

# Small model for experimentation (84M parameters)
config_small = GPT2Config(
    vocab_size=len(tokenizer),
    n_positions=1024,
    n_embd=768,
    n_layer=12,
    n_head=12
)

# Medium model (350M parameters, similar to original CodeParrot)
config_medium = GPT2Config(
    vocab_size=len(tokenizer),
    n_positions=1024,
    n_embd=1024,
    n_layer=24,
    n_head=16
)

# Initialize model
model = GPT2LMHeadModel(config_medium)
print(f"Model parameters: {model.num_parameters():,}")

Training Setup:

from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir="./codeparrot-small",
    overwrite_output_dir=True,
    num_train_epochs=3,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    gradient_accumulation_steps=16,  # Effective batch size: 64
    learning_rate=5e-4,
    weight_decay=0.1,
    warmup_steps=1000,
    lr_scheduler_type="cosine",
    logging_steps=100,
    save_steps=1000,
    eval_steps=1000,
    save_total_limit=3,
    fp16=True,  # Mixed precision training
    evaluation_strategy="steps",
    load_best_model_at_end=True,
    metric_for_best_model="eval_loss",
    report_to="tensorboard",
    push_to_hub=False
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=lm_dataset['train'],
    eval_dataset=lm_dataset['valid'],
    data_collator=data_collator,
)

# Train model
trainer.train()

Inference Example:

from transformers import pipeline

# Load trained model
code_generator = pipeline(
    "text-generation",
    model="./codeparrot-small",
    tokenizer=tokenizer
)

# Generate code
prompt = """
def calculate_fibonacci(n):
    """Calculate the nth Fibonacci number."""
"""

generated = code_generator(
    prompt,
    max_length=150,
    num_return_sequences=3,
    temperature=0.7,
    top_p=0.95,
    do_sample=True
)

for i, output in enumerate(generated):
    print(f"
--- Generation {i+1} ---")
    print(output['generated_text'])

Evaluation Strategies

Perplexity (Primary Metric):

import torch
from torch.nn import CrossEntropyLoss

def calculate_perplexity(model, dataset):
    model.eval()
    losses = []
    
    for batch in dataset:
        with torch.no_grad():
            outputs = model(**batch)
            loss = outputs.loss
            losses.append(loss.item())
    
    perplexity = torch.exp(torch.tensor(losses).mean())
    return perplexity.item()

Code Completion Accuracy:

• Exact match percentage
• Edit distance from ground truth
• Syntax validity of generated code
• Pass@k metric (if test cases available)

Human Evaluation:

• Code quality and readability
• Correctness of generated functions
• Usefulness of suggestions
• Adherence to Python best practices

Benchmark Datasets:

• HumanEval: Python programming problems (OpenAI)
• MBPP: Mostly Basic Python Problems (Google)
• APPS: Competitive programming problems

Python-Specific Considerations

Common Python Patterns to Learn:

• List comprehensions
• Generator expressions
• Decorators
• Context managers (with statements)
• Exception handling
• Class definitions and inheritance
• Async/await patterns
• Type hints (Python 3.5+)

Popular Libraries Represented:

• Data Science: pandas, numpy, scikit-learn, matplotlib
• Web Development: Django, Flask, FastAPI
• Deep Learning: PyTorch, TensorFlow, Keras
• Testing: pytest, unittest
• Utilities: requests, click, argparse

Python Version Considerations:

• Mix of Python 2 and Python 3 code
• Evolution of syntax (f-strings, walrus operator, match statements)
• Deprecated features vs. modern patterns

Comparison with Other Code Datasets

vs. The Stack:

• Size: CodeParrot (180GB Python) vs. The Stack (6.4TB multi-language)
• Focus: Python-only vs. 358 languages
• Manageability: More approachable for learning vs. production-scale
• Use Case: Education and experimentation vs. state-of-the-art models

vs. GitHub Code:

• Availability: Public vs. proprietary (GitHub Copilot)
• Scale: Smaller vs. larger
• Licensing: Clear permissive licenses vs. mixed

vs. CodeSearchNet:

• Size: Larger (180GB) vs. smaller (2M functions)
• Task: General code generation vs. code search focus
• Languages: Python-only vs. 6 languages

CodeParrot Advantages:

• Perfect for learning model training pipeline
• Python-focused for domain-specific needs
• Manageable size for individual researchers
• Well-documented curation process
• Educational tutorials available

Educational Value

Learning Objectives:

1. Dataset Curation: Understanding data collection and filtering
2. Tokenization: Custom tokenizer training for code
3. Model Training: Causal language modeling from scratch
4. Evaluation: Metrics for code generation quality
5. Deployment: Serving models for code completion
6. Scaling: Techniques for handling large datasets

Accompanying Resources:

• "Natural Language Processing with Transformers" book chapter
• Hugging Face course materials
• Blog posts documenting the creation process
• Example training scripts and notebooks

Computational Requirements

Storage:

• Dataset: ~180GB disk space
• Tokenized dataset: Additional 200-300GB
• Model checkpoints: 1-4GB per checkpoint (depending on size)

Training:

• Small model (84M): 1-2 GPUs (16GB VRAM), ~1 week
• Medium model (350M): 4-8 GPUs (32GB VRAM), ~2 weeks
• Large model (1B+): 16+ GPUs or TPUs, several weeks

Inference:

• CPU: Possible but slow for interactive use
• GPU: 1 GPU (8-16GB) for real-time completion
• Optimization: Quantization, distillation for production deployment

Ethical and Licensing Considerations

Licensing:

• Only permissive open-source licenses included
• MIT, Apache 2.0, BSD licenses most common
• Opt-out mechanism available for repository owners

Attribution:

• Generated code may resemble training data
• Consider adding attribution in generated output
• Respect original authors and licenses

Privacy:

• Potential for API keys, credentials in code
• Personal information in comments
• Filtering applied but not perfect

Responsible Use:

• Don't claim AI-generated code as original work
• Review generated code for security vulnerabilities
• Test thoroughly before production use
• Consider implications for developers and employment