Upload patchjudge/data_loader.py with huggingface_hub

patchjudge/data_loader.py

@@ -0,0 +1,380 @@
+"""Data loading pipeline for PatchJudge.
+
+Loads SWE-bench Verified gold patches and agent-generated patches from:
+1. HuggingFace datasets (AlexCuadron O1, CoderForge)
+2. SWE-bench S3 bucket (139 verified agent submissions)
+"""
+
+import json
+import os
+import re
+import logging
+from pathlib import Path
+from typing import Optional
+from collections import defaultdict
+
+from datasets import load_dataset
+from patchjudge.models import PatchExample
+
+logger = logging.getLogger(__name__)
+
+
+class SWEBenchLoader:
+    """Loads SWE-bench Verified data and agent patches."""
+    
+    def __init__(self, cache_dir: str = "data"):
+        self.cache_dir = Path(cache_dir)
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+        self._gold_data = None  # Lazy loaded
+    
+    def load_gold_data(self) -> dict:
+        """Load SWE-bench Verified dataset. Returns {instance_id: row_dict}."""
+        if self._gold_data is not None:
+            return self._gold_data
+        
+        logger.info("Loading SWE-bench Verified dataset...")
+        ds = load_dataset("princeton-nlp/SWE-bench_Verified", split="test")
+        self._gold_data = {}
+        for row in ds:
+            self._gold_data[row["instance_id"]] = {
+                "instance_id": row["instance_id"],
+                "repo": row["repo"],
+                "problem_statement": row["problem_statement"],
+                "gold_patch": row["patch"],
+                "base_commit": row["base_commit"],
+                "test_patch": row["test_patch"],
+                "difficulty": row.get("difficulty", ""),
+                "hints_text": row.get("hints_text", ""),
+            }
+        logger.info(f"Loaded {len(self._gold_data)} SWE-bench Verified instances")
+        return self._gold_data
+    
+    def load_coderforge_patches(self) -> list[PatchExample]:
+        """Load agent patches from CoderForge (Qwen3-Coder-32B, 500 instances)."""
+        logger.info("Loading CoderForge agent patches...")
+        ds = load_dataset(
+            "togethercomputer/CoderForge-Preview-32B-SWE-Bench-Verified-Evaluation-trajectories",
+            "trajectory", split="train"
+        )
+        gold = self.load_gold_data()
+        examples = []
+        
+        for row in ds:
+            # Extract instance_id from ds JSON field
+            try:
+                ds_info = json.loads(row["ds"])
+                instance_id = ds_info["instance_id"]
+            except (json.JSONDecodeError, KeyError):
+                # Try extracting from trajectory_id
+                tid = row.get("trajectory_id", "")
+                instance_id = tid.rsplit("_run", 1)[0] if "_run" in tid else tid
+            
+            if instance_id not in gold:
+                continue
+            
+            agent_patch = row.get("output_patch", "")
+            if not agent_patch or agent_patch.strip() == "":
+                continue
+            
+            g = gold[instance_id]
+            ex = PatchExample(
+                instance_id=instance_id,
+                repo=g["repo"],
+                problem_statement=g["problem_statement"],
+                gold_patch=g["gold_patch"],
+                agent_patch=agent_patch,
+                agent_name="CoderForge-Qwen3-32B",
+                test_passed=row.get("reward", 0.0) == 1.0,
+                base_commit=g["base_commit"],
+                difficulty=g["difficulty"],
+            )
+            examples.append(ex)
+        
+        logger.info(f"Loaded {len(examples)} CoderForge patches "
+                     f"({sum(1 for e in examples if e.test_passed)} passed)")
+        return examples
+    
+    def load_o1_patches(self) -> list[PatchExample]:
+        """Load agent patches from OpenHands+O1 (500 instances)."""
+        logger.info("Loading OpenHands+O1 agent patches...")
+        ds = load_dataset(
+            "AlexCuadron/SWE-Bench-Verified-O1-native-tool-calling-reasoning-high-results",
+            split="test"
+        )
+        gold = self.load_gold_data()
+        examples = []
+        
+        for row in ds:
+            issue_name = row.get("issue_name", "")
+            # issue_name format: "django__django-16454" — same as instance_id
+            instance_id = issue_name
+            
+            if instance_id not in gold:
+                continue
+            
+            agent_patch = row.get("patch", "")
+            if not agent_patch or agent_patch.strip() == "":
+                continue
+            
+            g = gold[instance_id]
+            ex = PatchExample(
+                instance_id=instance_id,
+                repo=g["repo"],
+                problem_statement=g["problem_statement"],
+                gold_patch=g["gold_patch"],
+                agent_patch=agent_patch,
+                agent_name="OpenHands-O1-reasoning-high",
+                test_passed=row.get("resolved", False),
+                base_commit=g["base_commit"],
+                difficulty=g["difficulty"],
+            )
+            examples.append(ex)
+        
+        logger.info(f"Loaded {len(examples)} O1 patches "
+                     f"({sum(1 for e in examples if e.test_passed)} passed)")
+        return examples
+    
+    def load_s3_agent_patches(
+        self,
+        agents: list[str] = None,
+        max_per_agent: int = 500,
+    ) -> list[PatchExample]:
+        """Load agent patches from SWE-bench S3 bucket.
+        
+        Args:
+            agents: List of agent directory names in S3. Defaults to a curated set.
+            max_per_agent: Max patches per agent.
+        """
+        try:
+            import boto3
+            from botocore import UNSIGNED
+            from botocore.config import Config
+            import requests
+        except ImportError:
+            logger.warning("boto3 not available, skipping S3 patches")
+            return []
+        
+        if agents is None:
+            agents = [
+                "20250225_sweagent_claude-3-7-sonnet",
+                "20241029_OpenHands-CodeAct-2.1-sonnet-20241022",
+                "20241028_agentless-1.5_gpt4o",
+                "20241108_autocoderover-v2.0-claude-3-5-sonnet-20241022",
+                "20240620_sweagent_claude3.5sonnet",
+            ]
+        
+        gold = self.load_gold_data()
+        s3 = boto3.client('s3', config=Config(signature_version=UNSIGNED))
+        BUCKET = 'swe-bench-submissions'
+        examples = []
+        
+        for agent_dir in agents:
+            logger.info(f"Loading patches from S3 agent: {agent_dir}")
+            
+            # Get resolve labels from GitHub
+            resolved_ids = set()
+            try:
+                url = (
+                    f"https://raw.githubusercontent.com/SWE-bench/experiments/"
+                    f"main/evaluation/verified/{agent_dir}/results/results.json"
+                )
+                import requests
+                r = requests.get(url, timeout=10)
+                if r.status_code == 200:
+                    resolved_ids = set(r.json().get("resolved", []))
+                    logger.info(f"  {agent_dir}: {len(resolved_ids)} resolved")
+            except Exception as e:
+                logger.warning(f"  Could not load resolve labels for {agent_dir}: {e}")
+            
+            # List instance directories
+            paginator = s3.get_paginator('list_objects_v2')
+            count = 0
+            try:
+                for page in paginator.paginate(
+                    Bucket=BUCKET,
+                    Prefix=f'verified/{agent_dir}/logs/',
+                    Delimiter='/'
+                ):
+                    for prefix_info in page.get('CommonPrefixes', []):
+                        if count >= max_per_agent:
+                            break
+                        
+                        prefix = prefix_info['Prefix']
+                        instance_id = prefix.rstrip('/').split('/')[-1]
+                        
+                        if instance_id not in gold:
+                            continue
+                        
+                        # Download patch.diff
+                        try:
+                            obj = s3.get_object(
+                                Bucket=BUCKET,
+                                Key=f'verified/{agent_dir}/logs/{instance_id}/patch.diff'
+                            )
+                            agent_patch = obj['Body'].read().decode('utf-8')
+                        except Exception:
+                            continue
+                        
+                        if not agent_patch.strip():
+                            continue
+                        
+                        g = gold[instance_id]
+                        ex = PatchExample(
+                            instance_id=instance_id,
+                            repo=g["repo"],
+                            problem_statement=g["problem_statement"],
+                            gold_patch=g["gold_patch"],
+                            agent_patch=agent_patch,
+                            agent_name=agent_dir,
+                            test_passed=instance_id in resolved_ids,
+                            base_commit=g["base_commit"],
+                            difficulty=g["difficulty"],
+                        )
+                        examples.append(ex)
+                        count += 1
+            except Exception as e:
+                logger.warning(f"  Error loading from S3 for {agent_dir}: {e}")
+            
+            logger.info(f"  Loaded {count} patches from {agent_dir}")
+        
+        logger.info(f"Total S3 patches: {len(examples)} "
+                     f"({sum(1 for e in examples if e.test_passed)} passed)")
+        return examples
+    
+    def build_dataset(
+        self,
+        sources: list[str] = None,
+        min_examples: int = 100,
+        include_repo_context: bool = False,
+        s3_agents: list[str] = None,
+    ) -> list[PatchExample]:
+        """Build the unified PatchExample dataset from multiple sources.
+        
+        Args:
+            sources: List of sources to use. Options: 'coderforge', 'o1', 's3'.
+                     Defaults to ['coderforge', 'o1'].
+            min_examples: Minimum examples to collect.
+            include_repo_context: If True, attempt to clone repos and gather context.
+            s3_agents: Agent list for S3 source.
+        """
+        if sources is None:
+            sources = ["coderforge", "o1"]
+        
+        all_examples = []
+        
+        if "coderforge" in sources:
+            all_examples.extend(self.load_coderforge_patches())
+        
+        if "o1" in sources:
+            all_examples.extend(self.load_o1_patches())
+        
+        if "s3" in sources:
+            all_examples.extend(self.load_s3_agent_patches(agents=s3_agents))
+        
+        # Deduplicate by (instance_id, agent_name)
+        seen = set()
+        unique = []
+        for ex in all_examples:
+            key = (ex.instance_id, ex.agent_name)
+            if key not in seen:
+                seen.add(key)
+                unique.append(ex)
+        
+        logger.info(f"Total unique examples: {len(unique)} "
+                     f"(passed: {sum(1 for e in unique if e.test_passed)}, "
+                     f"failed: {sum(1 for e in unique if not e.test_passed)})")
+        
+        if len(unique) < min_examples:
+            logger.warning(
+                f"Only {len(unique)} examples collected, "
+                f"below minimum of {min_examples}. "
+                f"Consider adding more sources."
+            )
+        
+        return unique
+    
+    def save_dataset(self, examples: list[PatchExample], filename: str = "patch_examples.jsonl"):
+        """Save examples to JSONL."""
+        path = self.cache_dir / filename
+        with open(path, 'w') as f:
+            for ex in examples:
+                f.write(json.dumps(ex.to_dict()) + "\n")
+        logger.info(f"Saved {len(examples)} examples to {path}")
+        return path
+    
+    def load_saved_dataset(self, filename: str = "patch_examples.jsonl") -> list[PatchExample]:
+        """Load previously saved examples."""
+        path = self.cache_dir / filename
+        examples = []
+        with open(path) as f:
+            for line in f:
+                if line.strip():
+                    examples.append(PatchExample.from_dict(json.loads(line)))
+        logger.info(f"Loaded {len(examples)} examples from {path}")
+        return examples
+
+
+def extract_repo_context_from_diff(diff: str) -> list[str]:
+    """Extract filenames mentioned in a diff."""
+    files = []
+    for line in diff.split('\n'):
+        if line.startswith('diff --git'):
+            # Extract b/path
+            match = re.search(r'b/(.+)$', line)
+            if match:
+                files.append(match.group(1))
+        elif line.startswith('---') and not line.startswith('--- /dev/null'):
+            match = re.search(r'a/(.+)$', line)
+            if match:
+                files.append(match.group(1))
+    return list(set(files))
+
+
+def get_diff_stats(diff: str) -> dict:
+    """Get basic stats from a unified diff."""
+    lines = diff.split('\n')
+    added = sum(1 for l in lines if l.startswith('+') and not l.startswith('+++'))
+    removed = sum(1 for l in lines if l.startswith('-') and not l.startswith('---'))
+    files = len(extract_repo_context_from_diff(diff))
+    hunks = sum(1 for l in lines if l.startswith('@@'))
+    return {
+        "lines_added": added,
+        "lines_removed": removed,
+        "files_changed": files,
+        "hunks": hunks,
+    }
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+    loader = SWEBenchLoader()
+    
+    # Load from HF datasets (no S3 dependency)
+    examples = loader.build_dataset(sources=["coderforge", "o1"])
+    
+    # Stats
+    passed = sum(1 for e in examples if e.test_passed)
+    failed = len(examples) - passed
+    repos = set(e.repo for e in examples)
+    agents = set(e.agent_name for e in examples)
+    
+    print(f"\n{'='*60}")
+    print(f"PatchJudge Dataset Summary")
+    print(f"{'='*60}")
+    print(f"Total examples:      {len(examples)}")
+    print(f"  Test passed:       {passed}")
+    print(f"  Test failed:       {failed}")
+    print(f"Unique instances:    {len(set(e.instance_id for e in examples))}")
+    print(f"Unique repos:        {len(repos)}")
+    print(f"Agent sources:       {agents}")
+    print(f"\nDifficulty distribution:")
+    
+    diff_counts = defaultdict(int)
+    for e in examples:
+        diff_counts[e.difficulty] += 1
+    for d, c in sorted(diff_counts.items()):
+        print(f"  {d}: {c}")
+    
+    # Save
+    path = loader.save_dataset(examples)
+    print(f"\nSaved to: {path}")