Upload data3/extract_functions_v2.py with huggingface_hub

data3/extract_functions_v2.py

@@ -0,0 +1,308 @@
+#!/usr/bin/env python3
+"""
+Extract individual functions from enhanced_dataset.csv and create a new dataset.
+Each function becomes a separate row in the new dataset.
+Version 2: Better handling of malformed CSV/JSON
+"""
+
+import csv
+import json
+import re
+from collections import defaultdict
+import sys
+
+def clean_json_string(json_str):
+    """
+    Clean up malformed JSON strings that may have been corrupted by CSV formatting.
+    """
+    # Remove extra spaces in key names that might have been inserted
+    # This is a bit risky but we'll try to handle common cases
+    
+    # Replace common malformed patterns
+    json_str = re.sub(r'"\s*function_nam\s*e\s*"', '"function_name"', json_str)
+    json_str = re.sub(r'"\s*function_start_line\s*"', '"function_start_line"', json_str)
+    json_str = re.sub(r'"\s*function_end_line\s*"', '"function_end_line"', json_str)
+    json_str = re.sub(r'"\s*relevance_score\s*"', '"relevance_score"', json_str)
+    json_str = re.sub(r'"\s*relevance_reason\s*"', '"relevance_reason"', json_str)
+    json_str = re.sub(r'"\s*doc_start_line\s*"', '"doc_start_line"', json_str)
+    json_str = re.sub(r'"\s*doc_end_line\s*"', '"doc_end_line"', json_str)
+    
+    # Remove markdown bold markers that might have been inserted
+    json_str = json_str.replace('**', '')
+    
+    # Fix spacing issues in keys
+    json_str = re.sub(r'"\s*([a-z_]+)\s*([a-z_]+)\s*([a-z_]*)\s*":', 
+                      lambda m: '"' + m.group(1) + m.group(2) + (m.group(3) if m.group(3) else '') + '":', 
+                      json_str)
+    
+    return json_str
+
+
+def extract_function_content(text, start_line, end_line):
+    """
+    Extract function content from text based on line number range.
+    
+    Args:
+        text: The full code text
+        start_line: Starting line number (1-indexed)
+        end_line: Ending line number (1-indexed)
+    
+    Returns:
+        Extracted function content as string
+    """
+    lines = text.split('\n')
+    # Convert to 0-indexed (since start_line is 1-indexed, we subtract 1)
+    start_idx = max(0, start_line - 1)
+    end_idx = min(len(lines), end_line)  # end_line is inclusive, so we don't subtract 1
+    
+    function_lines = lines[start_idx:end_idx]
+    return '\n'.join(function_lines)
+
+
+def process_dataset(input_file, output_file):
+    """
+    Process enhanced_dataset.csv and extract functions.
+    
+    Args:
+        input_file: Path to enhanced_dataset.csv
+        output_file: Path to output CSV file
+    """
+    print(f"Reading from: {input_file}")
+    print(f"Writing to: {output_file}")
+    
+    # Statistics
+    total_rows = 0
+    total_functions = 0
+    score_distribution = defaultdict(int)
+    skipped_rows = 0
+    parse_errors = 0
+    empty_function_info = 0
+    
+    with open(input_file, 'r', encoding='utf-8') as infile, \
+         open(output_file, 'w', encoding='utf-8', newline='') as outfile:
+        
+        reader = csv.DictReader(infile)
+        
+        # Define output columns
+        fieldnames = [
+            'original_index',      # Original row number
+            'function_index',      # Index within the file
+            'repo_name',
+            'path',
+            'language',
+            'license',
+            'keyword',
+            'text_hash',
+            'config',
+            'split',
+            'repo_path',
+            'ds_source',
+            'function_name',
+            'function_start_line',
+            'function_end_line',
+            'doc_start_line',
+            'doc_end_line',
+            'relevance_score',
+            'relevance_reason',
+            'function_content'
+        ]
+        
+        writer = csv.DictWriter(outfile, fieldnames=fieldnames)
+        writer.writeheader()
+        
+        # Store all function rows for later sorting
+        all_function_rows = []
+        
+        print("\nProcessing rows...")
+        for row in reader:
+            total_rows += 1
+            
+            if total_rows % 1000 == 0:
+                print(f"Processed {total_rows} rows, extracted {total_functions} functions, errors: {parse_errors}...", end='\r')
+            
+            # Parse function_info JSON
+            function_info_str = row.get('function_info', '[]')
+            if not function_info_str or function_info_str.strip() == '':
+                empty_function_info += 1
+                skipped_rows += 1
+                continue
+            
+            # Clean the JSON string
+            function_info_str = clean_json_string(function_info_str)
+            
+            # Handle potential CSV escaping issues
+            try:
+                # First try direct JSON parsing
+                function_info_list = json.loads(function_info_str)
+            except (json.JSONDecodeError, ValueError) as e:
+                # If that fails, try with ast.literal_eval as backup
+                try:
+                    import ast
+                    function_info_list = ast.literal_eval(function_info_str)
+                except:
+                    # If still fails, skip this row
+                    parse_errors += 1
+                    if parse_errors <= 5:  # Only print first 5 errors
+                        print(f"\nWarning: Failed to parse function_info in row {total_rows}")
+                        if parse_errors == 5:
+                            print("(Suppressing further parse error messages...)")
+                    skipped_rows += 1
+                    continue
+            
+            # Validate that we got a list
+            if not isinstance(function_info_list, list):
+                skipped_rows += 1
+                continue
+            
+            # Get the original text
+            text = row.get('text', '')
+            
+            # Extract each function
+            for func_idx, func_info in enumerate(function_info_list):
+                # Validate func_info is a dictionary
+                if not isinstance(func_info, dict):
+                    continue
+                
+                # Extract function content
+                start_line = func_info.get('function_start_line', 0)
+                end_line = func_info.get('function_end_line', 0)
+                
+                # Ensure they are integers
+                try:
+                    start_line = int(start_line) if start_line else 0
+                    end_line = int(end_line) if end_line else 0
+                except (ValueError, TypeError):
+                    start_line = 0
+                    end_line = 0
+                
+                if start_line > 0 and end_line > 0:
+                    function_content = extract_function_content(text, start_line, end_line)
+                else:
+                    function_content = ""
+                
+                # Get relevance score
+                relevance_score = func_info.get('relevance_score', 0)
+                
+                # Ensure it's an integer
+                try:
+                    relevance_score = int(relevance_score) if relevance_score else 0
+                except (ValueError, TypeError):
+                    relevance_score = 0
+                
+                # Track score distribution (in buckets of 10)
+                score_bucket = (relevance_score // 10) * 10
+                score_distribution[score_bucket] += 1
+                
+                # Create new row
+                new_row = {
+                    'original_index': row.get('Unnamed: 0', row.get('Unnamed: 0.1', total_rows - 1)),
+                    'function_index': func_idx,
+                    'repo_name': row.get('repo_name', ''),
+                    'path': row.get('path', ''),
+                    'language': row.get('language', ''),
+                    'license': row.get('license', ''),
+                    'keyword': row.get('keyword', ''),
+                    'text_hash': row.get('text_hash', ''),
+                    'config': row.get('config', ''),
+                    'split': row.get('split', ''),
+                    'repo_path': row.get('repo_path', ''),
+                    'ds_source': row.get('ds_source', ''),
+                    'function_name': func_info.get('function_name', ''),
+                    'function_start_line': start_line,
+                    'function_end_line': end_line,
+                    'doc_start_line': func_info.get('doc_start_line', ''),
+                    'doc_end_line': func_info.get('doc_end_line', ''),
+                    'relevance_score': relevance_score,
+                    'relevance_reason': func_info.get('relevance_reason', ''),
+                    'function_content': function_content
+                }
+                
+                all_function_rows.append(new_row)
+                total_functions += 1
+        
+        print(f"\n\nTotal rows processed: {total_rows}")
+        print(f"Total functions extracted: {total_functions}")
+        print(f"Skipped rows:")
+        print(f"  - Empty function_info: {empty_function_info}")
+        print(f"  - Parse errors: {parse_errors}")
+        print(f"  - Total skipped: {skipped_rows}")
+        
+        # Sort by relevance_score (descending - highest first)
+        print("\nSorting by relevance score...")
+        all_function_rows.sort(key=lambda x: x['relevance_score'], reverse=True)
+        
+        # Write sorted rows
+        print("Writing sorted data to output file...")
+        for row in all_function_rows:
+            writer.writerow(row)
+        
+        print(f"\nSuccessfully written {total_functions} functions to {output_file}")
+    
+    # Print score distribution
+    print("\n" + "="*70)
+    print("SCORE DISTRIBUTION")
+    print("="*70)
+    print(f"{'Score Range':<15} {'Count':<12} {'Percentage':<12} {'Visualization'}")
+    print("-"*70)
+    
+    # Sort by score range (descending)
+    sorted_scores = sorted(score_distribution.items(), reverse=True)
+    
+    # Filter out anomalous scores (very negative values)
+    normal_scores = [(k, v) for k, v in sorted_scores if k >= 0]
+    anomalous_scores = [(k, v) for k, v in sorted_scores if k < 0]
+    
+    for score_bucket, count in normal_scores:
+        percentage = (count / total_functions * 100) if total_functions > 0 else 0
+        bar = '█' * min(50, int(percentage / 2))  # Scale bar to fit
+        print(f"{score_bucket:>3}-{score_bucket+9:<9} {count:<12} {percentage:>6.2f}%      {bar}")
+    
+    if anomalous_scores:
+        print("\nAnomalous scores (negative or out of range):")
+        for score_bucket, count in anomalous_scores:
+            percentage = (count / total_functions * 100) if total_functions > 0 else 0
+            print(f"{score_bucket:>15} {count:<12} {percentage:>6.2f}%")
+    
+    print("-"*70)
+    print(f"{'Total':<15} {total_functions:<12} {'100.00%':<12}")
+    print("="*70)
+    
+    # Additional statistics
+    if total_functions > 0:
+        # Filter out anomalous scores for statistics
+        valid_scores = [row['relevance_score'] for row in all_function_rows 
+                       if 0 <= row['relevance_score'] <= 100]
+        
+        if valid_scores:
+            avg_score = sum(valid_scores) / len(valid_scores)
+            max_score = max(valid_scores)
+            min_score = min(valid_scores)
+            
+            print(f"\nScore Statistics (valid scores 0-100 only):")
+            print(f"  Average Score: {avg_score:.2f}")
+            print(f"  Maximum Score: {max_score}")
+            print(f"  Minimum Score: {min_score}")
+            print(f"  Valid Functions: {len(valid_scores)} / {total_functions}")
+
+
+if __name__ == "__main__":
+    input_file = "enhanced_dataset.csv"
+    output_file = "function_dataset_v2.csv"
+    
+    # Allow command line arguments
+    if len(sys.argv) > 1:
+        input_file = sys.argv[1]
+    if len(sys.argv) > 2:
+        output_file = sys.argv[2]
+    
+    try:
+        process_dataset(input_file, output_file)
+        print("\n✅ Processing complete!")
+    except FileNotFoundError:
+        print(f"❌ Error: File '{input_file}' not found.")
+        sys.exit(1)
+    except Exception as e:
+        print(f"❌ Error: {e}")
+        import traceback
+        traceback.print_exc()
+        sys.exit(1)