#!/usr/bin/env python3
"""
Collect real-world code dataset from Hugging Face for fine-tuning a dev-focused model.
Sources:
1. codeparrot/codeclean (Refactoring pairs: messy -> clean)
2. nampd-rr/stack-dataset (StackOverflow Q&A, filtered by score)
3. HuggingFaceH4/codelion (Code instruction following)
4. bigcode/programmingdemos (Real code snippets)
Output: JSON format matching training pipeline (messages format)
"""
import json
import os
import random
import sys
from pathlib import Path
from datasets import load_dataset
BASE_DIR = Path(__file__).parent.parent
PROCESSED_DIR = BASE_DIR / "data" / "processed"
OUTPUT_FILE = PROCESSED_DIR / "dev_finetuning_real.json"
# Target counts per source
TARGET_PER_SOURCE = 300 # Aim for ~1200 total
def load_codeclean():
"""Load code refactoring pairs from CodeClean."""
print("[*] Loading codeparrot/codeclean (Refactoring pairs)...")
try:
ds = load_dataset("codeparrot/codeclean", split="train", streaming=True)
data = []
count = 0
for row in ds:
if count >= TARGET_PER_SOURCE:
break
# CodeClean format: {'source': 'bad_code', 'target': 'clean_code'}
source = row.get('source', '')
target = row.get('target', '')
if not source or not target or len(source) < 50:
continue
entry = {
"messages": [
{
"role": "user",
"content": f"Review this code and suggest improvements for readability and best practices:\n\n```python\n{source}\n```"
},
{
"role": "assistant",
"content": f"Here is the refactored code with improved readability and structure:\n\n```python\n{target}\n```"
}
],
"source": "codeclean",
"tags": ["python", "refactoring", "best_practice", "review"]
}
data.append(entry)
count += 1
print(f" [*] Collected {len(data)} refactoring pairs")
return data
except Exception as e:
print(f" [!] Error loading codeclean: {e}")
return []
def load_stackoverflow():
"""Load high-quality StackOverflow Q&A."""
print("[*] Loading nampd-rr/stack-dataset (StackOverflow)...")
try:
# Load a subset of the dataset
ds = load_dataset("nampd-rr/stack-dataset", "stackoverflow", split="train", streaming=True)
data = []
count = 0
for row in ds:
if count >= TARGET_PER_SOURCE:
break
question = row.get('question', '')
answer = row.get('answer', '')
score = row.get('answer_score', 0)
# Filter for high-score answers (top quality)
if score < 5 or not question or not answer:
continue
# Truncate very long entries
if len(question) > 3000:
question = question[:3000] + "..."
if len(answer) > 5000:
answer = answer[:5000] + "..."
entry = {
"messages": [
{"role": "user", "content": f"Question: {question}\n\nPlease provide a clear, code-based answer."},
{"role": "assistant", "content": f"Answer: {answer}\n\nScore: {score} (High Quality)"}
],
"source": "stackoverflow",
"tags": ["q&a", "explanation", "problem_solving"]
}
data.append(entry)
count += 1
print(f" [*] Collected {len(data)} StackOverflow pairs")
return data
except Exception as e:
print(f" [!] Error loading stackoverflow: {e}")
return []
def load_codelion():
"""Load CodeLion instruction-following data."""
print("[*] Loading HuggingFaceH4/codelion (Code Instructions)...")
try:
ds = load_dataset("HuggingFaceH4/codelion", split="train", streaming=True)
data = []
count = 0
for row in ds:
if count >= TARGET_PER_SOURCE:
break
instruction = row.get('instruction', '')
output = row.get('output', '')
if not instruction or not output:
continue
entry = {
"messages": [
{"role": "user", "content": instruction},
{"role": "assistant", "content": output}
],
"source": "codelion",
"tags": ["instruction_following", "code_generation"]
}
data.append(entry)
count += 1
print(f" [*] Collected {len(data)} CodeLion pairs")
return data
except Exception as e:
print(f" [!] Error loading codelion: {e}")
return []
def load_bigcode_demos():
"""Load BigCode programming demos."""
print("[*] Loading bigcode/programmingdemos (Real Code Snippets)...")
try:
ds = load_dataset("bigcode/programmingdemos", split="train", streaming=True)
data = []
count = 0
for row in ds:
if count >= TARGET_PER_SOURCE // 2: # Less of this, it's raw code
break
code = row.get('code', '')
language = row.get('language', 'unknown')
if not code or len(code) < 100:
continue
# Create a "explain this code" or "improve this code" prompt
entry = {
"messages": [
{"role": "user", "content": f"Review this {language} code for potential issues and improvements:\n\n```{language}\n{code[:2000]}\n```"},
{"role": "assistant", "content": f"Analysis of {language} code:\n\n{code[:2000]}"}
],
"source": "bigcode_demos",
"tags": [language, "code_analysis", "review"]
}
data.append(entry)
count += 1
print(f" [*] Collected {len(data)} BigCode demos")
return data
except Exception as e:
print(f" [!] Error loading bigcode_demos: {e}")
return []
def compile_real_dataset():
"""Compile all real-world sources into a single training dataset."""
print("=" * 60)
print("Dev-AI Real-World Dataset Collection")
print("=" * 60)
all_data = []
# Load from all sources
all_data.extend(load_codeclean())
all_data.extend(load_stackoverflow())
all_data.extend(load_codelion())
all_data.extend(load_bigcode_demos())
# Shuffle the data
random.shuffle(all_data)
# Save
os.makedirs(PROCESSED_DIR, exist_ok=True)
with open(OUTPUT_FILE, "w") as f:
json.dump(all_data, f, indent=2)
# Stats
total_tokens = sum(len(str(d)) for d in all_data)
file_size = OUTPUT_FILE.stat().st_size / (1024 * 1024)
print(f"\n{'=' * 60}")
print(f"Total entries: {len(all_data)}")
print(f"Sources: {', '.join(sorted(set(d['source'] for d in all_data)))}")
print(f"File size: {file_size:.2f} MB")
print(f"Saved to: {OUTPUT_FILE}")
print(f"{'=' * 60}")
return all_data
if __name__ == "__main__":
compile_real_dataset()