vinsblack/CodeReality

1#!/usr/bin/env python32"""3Code Completion Benchmark for CodeReality-1T Dataset4Evaluates code completion models using Pass@k metrics5"""6 7import json8import os9import re10import random11from typing import Dict, List, Tuple, Optional12from collections import defaultdict13 14def load_dataset_sample(data_dir: str, sample_size: int = 200) -> List[Dict]:15    """Load sample of repositories with code files."""16    print(f"🔍 Loading sample of {sample_size} repositories with code files...")17 18    repositories = []19    files = [f for f in os.listdir(data_dir) if f.endswith('.jsonl')]20    random.shuffle(files)21 22    for filename in files[:15]:  # Sample from first 15 files23        file_path = os.path.join(data_dir, filename)24        try:25            with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:26                for line in f:27                    if len(repositories) >= sample_size:28                        break29                    try:30                        repo_data = json.loads(line)31                        # Filter repositories with code files32                        if has_code_files(repo_data):33                            repositories.append(repo_data)34                    except json.JSONDecodeError:35                        continue36        except Exception as e:37            continue38 39        if len(repositories) >= sample_size:40            break41 42    print(f"✅ Loaded {len(repositories)} repositories with code files")43    return repositories44 45def has_code_files(repo: Dict) -> bool:46    """Check if repository contains code files."""47    code_extensions = {'.py', '.js', '.java', '.cpp', '.c', '.go', '.rs', '.ts'}48 49    files = repo.get('files', [])50    for file_obj in files:51        if isinstance(file_obj, dict):52            file_path = file_obj.get('path', '')53            if any(file_path.endswith(ext) for ext in code_extensions):54                return True55    return False56 57def extract_function_snippets(repo: Dict, language: str = 'python') -> List[Dict]:58    """Extract function definitions for completion tasks."""59    snippets = []60 61    # Language-specific patterns62    patterns = {63        'python': r'def\s+(\w+)\s*\([^)]*\):\s*',64        'javascript': r'function\s+(\w+)\s*\([^)]*\)\s*{',65        'java': r'(?:public|private|protected)?\s*(?:static)?\s*\w+\s+(\w+)\s*\([^)]*\)\s*{',66        'cpp': r'\w+\s+(\w+)\s*\([^)]*\)\s*{',67    }68 69    if language not in patterns:70        return snippets71 72    pattern = patterns[language]73    extension_map = {74        'python': '.py',75        'javascript': '.js',76        'java': '.java',77        'cpp': '.cpp'78    }79 80    target_ext = extension_map[language]81 82    files = repo.get('files', [])83    for file_obj in files:84        if isinstance(file_obj, dict):85            file_path = file_obj.get('path', '')86            content = file_obj.get('content', '')87 88            if file_path.endswith(target_ext) and content:89                matches = list(re.finditer(pattern, content, re.MULTILINE))90 91                for match in matches:92                    start_pos = match.start()93                    function_name = match.group(1)94 95                    # Get context before function96                    lines_before = content[:start_pos].split('\n')97                    context_lines = lines_before[-5:] if len(lines_before) >= 5 else lines_before98                    context = '\n'.join(context_lines)99 100                    # Get function body (simplified - until next function or end)101                    remaining_content = content[start_pos:]102                    lines = remaining_content.split('\n')103 104                    function_lines = []105                    indent_level = None106 107                    for i, line in enumerate(lines):108                        if i == 0:109                            function_lines.append(line)110                            continue111 112                        # Detect indentation level113                        if indent_level is None and line.strip():114                            indent_level = len(line) - len(line.lstrip())115 116                        # Stop if we hit same or lower indentation level (end of function)117                        if line.strip() and indent_level is not None:118                            current_indent = len(line) - len(line.lstrip())119                            if current_indent <= indent_level and not line.strip().startswith(('if', 'for', 'while', 'try', 'except', 'else', 'elif')):120                                break121 122                        function_lines.append(line)123 124                        # Limit function length125                        if len(function_lines) > 20:126                            break127 128                    function_body = '\n'.join(function_lines)129 130                    # Create completion task: provide function signature, expect body131                    if len(function_lines) > 3:  # Only meaningful functions132                        snippets.append({133                            'function_name': function_name,134                            'context': context,135                            'prompt': function_lines[0],  # Function signature136                            'completion': '\n'.join(function_lines[1:]),  # Function body137                            'file_path': file_path,138                            'language': language139                        })140 141    return snippets142 143def simple_code_completion_model(prompt: str, language: str) -> List[str]:144    """Simple rule-based code completion for demonstration."""145    completions = []146 147    # Generate multiple completions (for Pass@k evaluation)148    templates = {149        'python': [150            "    pass",151            "    return None",152            "    # TODO: implement this function\n    pass",153            "    result = None\n    return result",154            "    # Implementation needed\n    raise NotImplementedError()"155        ],156        'javascript': [157            "    return null;",158            "    // TODO: implement\n    return;",159            "    throw new Error('Not implemented');",160            "    var result = null;\n    return result;",161            "    console.log('Function called');\n    return;"162        ],163        'java': [164            "    return null;",165            "    // TODO: implement this method\n    return null;",166            "    throw new UnsupportedOperationException();",167            "    Object result = null;\n    return result;",168            "    System.out.println(\"Method called\");\n    return null;"169        ]170    }171 172    if language in templates:173        # Return multiple variations for Pass@k evaluation174        return templates[language]175    else:176        return ["// TODO: implement"]177 178def evaluate_completion_quality(predicted: str, actual: str) -> float:179    """Simple evaluation of completion quality."""180    # Normalize strings181    pred_lines = [line.strip() for line in predicted.split('\n') if line.strip()]182    actual_lines = [line.strip() for line in actual.split('\n') if line.strip()]183 184    if not actual_lines:185        return 0.0186 187    # Check for basic structural similarity188    score = 0.0189 190    # Check if both are empty implementations191    empty_indicators = ['pass', 'todo', 'not implemented', 'null', 'return;', 'return null'}192    pred_empty = any(indicator in predicted.lower() for indicator in empty_indicators)193    actual_empty = any(indicator in actual.lower() for indicator in empty_indicators)194 195    if pred_empty and actual_empty:196        score += 0.8197    elif not pred_empty and not actual_empty:198        # Check for keyword similarity199        pred_keywords = set(re.findall(r'\b\w+\b', predicted.lower()))200        actual_keywords = set(re.findall(r'\b\w+\b', actual.lower()))201 202        if actual_keywords:203            keyword_overlap = len(pred_keywords & actual_keywords) / len(actual_keywords)204            score += keyword_overlap * 0.6205 206        # Check for similar line count207        line_ratio = min(len(pred_lines), len(actual_lines)) / max(len(pred_lines), len(actual_lines))208        score += line_ratio * 0.4209 210    return min(score, 1.0)211 212def calculate_pass_at_k(completion_results: List[Tuple[List[str], str]], k: int = 1) -> float:213    """Calculate Pass@k metric."""214    if k <= 0:215        return 0.0216 217    total_passed = 0218 219    for completions, ground_truth in completion_results:220        # Take top k completions221        top_k_completions = completions[:k]222 223        # Check if any completion passes224        passed = False225        for completion in top_k_completions:226            quality_score = evaluate_completion_quality(completion, ground_truth)227            if quality_score > 0.5:  # Threshold for "passing"228                passed = True229                break230 231        if passed:232            total_passed += 1233 234    return total_passed / len(completion_results) if completion_results else 0.0235 236def run_completion_benchmark(repositories: List[Dict]) -> Dict:237    """Run code completion benchmark."""238    print("🧮 Running code completion benchmark...")239 240    results = {241        'total_repositories': len(repositories),242        'completion_tasks': [],243        'language_stats': defaultdict(int),244        'pass_at_1': 0.0,245        'pass_at_3': 0.0,246        'pass_at_5': 0.0,247        'average_quality': 0.0248    }249 250    completion_results = []251    quality_scores = []252 253    # Extract function snippets from repositories254    for repo in repositories:255        for language in ['python', 'javascript', 'java']:256            snippets = extract_function_snippets(repo, language)257 258            for snippet in snippets[:2]:  # Limit per repo for performance259                results['language_stats'][language] += 1260 261                # Generate completions262                completions = simple_code_completion_model(snippet['prompt'], language)263                ground_truth = snippet['completion']264 265                completion_results.append((completions, ground_truth))266 267                # Calculate quality for first completion268                if completions:269                    quality = evaluate_completion_quality(completions[0], ground_truth)270                    quality_scores.append(quality)271 272                results['completion_tasks'].append({273                    'function_name': snippet['function_name'],274                    'language': language,275                    'prompt_length': len(snippet['prompt']),276                    'completion_length': len(ground_truth)277                })278 279    # Calculate metrics280    results['pass_at_1'] = calculate_pass_at_k(completion_results, 1)281    results['pass_at_3'] = calculate_pass_at_k(completion_results, 3)282    results['pass_at_5'] = calculate_pass_at_k(completion_results, 5)283    results['average_quality'] = sum(quality_scores) / len(quality_scores) if quality_scores else 0.0284 285    return results286 287def print_benchmark_results(results: Dict):288    """Print formatted benchmark results."""289    print("=" * 60)290    print("🎯 CODE COMPLETION BENCHMARK RESULTS")291    print("=" * 60)292 293    print(f"Total repositories: {results['total_repositories']}")294    print(f"Completion tasks: {len(results['completion_tasks'])}")295 296    print(f"\n📊 Pass@k Metrics:")297    print(f"  Pass@1: {results['pass_at_1']:.3f}")298    print(f"  Pass@3: {results['pass_at_3']:.3f}")299    print(f"  Pass@5: {results['pass_at_5']:.3f}")300    print(f"  Average Quality: {results['average_quality']:.3f}")301 302    print(f"\n🔤 Language Distribution:")303    for language, count in sorted(results['language_stats'].items(), key=lambda x: x[1], reverse=True):304        print(f"  {language}: {count} functions")305 306    print(f"\n💡 Insights:")307    print("- This is a simplified demonstration benchmark")308    print("- Real evaluation requires more sophisticated code execution")309    print("- CodeReality-1T provides diverse, noisy code for robust testing")310    print("- Consider functional correctness testing for production models")311 312def main():313    """Run code completion benchmark."""314    print("🚀 CodeReality-1T Code Completion Benchmark")315    print("=" * 60)316 317    # Configuration318    data_dir = "/mnt/z/CodeReality_Final/unified_dataset"319    sample_size = 100320 321    if not os.path.exists(data_dir):322        print(f"❌ Dataset directory not found: {data_dir}")323        print("Please update the data_dir path to point to your CodeReality-1T dataset")324        return325 326    # Load dataset sample327    repositories = load_dataset_sample(data_dir, sample_size)328 329    if not repositories:330        print("❌ No repositories loaded. Check dataset path.")331        return332 333    # Run benchmark334    results = run_completion_benchmark(repositories)335 336    # Print results337    print_benchmark_results(results)338 339    # Save results340    output_file = "code_completion_results.json"341    with open(output_file, 'w') as f:342        # Convert defaultdict to regular dict for JSON serialization343        results_json = {344            'total_repositories': results['total_repositories'],345            'completion_tasks': results['completion_tasks'],346            'language_stats': dict(results['language_stats']),347            'pass_at_1': results['pass_at_1'],348            'pass_at_3': results['pass_at_3'],349            'pass_at_5': results['pass_at_5'],350            'average_quality': results['average_quality']351        }352        json.dump(results_json, f, indent=2)353 354    print(f"\n💾 Results saved to: {output_file}")355 356if __name__ == "__main__":357    main()