SWE-Bench: Can Language Models Resolve Real-world Github Issues?
import argparse
import json
import logging
import os
from typing import Dict, Optional


from utils import Repo, extract_patches, extract_problem_statement_and_hints

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


def create_instance(repo: Repo, pull: Dict) -> Dict:
    """
    Create a single task instance from a pull request, where task instance is:

    {
        repo (str): owner/repo this task instance is from,
        pull_number (int): number of PR this task instance is from,
        base_commit (str): SHA of the base commit PR is based on,
        patch (str): reference solution as .patch (apply to base commit),
        test_patch (str): test suite as .patch (apply to base commit),
    }
    """
    patch, test_patch = extract_patches(pull, repo)
    problem_statement, hints = extract_problem_statement_and_hints(pull, repo)
    return {
        "repo": repo.repo.full_name,
        "pull_number": pull["number"],
        "instance_id": (repo.repo.full_name + "-" + str(pull["number"])).replace(
            "/", "__"
        ),
        "issue_numbers": pull["resolved_issues"],
        "base_commit": pull["base"]["sha"],
        "patch": patch,
        "test_patch": test_patch,
        "problem_statement": problem_statement,
        "hints_text": hints,
        "created_at": pull["created_at"],
    }


def is_valid_pull(pull: Dict) -> bool:
    """
    Check whether PR has an associated issue and is merged

    Args:
        pull (dict): pull request object
    Returns:
        bool: whether PR is valid
    """
    if pull["merged_at"] is None:
        return False
    if "resolved_issues" not in pull or len(pull["resolved_issues"]) < 1:
        return False
    return True


def is_valid_instance(instance: Dict) -> bool:
    """
    Check whether task instance has all required fields for task instance creation

    Args:
        instance (dict): task instance object
    Returns:
        bool: whether task instance is valid
    """
    if instance["patch"] is None or instance["patch"] == "":
        return False
    if instance["problem_statement"] is None or instance["problem_statement"] == "":
        return False
    return True


def has_test_patch(instance: Dict) -> bool:
    """
    Check whether task instance has a test suite

    Args:
        instance (dict): task instance object
    Returns:
        bool: whether task instance has a test suite
    """
    if instance["test_patch"] is None or instance["test_patch"].strip() == "":
        return False
    return True


def main(pr_file: str, output: str, token: Optional[str] = None):
    """
    Main thread for creating task instances from pull requests

    Args:
        pr_file (str): path to pull request JSONL file
        output (str): output file name
        token (str): GitHub token
    """
    if token is None:
        # Get GitHub token from environment variable if not provided
        token = os.environ["GITHUB_TOKEN"]

    def load_repo(repo_name):
        # Return repo object for a given repo name
        owner, repo = repo_name.split("/")
        return Repo(owner, repo, token=token)

    repos = dict()
    completed = 0
    with_tests = 0
    total_instances = 0
    all_output = output + ".all"
    seen_prs = set()

    # Continue where we left off if output file already exists
    if os.path.exists(all_output):
        with open(all_output) as f:
            for line in f:
                pr = json.loads(line)
                if "instance_id" not in pr:
                    pr["instance_id"] = (
                        pr["repo"] + "-" + str(pr["pull_number"])
                    ).replace("/", "__")
                instance_id = pr["instance_id"]
                seen_prs.add(instance_id)
                if is_valid_instance(pr):
                    completed += 1
                    if has_test_patch(pr):
                        with_tests += 1
    logger.info(f"{len(seen_prs)} instance_ids previously recorded")

    # Write to .all file for all PRs
    with open(all_output, "a") as all_output:
        # Write to output file for PRs with test suites
        with open(output, "w") as output:
            for ix, line in enumerate(open(pr_file)):
                total_instances += 1
                pull = json.loads(line)
                if ix % 100 == 0:
                    logger.info(
                        f"[{pull['base']['repo']['full_name']}] ( {ix} / {total_instances} ) {completed} valid, {with_tests} with tests."
                    )
                # Construct instance fields
                instance_id = (
                    pull["base"]["repo"]["full_name"] + "-" + str(pull["number"])
                )
                instance_id = instance_id.replace("/", "__")
                if instance_id in seen_prs:
                    seen_prs -= {instance_id}
                    continue
                if not is_valid_pull(pull):
                    # Throw out invalid PRs
                    continue
                # Create task instance
                repo_name = pull["base"]["repo"]["full_name"]
                if repo_name not in repos:
                    repos[repo_name] = load_repo(repo_name)
                repo = repos[repo_name]
                instance = create_instance(repo, pull)
                if is_valid_instance(instance):
                    # If valid, write to .all output file
                    print(
                        json.dumps(instance), end="\n", flush=True, file=all_output
                    )  # write all instances to a separate file
                    completed += 1
                    if has_test_patch(instance):
                        # If has test suite, write to output file
                        print(json.dumps(instance), end="\n", flush=True, file=output)
                        with_tests += 1
    logger.info(
        f"Total instances: {total_instances}, completed: {completed}, with tests: {with_tests}"
    )
    logger.info(f"Didn't see {len(seen_prs)} instances previously recorded")
    logger.info("\n".join(sorted(seen_prs)))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("pr_file", type=str, help="Path to pull request JSONL file")
    parser.add_argument("output", type=str, help="Output file name")
    parser.add_argument("--token", type=str, help="GitHub token")
    args = parser.parse_args()
    main(**vars(args))

collect/build_dataset.py
princeton-nlp/SWE-bench
import argparse
import glob
import json
import os
import random

from tqdm import tqdm
from datetime import datetime


def main(instances_path: str, output_path: str, eval_path: str, seed: int):
    """
    Combine all non-eval task instances into a single fine tuning dataset

    Args:
        instances_path (str): Path to directory containing all candidate task instances
        output_path (str): Path to save output fine tuning dataset to
        eval_path (str): Path to directory containing all eval task instances
        seed (int): Random seed
    """
    # Define output file name
    random.seed(seed)
    SWE_PRS_FT_DATASET = (
        f"SWE_PRS_FT_DATASET_{datetime.now().strftime('%Y%m%d%H')}_{seed}.jsonl"
    )
    destination = os.path.join(output_path, SWE_PRS_FT_DATASET)
    total_insts, total_repos = 0, 0

    # Gather Evaluation Set Task Instances
    eval_instances = []
    for x in glob.glob(os.path.join(eval_path, "*-task-instances.jsonl")):
        with open(x, "r") as f:
            eval_instances.extend(f.readlines())
    eval_instances = set(eval_instances)

    # Create fine tuning dataset
    with open(destination, "w") as f_out:
        for dataset_path in tqdm(
            glob.glob(os.path.join(instances_path, "*-task-instances.jsonl.all"))
        ):
            total_repos += 1
            with open(dataset_path, "r") as f:
                lines = f.readlines()

                # Remove data from evaluation dataset
                lines = [line for line in lines if line not in eval_instances]

                # Shuffle lines
                random.shuffle(lines)

                # Keep 500 lines per dataset
                for line in lines[:500]:
                    line = json.loads(line)
                    if "test_patch" in line:
                        del line["test_patch"]
                    f_out.write(json.dumps(line) + "\n")
                    total_insts += 1

    print(
        f"Fine tuning dataset saved to {destination} ({total_insts} instances from {total_repos} repos)"
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--instances_path",
        type=str,
        help="Path to directory containing all candidate task instances",
    )
    parser.add_argument(
        "--output_path", type=str, help="Path to save output fine tuning dataset to"
    )
    parser.add_argument(
        "--eval_path",
        type=str,
        help="Path to directory containing all eval task instances",
    )
    parser.add_argument("--seed", type=int, default=42, help="Random seed")
    args = parser.parse_args()
    main(**vars(args))

collect/build_dataset_ft.py
princeton-nlp/SWE-bench
import argparse
import os
import subprocess


def main(repo_url):
    """
    Remove .github/workflows folder from all branches of a repo

    Args:
        repo_url (str): URL of the target repo
    """
    # Get list of remote branches
    branches_command = subprocess.run(
        ["git", "ls-remote", "--heads", repo_url], capture_output=True, text=True
    )
    branches = branches_command.stdout.strip().split("\n")
    branches = [branch.split()[1] for branch in branches]
    subprocess.run(
        ["git", "clone", repo_url, "temp_repo"],
        stderr=subprocess.DEVNULL,
        stdout=subprocess.DEVNULL,
    )

    # Iterate through all branches
    os.chdir("temp_repo")
    for branch in branches:
        # Switch to branch
        print(f"--------------\nProcessing branch: {branch}")
        branch = branch.split("/")[-1]
        subprocess.run(["git", "checkout", branch])

        workflows_path = os.path.join(".github", "workflows")
        if os.path.exists(workflows_path):
            # Remove .github/workflows folder if it exists
            print(f"Deleting .github/workflows folder from branch: {branch}")
            subprocess.run(["rm", "-rf", workflows_path])
            subprocess.run(["git", "add", "-A"])
            subprocess.run(["git", "commit", "-m", "Remove .github/workflows folder"])
            subprocess.run(["git", "push"])
        else:
            print(f".github/workflows folder not found in branch: {branch}")

    os.chdir("..")
    subprocess.run(["rm", "-rf", "temp_repo"])


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--repo_url", type=str, required=True)
    args = parser.parse_args()
    main(**vars(args))

collect/cleanup/delete_gh_workflows.py
princeton-nlp/SWE-bench
import argparse
import os
import subprocess

from multiprocessing import Pool


def get_conda_env_names(output: str) -> list:
    """
    Parse conda environments (`conda env list`) created for a particular conda installation

    Args:
        output (str): Output of `conda env list` command
    """
    lines = output.split("\n")
    env_names = []
    for line in lines:
        if line.startswith("#"):
            continue
        if line.strip() == "":
            continue
        if " " in line:
            env_name = line.split(" ")[0]
            env_names.append(env_name)
    return [x for x in env_names if len(x) > 0]


def delete_folders_with_prefix(prefix, conda_path):
    """
    Find and rm folders with a particular prefix in the conda installation's env folder

    Args:
        prefix (str): Prefix of folders to remove
        conda_path (str): Path to conda installation
    """
    envs_folder = os.path.join(conda_path, "envs")
    command = f'find {envs_folder} -type d -name "{prefix}*" -exec rm -rf {{}} +'
    subprocess.run(command, shell=True)


def remove_environment(env_name, prefix):
    """
    Remove all conda environments with a particular prefix from a conda installation
    """
    if env_name.startswith(prefix):
        print(f"Removing {env_name}")
        conda_cmd = "conda remove -n " + env_name + " --all -y"
        cmd = conda_source + " && " + conda_cmd
        conda_create_output = subprocess.run(cmd, shell=True, check=True)
        print(f"Output: {conda_create_output}")


if __name__ == "__main__":
    """
    Logic for removing conda environments and their folders from a conda installation
    """
    parser = argparse.ArgumentParser()
    parser.add_argument("prefix", type=str, help="Prefix for environments to delete")
    parser.add_argument(
        "--conda_path",
        type=str,
        help="Path to miniconda installation",
    )
    args = parser.parse_args()

    # Remove conda environments with a specific prefix
    conda_source = "source " + os.path.join(args.conda_path, "etc/profile.d/conda.sh")
    check_env = conda_source + " && " + "conda env list"
    conda_envs = subprocess.run(check_env, shell=True, check=True, capture_output=True)
    conda_envs_names = get_conda_env_names(conda_envs.stdout.decode("utf-8"))

    # Remove conda environments in parallel
    num_processes = 25
    pool = Pool(num_processes)
    pool.starmap(
        remove_environment, zip(conda_envs_names, [args.prefix] * len(conda_envs_names))
    )

    # Remove env folder with the same prefix
    print(
        f"Removing miniconda folder for environments with {args.prefix} from {args.conda_path}"
    )
    delete_folders_with_prefix(args.prefix, args.conda_path)
    print(f"Done!")

collect/cleanup/remove_envs.py
princeton-nlp/SWE-bench
import argparse, os

from dotenv import load_dotenv
from build_dataset import main as build_dataset
from print_pulls import main as print_pulls
from multiprocessing import Pool
from typing import Dict, List


load_dotenv()


def split_instances(input_list: List, n: int) -> List:
    """
    Split a list into n approximately equal length sublists

    Args:
        input_list (list): List to split
        n (int): Number of sublists to split into
    Returns:
        result (list): List of sublists
    """
    avg_length = len(input_list) // n
    remainder = len(input_list) % n
    result, start = [], 0

    for i in range(n):
        length = avg_length + 1 if i < remainder else avg_length
        sublist = input_list[start : start + length]
        result.append(sublist)
        start += length

    return result


def construct_data_files(data: Dict):
    """
    Logic for combining multiple .all PR files into a single fine tuning dataset

    Args:
        data (dict): Dictionary containing the following keys:
            repos (list): List of repositories to retrieve instruction data for
            path_prs (str): Path to save PR data files to
            path_tasks (str): Path to save task instance data files to
            token (str): GitHub token to use for API requests
    """
    repos, path_prs, path_tasks, token = (
        data["repos"],
        data["path_prs"],
        data["path_tasks"],
        data["token"],
    )
    for repo in repos:
        repo = repo.strip(",").strip()
        repo_name = repo.split("/")[1]
        try:
            path_pr = os.path.join(path_prs, f"{repo_name}-prs.jsonl")
            if not os.path.exists(path_pr):
                print(f"Pull request data for {repo} not found, creating...")
                print_pulls(repo, path_pr, token)
                print(f"Successfully saved PR data for {repo} to {path_pr}")
            else:
                print(
                    f"Pull request data for {repo} already exists at {path_pr}, skipping..."
                )

            path_task = os.path.join(path_tasks, f"{repo_name}-task-instances.jsonl")
            if not os.path.exists(path_task):
                print(f"Task instance data for {repo} not found, creating...")
                build_dataset(path_pr, path_task, token)
                print(
                    f"Successfully saved task instance data for {repo} to {path_task}"
                )
            else:
                print(
                    f"Task instance data for {repo} already exists at {path_task}, skipping..."
                )
        except Exception as e:
            print(f"Something went wrong for {repo}, skipping: {e}")
            pass


def main(repos, path_prs, path_tasks):
    """
    Spawns multiple threads given multiple GitHub tokens for collecting fine tuning data

    Args:
        repos (list): List of repositories to retrieve instruction data for
        path_prs (str): Path to save PR data files to
        path_tasks (str): Path to save task instance data files to
    """
    path_prs, path_tasks = os.path.abspath(path_prs), os.path.abspath(path_tasks)
    print(f"Will save PR data to {path_prs}")
    print(f"Will save task instance data to {path_tasks}")
    print(f"Received following repos to create task instances for: {repos}")

    tokens = os.getenv("GITHUB_TOKENS").split(",")
    data_task_lists = split_instances(repos, len(tokens))

    data_pooled = [
        {"repos": repos, "path_prs": path_prs, "path_tasks": path_tasks, "token": token}
        for repos, token in zip(data_task_lists, tokens)
    ]

    with Pool(len(tokens)) as p:
        p.map(construct_data_files, data_pooled)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--repos", nargs="+", help="List of repositories to create task instances for"
    )
    parser.add_argument("--path_prs", type=str, help="Path to save PR data files to")
    parser.add_argument(
        "--path_tasks", type=str, help="Path to save task instance data files to"
    )
    args = parser.parse_args()
    main(**vars(args))

collect/get_tasks_pipeline.py
princeton-nlp/SWE-bench
import os, json

from bs4 import BeautifulSoup
from ghapi.core import GhApi
from selenium import webdriver
from selenium.webdriver.common.by import By


api = GhApi(token="GITHUB TOKEN HERE")


def get_package_stats(data_tasks, f):
    """
    Get package stats from pypi page

    Args:
        data_tasks (list): List of packages + HTML
        f (str): File to write to
    """
    # Adjust access type if file already exists
    content = None
    access_type = "w"
    if os.path.exists(f):
        with open(f, "r") as fp_:
            content = fp_.read()
            access_type = "a"
            fp_.close()

    # Extra package title, pypi URL, stars, pulls, and github URL
    with open(f, access_type) as fp_:
        for idx, chunk in enumerate(data_tasks):
            # Get package name and pypi URL
            package_name = chunk["title"]
            package_url = chunk["href"]
            if content is not None and package_url in content:
                continue

            # Get github URL
            package_github = None
            driver.get(package_url)
            soup = BeautifulSoup(driver.page_source, "html.parser")
            for link in soup.find_all("a", class_="vertical-tabs__tab--with-icon"):
                found = False
                for x in ["Source", "Code", "Homepage"]:
                    if (
                        x.lower() in link.get_text().lower()
                        and "github" in link["href"].lower()
                    ):
                        package_github = link["href"]
                        found = True
                        break
                if found:
                    break

            # Get stars and pulls from github API
            stars_count, pulls_count = None, None
            if package_github is not None:
                repo_parts = package_github.split("/")[-2:]
                owner, name = repo_parts[0], repo_parts[1]

                try:
                    repo = api.repos.get(owner, name)
                    stars_count = int(repo["stargazers_count"])
                    issues = api.issues.list_for_repo(owner, name)
                    pulls_count = int(issues[0]["number"])
                except:
                    pass

            # Write to file
            print(
                json.dumps(
                    {
                        "rank": idx,
                        "name": package_name,
                        "url": package_url,
                        "github": package_github,
                        "stars": stars_count,
                        "pulls": pulls_count,
                    }
                ),
                file=fp_,
                flush=True,
            )


if __name__ == "__main__":
    # Start selenium driver to get top 5000 pypi page
    url_top_pypi = "https://hugovk.github.io/top-pypi-packages/"
    driver = webdriver.Chrome()
    driver.get(url_top_pypi)
    button = driver.find_element(By.CSS_SELECTOR, 'button[ng-click="show(5000)"]')
    button.click()

    # Retrieve HTML for packages from page
    soup = BeautifulSoup(driver.page_source, "html.parser")
    package_list = soup.find("div", {"class": "list"})
    packages = package_list.find_all("a", class_="ng-scope")

    get_package_stats(packages, "pypi_rankings.jsonl")

collect/get_top_pypi.py
princeton-nlp/SWE-bench
import subprocess

repos = ["Repos here"]

for repo in repos:
    print(f"Making mirror repo for {repo}")
    out_make = subprocess.run(
        f"./make_repo.sh {repo}",
        shell=True,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    if out_make.returncode != 0:
        print(f"Error making mirror repo for {repo}")
    else:
        print(f"Success making mirror repo for {repo}")

collect/make_repo/call_make_repo.py
princeton-nlp/SWE-bench
import argparse
import json
import logging
import os
from typing import Optional

from fastcore.xtras import obj2dict
from utils import Repo

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


def log_all_pulls(repo: Repo, output: str):
    """
    Iterate over all pull requests in a repository and log them to a file

    Args:
        repo (Repo): repository object
        output (str): output file name
    """
    with open(output, "w") as output:
        for pull in repo.get_all_pulls():
            setattr(pull, "resolved_issues", repo.extract_resolved_issues(pull))
            print(json.dumps(obj2dict(pull)), end="\n", flush=True, file=output)


def main(repo_name: str, output: str, token: Optional[str] = None):
    """
    Logic for logging all pull requests in a repository

    Args:
        repo_name (str): name of the repository
        output (str): output file name
        token (str, optional): GitHub token
    """
    if token is None:
        token = os.environ["GITHUB_TOKEN"]
    owner, repo = repo_name.split("/")
    repo = Repo(owner, repo, token=token)
    log_all_pulls(repo, output)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("repo_name", type=str, help="Name of the repository")
    parser.add_argument("output", type=str, help="Output file name")
    parser.add_argument("--token", type=str, help="GitHub token")
    args = parser.parse_args()
    main(**vars(args))

collect/print_pulls.py
princeton-nlp/SWE-bench
import logging
import re
import requests
import time

from bs4 import BeautifulSoup
from ghapi.core import GhApi
from fastcore.net import HTTP404NotFoundError, HTTP403ForbiddenError
from typing import Dict, List, Tuple, Optional

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


class Repo:
    def __init__(self, owner: str, name: str, token: Optional[str] = None):
        """
        Init to retrieve target repository and create ghapi tool

        Args:
            owner (str): owner of target repository
            name (str): name of target repository
            token (str): github token
        """
        self.owner = owner
        self.name = name
        self.token = token
        self.api = GhApi(token=token)
        self.repo = self.call_api(self.api.repos.get, owner=owner, repo=name)

    def call_api(self, func: callable, **kwargs) -> Dict:
        """
        API call wrapper with rate limit handling (checks every 5 minutes if rate limit is reset)

        Args:
            func (callable): API function to call
            **kwargs: keyword arguments to pass to API function
        Return:
            values (dict): response object of `func`
        """
        while True:
            try:
                values = func(**kwargs)
                return values
            except HTTP403ForbiddenError as e:
                while True:
                    rl = self.api.rate_limit.get()
                    logger.info(
                        f"[{self.owner}/{self.name}] Rate limit exceeded, waiting for 5 minutes, remaining: {rl.resources.core.remaining}"
                    )
                    if rl.resources.core.remaining > 0:
                        break
                    time.sleep(60 * 5)
            except HTTP404NotFoundError as e:
                logger.info(f"[{self.owner}/{self.name}] Resource not found {kwargs}")
                return None

    def extract_resolved_issues(self, pull: Dict) -> List[str]:
        """
        Extract list of issues referenced by a PR

        Args:
            pull (dict): PR dictionary object from GitHub
        Return:
            resolved_issues (list): list of issue numbers referenced by PR
        """
        # Define 1. issue number regex pattern 2. comment regex pattern 3. keywords
        issues_pat = re.compile(r"(\w+)\s+\#(\d+)")
        comments_pat = re.compile(r"(?s)<!--.*?-->")
        keywords = {
            "close",
            "closes",
            "closed",
            "fix",
            "fixes",
            "fixed",
            "resolve",
            "resolves",
            "resolved",
        }

        # Construct text to search over for issue numbers from PR body and commit messages
        text = pull.title if pull.title else ""
        text += "\n" + (pull.body if pull.body else "")
        commits = self.get_all_loop(
            self.api.pulls.list_commits, pull_number=pull.number, quiet=True
        )
        commit_messages = [commit.commit.message for commit in commits]
        commit_text = "\n".join(commit_messages) if commit_messages else ""
        text += "\n" + commit_text
        # Remove comments from text
        text = comments_pat.sub("", text)
        # Look for issue numbers in text via scraping <keyword, number> patterns
        references = dict(issues_pat.findall(text))
        resolved_issues = list()
        if references:
            for word, issue_num in references.items():
                if word.lower() in keywords:
                    resolved_issues.append(issue_num)
        return resolved_issues

    def get_all_loop(
        self,
        func: callable,
        per_page: int = 100,
        num_pages: Optional[int] = None,
        quiet: bool = False,
        **kwargs,
    ) -> List:
        """
        Return all values from a paginated API endpoint.
        
        Args:
            func (callable): API function to call
            per_page (int): number of values to return per page
            num_pages (int): number of pages to return
            quiet (bool): whether to print progress
            **kwargs: keyword arguments to pass to API function
        """
        page = 1
        args = {
            "owner": self.owner,
            "repo": self.name,
            "per_page": per_page,
            **kwargs,
        }
        while True:
            try:
                # Get values from API call
                values = func(**args, page=page)
                for value in values:
                    yield value
                if len(values) == 0:
                    break
                if not quiet:
                    rl = self.api.rate_limit.get()
                    logger.info(
                        f"[{self.owner}/{self.name}] Processed page {page} ({per_page} values per page). Remaining calls: {rl.resources.core.remaining}"
                    )
                if num_pages is not None and page >= num_pages:
                    break
                page += 1
            except Exception as e:
                # Rate limit handling
                logger.error(f"Error processing page {page}: {e}")
                while True:
                    rl = self.api.rate_limit.get()
                    if rl.resources.core.remaining > 0:
                        break
                    logger.info(
                        f"[{self.owner}/{self.name}] Waiting for rate limit reset, checking again in 5 minutes"
                    )
                    time.sleep(60 * 5)
        if not quiet:
            logger.info(
                f"[{self.owner}/{self.name}] Processed {(page-1)*per_page + len(values)} values"
            )

    def get_all_issues(
        self,
        per_page: int = 100,
        num_pages: Optional[int] = None,
        direction: str = "asc",
        sort: str = "created",
        state: str = "closed",
        quiet: bool = False,
    ) -> List:
        """
        Wrapper for API call to get all issues from repo

        Args:
            per_page (int): number of issues to return per page
            num_pages (int): number of pages to return
            direction (str): direction to sort issues
            sort (str): field to sort issues by
            state (str): state of issues to look for
            quiet (bool): whether to print progress
        """
        issues = self.get_all_loop(
            self.api.issues.list_for_repo,
            num_pages=num_pages,
            per_page=per_page,
            direction=direction,
            sort=sort,
            state=state,
            quiet=quiet,
        )
        return issues

    def get_all_pulls(
        self,
        per_page: int = 100,
        num_pages: Optional[int] = None,
        direction: str = "asc",
        sort: str = "created",
        state: str = "closed",
        quiet: str = False,
    ) -> List:
        """
        Wrapper for API call to get all PRs from repo

        Args:
            per_page (int): number of PRs to return per page
            num_pages (int): number of pages to return
            direction (str): direction to sort PRs
            sort (str): field to sort PRs by
            state (str): state of PRs to look for
            quiet (bool): whether to print progress
        """
        pulls = self.get_all_loop(
            self.api.pulls.list,
            num_pages=num_pages,
            direction=direction,
            per_page=per_page,
            sort=sort,
            state=state,
            quiet=quiet,
        )
        return pulls


def extract_problem_statement_and_hints(pull: Dict, repo: Repo) -> Tuple[str, str]:
    """
    Extract problem statement from issues associated with a pull request

    Args:
        pull (dict): PR dictionary object from GitHub
        repo (Repo): Repo object
    Return:
        text (str): problem statement
        hints (str): hints
    """
    if repo.name == "django":
        return extract_problem_statement_and_hints_django(pull, repo)
    text = ""
    all_hint_texts = list()
    for issue_number in pull["resolved_issues"]:
        issue = repo.call_api(
            repo.api.issues.get,
            owner=repo.owner,
            repo=repo.name,
            issue_number=issue_number,
        )
        if issue is None:
            continue
        title = issue.title if issue.title else ""
        body = issue.body if issue.body else ""
        text += f"{title}\n{body}\n"
        issue_number = issue.number
        hint_texts = _extract_hints(pull, repo, issue_number)
        hint_text = "\n".join(hint_texts)
        all_hint_texts.append(hint_text)
    return text, "\n".join(all_hint_texts) if all_hint_texts else ""


def _extract_hints(pull: dict, repo: Repo, issue_number: int) -> List[str]:
    """
    Extract hints from comments associated with a pull request (before first commit)

    Args:
        pull (dict): PR dictionary object from GitHub
        repo (Repo): Repo object
        issue_number (int): issue number
    Return:
        hints (list): list of hints
    """
    # Get all commits in PR
    commits = repo.get_all_loop(
        repo.api.pulls.list_commits, pull_number=pull["number"], quiet=True
    )
    commits = list(commits)
    if len(commits) == 0:
        # If there are no comments, return no hints
        return []
    # Get time of first commit in PR
    commit_time = commits[0].commit.author.date  # str
    commit_time = time.mktime(time.strptime(commit_time, "%Y-%m-%dT%H:%M:%SZ"))
    # Get all comments in PR
    all_comments = repo.get_all_loop(
        repo.api.issues.list_comments, issue_number=issue_number, quiet=True
    )
    all_comments = list(all_comments)
    # Iterate through all comments, only keep comments created before first commit
    comments = list()
    for comment in all_comments:
        comment_time = time.mktime(
            time.strptime(comment.updated_at, "%Y-%m-%dT%H:%M:%SZ")
        )  # use updated_at instead of created_at
        if comment_time < commit_time:
            comments.append(comment)
        else:
            break
        # only include information available before the first commit was created
    # Keep text from comments
    comments = [comment.body for comment in comments]
    return comments


def extract_patches(pull: Dict, repo: Repo) -> Tuple[str, str]:
    """
    Get patch and test patch from PR

    Args:
        pull (dict): PR dictionary object from GitHub
        repo (Repo): Repo object
    Return:
        patch_change_str (str): gold patch
        patch_test_str (str): test patch
    """
    # Convert diff to patch format with "index" lines removed
    patch = requests.get(pull["diff_url"]).text
    if patch.endswith("\n"):
        patch = patch[:-1]
    # Create change patch and test patch
    patch_change, patch_test = [], []

    # Flag to determine if current diff block is a test or general change
    # Values: 'test', 'diff', None
    flag = None

    for line in patch.split("\n"):
        # Exclude commit specific metadata
        if line.startswith("index "):
            continue
        # Determine if current diff block is a test or general change
        if line.startswith("diff --git a/"):
            words = set(re.split(r" |_|\/|\.", line.lower()))
            flag = (
                "test"
                if ("test" in words or "tests" in words or "testing" in words)
                else "diff"
            )
            if flag != "test" and not line.strip().endswith(".py"):
                flag = None
        # Append line to separate patch depending on flag status
        if flag == "test":
            patch_test.append(line)
        elif flag == "diff":
            patch_change.append(line)

    patch_change_str = "\n".join(patch_change) + "\n" if len(patch_change) > 0 else ""
    patch_test_str = "\n".join(patch_test) + "\n" if len(patch_test) > 0 else ""
    return patch_change_str, patch_test_str


### MARK: Repo Specific Parsing Functions ###
def extract_problem_statement_and_hints_django(
    pull: dict, repo: Repo
) -> Tuple[str, str]:
    """
    Get problem statement and hints from issues associated with a pull request

    Args:
        pull (dict): PR dictionary object from GitHub
        repo (Repo): Repo object
    Return:
        text (str): problem statement
        hints (str): hints
    """
    text = ""
    all_hints_text = list()
    for issue_number in pull["resolved_issues"]:
        url = f"https://code.djangoproject.com/ticket/{issue_number}"
        resp = requests.get(url)
        if resp.status_code != 200:
            continue
        soup = BeautifulSoup(resp.text, "html.parser")

        # Get problem statement (title + body)
        issue_desc = soup.find("div", {"id": "ticket"})
        title = issue_desc.find("h1", class_="searchable").get_text()
        title = re.sub(r"\s+", " ", title).strip()
        body = issue_desc.find("div", class_="description").get_text()
        body = re.sub(r"\s+", " ", body).strip()
        text += f"{title}\n{body}\n"

        # Get time of first commit in PR
        commits = repo.get_all_loop(
            repo.api.pulls.list_commits, pull_number=pull["number"], quiet=True
        )
        commits = list(commits)
        if len(commits) == 0:
            continue
        commit_time = commits[0].commit.author.date
        commit_time = time.mktime(time.strptime(commit_time, "%Y-%m-%dT%H:%M:%SZ"))

        # Get all comments before first commit
        comments_html = soup.find("div", {"id": "changelog"})
        div_blocks = comments_html.find_all("div", class_="change")
        comments = []
        # Loop through each div block
        for div_block in div_blocks:
            # Find the comment text and timestamp
            comment_resp = div_block.find("div", class_="comment")
            timestamp_resp = div_block.find("a", class_="timeline")
            if comment_resp is None or timestamp_resp is None:
                continue

            comment_text = re.sub(r"\s+", " ", comment_resp.text).strip()
            timestamp = timestamp_resp["title"]
            if timestamp.startswith("See timeline at "):
                timestamp = timestamp[len("See timeline at ") :]
            timestamp = time.mktime(time.strptime(timestamp, "%m/%d/%y %H:%M:%S"))

            # Append the comment and timestamp as a tuple to the comments list
            if timestamp < commit_time:
                all_hints_text.append((comment_text, timestamp))

    return text, all_hints_text

collect/utils.py
princeton-nlp/SWE-bench
MAP_VERSION_TO_INSTALL_SKLEARN = {
    k: {
        "python": "3.6",
        "packages": "numpy scipy cython pytest pandas matplotlib",
        "install": "pip install -v --no-use-pep517 --no-build-isolation -e .",
    }
    for k in ["0.20", "0.21", "0.22"]
}
MAP_VERSION_TO_INSTALL_SKLEARN.update(
    {
        k: {
            "python": "3.7",
            "packages": "numpy scipy cython pytest pandas matplotlib",
            "install": "pip install -v --no-use-pep517 --no-build-isolation -e .",
        }
        for k in ["0.23", "0.24"]
    }
)
MAP_VERSION_TO_INSTALL_SKLEARN.update(
    {
        k: {
            "python": "3.9",
            "packages": "numpy scipy cython pytest pandas matplotlib joblib threadpoolctl",
            "install": "pip install -v --no-use-pep517 --no-build-isolation -e .",
        }
        for k in ["1.0", "1.1", "1.2", "1.3", "1.4"]
    }
)

MAP_VERSION_TO_INSTALL_FLASK = {
    "0.11-dev": {
        "python": "3.6",
        "packages": "pytest",
        "pip_packages": "tox",
        "install": "python setup.py develop",
    },
    "1.0": {
        "python": "3.8",
        "packages": "pytest",
        "pip_packages": "tox",
        "install": 'pip install -e ".[dev]"',
    },
    "1.1": {
        "python": "3.9",
        "packages": "pytest",
        "pip_packages": "tox",
        "install": 'pip install -e ".[dev]"',
    },
    "2.0": {
        "python": "3.9",
        "packages": "requirements.txt",
        "install": "pip install -e .",
    },
    "2.1": {
        "python": "3.10",
        "packages": "requirements.txt",
        "install": "pip install -e .",
    },
}
MAP_VERSION_TO_INSTALL_FLASK.update(
    {
        k: {
            "python": "3.6",
            "packages": "pytest",
            "pip_packages": "tox",
            "install": "pip install -e .",
        }
        for k in ["0.11", "0.12", "0.12-dev", "0.13-dev"]
    }
)
MAP_VERSION_TO_INSTALL_FLASK.update(
    {
        k: {
            "python": "3.11",
            "packages": "requirements.txt",
            "install": "pip install -e .",
        }
        for k in ["2.2", "2.3"]
    }
)

MAP_VERSION_TO_INSTALL_DJANGO = {
    k: {
        "python": "3.5",
        "packages": "requirements.txt",
        "install": "python setup.py install",
    }
    for k in ["1.7", "1.8", "1.9", "1.10", "1.11", "2.0", "2.1", "2.2"]
}
MAP_VERSION_TO_INSTALL_DJANGO.update(
    {
        k: {"python": "3.5", "install": "python setup.py install"}
        for k in ["1.4", "1.5", "1.6"]
    }
)
MAP_VERSION_TO_INSTALL_DJANGO.update(
    {
        k: {
            "python": "3.6",
            "packages": "requirements.txt",
            "install": "python -m pip install -e .",
        }
        for k in ["3.0", "3.1", "3.2"]
    }
)
MAP_VERSION_TO_INSTALL_DJANGO.update(
    {
        k: {
            "python": "3.8",
            "packages": "requirements.txt",
            "install": "python -m pip install -e .",
        }
        for k in ["4.0"]
    }
)
MAP_VERSION_TO_INSTALL_DJANGO.update(
    {
        k: {
            "python": "3.9",
            "packages": "requirements.txt",
            "install": "python -m pip install -e .",
        }
        for k in ["4.1", "4.2"]
    }
)
MAP_VERSION_TO_INSTALL_DJANGO.update(
    {
        k: {
            "python": "3.11",
            "packages": "requirements.txt",
            "install": "python -m pip install -e .",
        }
        for k in ["5.0"]
    }
)

MAP_VERSION_TO_INSTALL_REQUESTS = {
    k: {"python": "3.9", "packages": "pytest", "install": "python -m pip install ."}
    for k in
        ["0.7", "0.8", "0.9", "0.11", "0.13", "0.14", "1.1", "1.2", "2.0", "2.2"] + \
        ["2.3", "2.4", "2.5", "2.7", "2.8", "2.9", "2.10", "2.11", "2.12", "2.17"] + \
        ["2.18", "2.19", "2.22", "2.26", "2.25", "2.27", "3.0"]
}

MAP_VERSION_TO_INSTALL_SEABORN = {
    k: {"python": "3.9", "install": "pip install -e .", "pip_packages": "pytest"}
    for k in ["0.3", "0.4", "0.5", "0.6", "0.7", "0.8", "0.9", "0.11"]
}
MAP_VERSION_TO_INSTALL_SEABORN.update(
    {k: {"python": "3.9", "install": "pip install -e .[dev]"} for k in ["0.12", "0.13"]}
)

MAP_VERSION_TO_INSTALL_PYTEST = {
    k: {"python": "3.9", "install": "pip install -e ."}
    for k in ["3.10", "6.0", "6.2", "6.3", "8.0"]
    + [
        str(round(0.1 * x, 1))
        for interval in [(30, 47), (50, 55), (70, 75)]
        for x in range(interval[0], interval[1], 1)
    ]
}

MAP_VERSION_TO_INSTALL_MATPLOTLIB = {
    k: {
        "python": "3.9",
        "packages": "environment.yml",
        "install": "python -m pip install -e .",
    }
    for k in ["3.5", "3.6", "3.7"]
}
MAP_VERSION_TO_INSTALL_MATPLOTLIB.update(
    {
        k: {
            "python": "3.8",
            "packages": "requirements.txt",
            "install": "python -m pip install -e .",
        }
        for k in ["3.1", "3.2", "3.3", "3.4"]
    }
)
MAP_VERSION_TO_INSTALL_MATPLOTLIB.update(
    {
        k: {
            "python": "3.7",
            "packages": "requirements.txt",
            "install": "python -m pip install -e .",
        }
        for k in ["3.0"]
    }
)
MAP_VERSION_TO_INSTALL_MATPLOTLIB.update(
    {
        k: {
            "python": "3.5",
            "install": "python setup.py build; python setup.py install",
        }
        for k in ["2.0", "2.1", "2.2", "1.0", "1.1", "1.2", "1.3", "1.4", "1.5"]
    }
)

MAP_VERSION_TO_INSTALL_SPHINX = {
    k: {
        "python": "3.9",
        "pip_packages": "tox",
        "install": "pip install -e .[test]",
        "pre_install": ["sed -i 's/pytest/pytest -rA/' tox.ini"],
    } for k in
        ["1.5", "1.6", "1.7", "1.8", "2.0", "2.1", "2.2", "2.3", "2.4", "3.0"] + \
        ["3.1", "3.2", "3.3", "3.4", "3.5", "4.0", "4.1", "4.2", "4.3", "4.4"] + \
        ["4.5", "5.0", "5.1", "5.2", "5.3", "6.0", "6.2", "7.0", "7.1", "7.2"]
}
for k in ["3.0", "3.1", "3.2", "3.3", "3.4", "3.5", "4.0"]:
    MAP_VERSION_TO_INSTALL_SPHINX[k][
        "pre_install"
    ].append("sed -i 's/Jinja2>=2.3/Jinja2<3.1/' setup.py")

MAP_VERSION_TO_INSTALL_ASTROPY = {
    k: {"python": "3.9", "install": "pip install -e .[test]"}
    for k in
        ["0.1", "0.2", "0.3", "0.4", "1.1", "1.2", "1.3", "3.0", "3.1", "3.2"] + \
        ["4.1", "4.2", "4.3", "5.0", "5.1", "5.2"]
}

MAP_VERSION_TO_INSTALL_SYMPY = {
    k: {
        "python": "3.9",
        "packages": "mpmath flake8",
        "pip_packages": "flake8-comprehensions",
        "install": "pip install -e .",
    }
    for k in
        ["0.7", "1.0", "1.1", "1.10", "1.11", "1.12", "1.2", "1.4", "1.5", "1.6"] + \
        ["1.7", "1.8", "1.9"]
}
MAP_VERSION_TO_INSTALL_SYMPY.update(
    {
        k: {
            "python": "3.9",
            "packages": "requirements.txt",
            "install": "pip install -e .",
        }
        for k in ["1.13"]
    }
)

MAP_VERSION_TO_INSTALL_PYLINT = {
    k: {"python": "3.9", "packages": "requirements.txt", "install": "pip install -e ."}
    for k in ["2.10", "2.11", "2.13", "2.14", "2.15", "2.16", "2.17", "2.8", "2.9", "3.0"]
}

MAP_VERSION_TO_INSTALL_XARRAY = {
    k: {
        "python": "3.10",
        "packages": "environment.yml",
        "install": "pip install -e .",
        "pip_packages": "pytest",
        "no_use_env": True,
    }
    for k in ["0.12", "0.18", "0.19", "0.20", "2022.03", "2022.06", "2022.09"]
}

# Constants - Task Instance Instllation Environment
MAP_VERSION_TO_INSTALL = {
    "astropy/astropy": MAP_VERSION_TO_INSTALL_ASTROPY,
    "django/django": MAP_VERSION_TO_INSTALL_DJANGO,
    "matplotlib/matplotlib": MAP_VERSION_TO_INSTALL_MATPLOTLIB,
    "mwaskom/seaborn": MAP_VERSION_TO_INSTALL_SEABORN,
    "pallets/flask": MAP_VERSION_TO_INSTALL_FLASK,
    "psf/requests": MAP_VERSION_TO_INSTALL_REQUESTS,
    "pydata/xarray": MAP_VERSION_TO_INSTALL_XARRAY,
    "pylint-dev/pylint": MAP_VERSION_TO_INSTALL_PYLINT,
    "pytest-dev/pytest": MAP_VERSION_TO_INSTALL_PYTEST,
    "scikit-learn/scikit-learn": MAP_VERSION_TO_INSTALL_SKLEARN,
    "sphinx-doc/sphinx": MAP_VERSION_TO_INSTALL_SPHINX,
    "sympy/sympy": MAP_VERSION_TO_INSTALL_SYMPY,
}

# Constants - Repository Specific Installation Instructions
MAP_REPO_TO_INSTALL = {}

# Constants - Task Instance Test Frameworks
TEST_PYTEST = "pytest --no-header -rA --tb=no -p no:cacheprovider"
MAP_REPO_TO_TEST_FRAMEWORK = {
    "astropy/astropy": TEST_PYTEST,
    "django/django": "./tests/runtests.py --verbosity 2",
    "matplotlib/matplotlib": TEST_PYTEST,
    "mwaskom/seaborn": "pytest --no-header -rA",
    "pallets/flask": TEST_PYTEST,
    "psf/requests": TEST_PYTEST,
    "pydata/xarray": TEST_PYTEST,
    "pylint-dev/pylint": TEST_PYTEST,
    "pytest-dev/pytest": "pytest -rA",
    "scikit-learn/scikit-learn": TEST_PYTEST,
    "sphinx-doc/sphinx": "tox -epy39 -v --",
    "sympy/sympy": "bin/test -C --verbose",
}

# Constants - Task Instance Version Extraction
MAP_REPO_TO_VERSION_PATHS = {
    "django/django": ["django/__init__.py"],
    "mwaskom/seaborn": ["seaborn/__init__.py"],
    "pallets/flask": ["src/flask/__init__.py", "flask/__init__.py"],
    "psf/requests": ["requests/__version__.py", "requests/__init__.py"],
    "pyca/cryptography": [
        "src/cryptography/__about__.py",
        "src/cryptography/__init__.py",
    ],
    "pylint-dev/pylint": ["pylint/__pkginfo__.py", "pylint/__init__.py"],
    "pytest-dev/pytest": ["src/_pytest/_version.py", "_pytest/_version.py"],
    "Qiskit/qiskit": ["qiskit/VERSION.txt"],
    "scikit-learn/scikit-learn": ["sklearn/__init__.py"],
    "sphinx-doc/sphinx": ["sphinx/__init__.py"],
    "sympy/sympy": ["sympy/release.py", "sympy/__init__.py"],
}

MAP_REPO_TO_VERSION_PATTERNS = {
    k: [r'__version__ = [\'"](.*)[\'"]', r"VERSION = \((.*)\)"]
    for k in [
        "scikit-learn/scikit-learn",
        "pallets/flask",
        "django/django",
        "psf/requests",
        "mwaskom/seaborn",
        "pyca/cryptography",
        "sphinx-doc/sphinx",
        "sympy/sympy",
        "pylint-dev/pylint",
    ]
}
MAP_REPO_TO_VERSION_PATTERNS.update(
    {
        k: [
            r'__version__ = [\'"](.*)[\'"]',
            r'__version__ = version = [\'"](.*)[\'"]',
            r"VERSION = \((.*)\)",
        ]
        for k in ["pytest-dev/pytest", "matplotlib/matplotlib"]
    }
)
MAP_REPO_TO_VERSION_PATTERNS.update({k: [r"(.*)"] for k in ["Qiskit/qiskit"]})

# Constants - Task Instance Requirements File Paths
MAP_REPO_TO_REQS_PATHS = {
    "django/django": ["tests/requirements/py3.txt"],
    "matplotlib/matplotlib": ["requirements/dev/dev-requirements.txt", "requirements/testing/travis_all.txt"],
    "pallets/flask": ["requirements/dev.txt"],
    "pylint-dev/pylint": ["requirements_test.txt"],
    "sympy/sympy": ["requirements-dev.txt"],
}

# Constants - Task Instance environment.yml File Paths
MAP_REPO_TO_ENV_YML_PATHS = {
    "matplotlib/matplotlib": ["environment.yml"],
    "pydata/xarray": ["ci/requirements/environment.yml", "environment.yml"],
}

# Constants - Evaluation Keys
KEY_INSTANCE_ID = "instance_id"
KEY_MODEL = "model"
KEY_PREDICTION = "prediction"

# Constants - Logging
APPLY_PATCH_FAIL = ">>>>> Patch Apply Failed"
APPLY_PATCH_PASS = ">>>>> Applied Patch"
INSTALL_FAIL = ">>>>> Init Failed"
INSTALL_PASS = ">>>>> Init Succeeded"
INSTALL_TIMEOUT = ">>>>> Init Timed Out"
RESET_FAILED = ">>>>> Reset Failed"
TESTS_ERROR = ">>>>> Tests Errored"
TESTS_FAILED = ">>>>> Some Tests Failed"
TESTS_PASSED = ">>>>> All Tests Passed"
TESTS_TIMEOUT = ">>>>> Tests Timed Out"

# Constants - Miscellaneous
NON_TEST_EXTS = [".json", ".png", "csv", ".txt", ".md", ".jpg", ".jpeg", ".pkl", ".yml", ".yaml", ".toml"]
SWE_BENCH_URL_RAW = "https://raw.githubusercontent.com/"
harness/constants.py
princeton-nlp/SWE-bench
import logging, os, subprocess

from constants import (
    APPLY_PATCH_FAIL,
    APPLY_PATCH_PASS,
    INSTALL_FAIL,
    INSTALL_PASS,
    INSTALL_TIMEOUT,
    KEY_INSTANCE_ID,
    KEY_MODEL,
    MAP_REPO_TO_INSTALL,
    MAP_REPO_TO_TEST_FRAMEWORK,
    MAP_VERSION_TO_INSTALL,
    RESET_FAILED,
    TESTS_FAILED,
    TESTS_PASSED,
    TESTS_TIMEOUT,
    TESTS_ERROR,
)
from tempfile import TemporaryDirectory
from typing import Dict, List
from utils import (
    clone_repo,
    get_conda_env_names,
    get_environment_yml,
    get_requirements,
    get_test_directives,
)

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger_testbed = logging.getLogger("testbed_context_manager")


class TestbedContextManager:
    def __init__(
        self,
        task_instances: List,
        log_dir: str,
        path_conda: str = None,
        testbed: str = None,
        verbose: bool = False,
        timeout: int = None,
        temp_dir: str = None,
    ):
        """
        Initialize testbed context. Creates temporary directories and groups task instances
        by repo/version.

        Args:
            task_instances (list): List of task instances
            log_dir (str): Path to log directory
            path_conda (str): Path to conda installation
            testbed (str): Path to testbed directory
            verbose (bool): Whether to show logs
            timeout (int): Timeout for actions
            temp_dir (str): Path to temporary directory
        """
        logger_testbed.propagate = verbose
        self.verbose = verbose
        self.old_dir = os.getcwd()
        self.log_dir = log_dir
        self.timeout = timeout
        self.subprocess_args = {"check": True, "stdout": subprocess.DEVNULL, "stderr": subprocess.DEVNULL}

        # Create log, temp directories if they don't exist
        if not os.path.exists(self.log_dir):
            logger_testbed.info(f"[Testbed] Creating log directory {self.log_dir}")
            os.makedirs(self.log_dir, exist_ok=True)
        if temp_dir is not None and not os.path.exists(temp_dir):
            logger_testbed.info(f"[Testbed] Creating temp directory {temp_dir}")
            os.makedirs(temp_dir, exist_ok=True)

        # Set up conda path, create in temp directory if None
        if path_conda is not None:
            self.temp_dir_conda = None
            self.path_conda = path_conda
        else:
            self.temp_dir_conda = TemporaryDirectory(dir=temp_dir)
            self.path_conda = self.temp_dir_conda.name
        logger_testbed.info(f"[Testbed] Using conda path {self.path_conda}")

        # Set up testbed path, create in temp directory if None
        if testbed is not None:
            self.temp_dir_work = None
            self.testbed = testbed
        else:
            self.temp_dir_work = TemporaryDirectory(dir=temp_dir)
            self.testbed = self.temp_dir_work.name
        logger_testbed.info(f"[Testbed] Using working directory {self.testbed} for testbed")

        # Sort task instances by created_at
        self.task_instances = sorted(task_instances, key=lambda x: x["created_at"], reverse=True)

        # Group repos by repo, then version
        self.task_instances_grouped = {}
        for instance in self.task_instances:
            # Create test command from framework + directives
            test_type = MAP_REPO_TO_TEST_FRAMEWORK[instance["repo"]]
            test_directives = get_test_directives(instance)
            instance["test_cmd"] = f"{test_type} {' '.join(test_directives)}"

            # Group task instances by repo, version
            repo = instance["repo"]
            version = instance["version"] if "version" in instance else None
            if repo not in self.task_instances_grouped:
                self.task_instances_grouped[repo] = {}
            if version not in self.task_instances_grouped[repo]:
                self.task_instances_grouped[repo][version] = []
            self.task_instances_grouped[repo][version].append(instance)

        # Log grouped task instances to be run
        self.setup_refs = {}
        for repo, map_version_to_instances in self.task_instances_grouped.items():
            logger_testbed.info(f"[Testbed] Repo {repo}: {len(map_version_to_instances)} versions")

            # Determine instances to use for environment installation
            self.setup_refs[repo] = {}
            for version, instances in map_version_to_instances.items():
                logger_testbed.info(f"[Testbed] \tVersion {version}: {len(instances)} instances")
                self.setup_refs[repo][version] = instances[0]

        # Remove None versions, versions not in MAP_VERSION_TO_INSTALL
        self._custom_restraints()

    def __enter__(self):
        """
        Set up testbed (conda environments, git repositories)
        """
        # If path_conda not provided, create temporary miniconda3 installation
        if self.temp_dir_conda is not None:
            # Set up the paths for Miniconda
            self.path_conda = os.path.join(self.path_conda, "miniconda3")
            os.mkdir(self.path_conda)
            miniconda_sh = os.path.join(self.path_conda, "miniconda.sh")
            logger_testbed.info(f"No conda path provided, creating temporary install in {self.path_conda}...")

            # Download Miniconda installer
            download_cmd = [
                "wget",
                "https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh",
                "-O",
                miniconda_sh,
            ]
            subprocess.run(download_cmd, **self.subprocess_args)

            # Install Miniconda
            install_cmd = ["bash", miniconda_sh, "-b", "-u", "-p", self.path_conda]
            subprocess.run(install_cmd, **self.subprocess_args)

            # Clean up the installer
            os.remove(miniconda_sh)
        logger_testbed.info(f"[Testbed] Using conda path {self.path_conda}")

        # Set up conda executables, get existing environments
        self.path_conda = os.path.abspath(self.path_conda)
        path_activate = os.path.join(self.path_conda, "bin", "activate")
        exec_type = "mamba" if "mamba" in self.path_conda else "conda"
        exec_cmd = os.path.join(self.path_conda, "bin", exec_type)
        env_list = get_conda_env_names(exec_cmd)

        # Set up testbed (environment, github repo) for each repo
        for repo, version_to_setup_ref in self.setup_refs.items():
            repo_prefix = repo.replace("/", "__")

            # Run any repo-level installation commands if provided
            if repo in MAP_REPO_TO_INSTALL:
                install_cmd = MAP_REPO_TO_INSTALL[repo]
                logger_testbed.info(f"[Testbed] Running custom install command for {repo}: {install_cmd}")
                subprocess.run(install_cmd, shell=True, **self.subprocess_args)

            # Create conda environment per version of the repo
            for version, install in MAP_VERSION_TO_INSTALL[repo].items():
                # Skip if none of the task instances are for this version
                if version not in version_to_setup_ref:
                    continue

                # Name for both environment and github repo
                env_name = f"{repo_prefix}__{version}"
                logger_testbed.info(f"[Testbed] Setting up testbed for {env_name}")

                # Clone github per repo/version
                repo_path = os.path.join(self.testbed, env_name)
                if not os.path.exists(repo_path):
                    clone_repo(repo, repo_path)
                    logger_testbed.info(f"[Testbed] Cloned {repo} to {repo_path}")
                else:
                    logger_testbed.info(
                        f"[Testbed] Repo for {repo_prefix} version {version} exists: {repo_path}; skipping"
                    )

                # Skip if conda environment already exists
                if env_name in env_list:
                    logger_testbed.info(
                        f"[Testbed] Environment {env_name} already exists; skipping"
                    )
                    continue

                # Get setup reference instance
                setup_ref_instance = version_to_setup_ref[version]

                # Create conda environment according to install instructinos
                pkgs = install["packages"] if "packages" in install else ""
                if pkgs == "requirements.txt":
                    # Create environment
                    cmd = f"{exec_cmd} create -n {env_name} python={install['python']} -y"
                    logger_testbed.info(f"[Testbed] Creating environment {env_name}; Command: {cmd}")
                    subprocess.run(cmd, shell=True, **self.subprocess_args)

                    # Install dependencies
                    path_to_reqs = get_requirements(setup_ref_instance, self.testbed)
                    cmd = f"source {path_activate} {env_name} && pip install -r {path_to_reqs}"
                    logger_testbed.info(f"[Testbed] Installing dependencies for {env_name}; Command: {cmd}")
                    subprocess.run(cmd, shell=True, **self.subprocess_args)
                    os.remove(path_to_reqs)
                elif pkgs == "environment.yml":
                    # Create environment from yml
                    path_to_reqs = get_environment_yml(setup_ref_instance, env_name, self.testbed)
                    if "no_use_env" in install and install["no_use_env"]:
                        # `conda create` based installation
                        cmd = f"{exec_cmd} create -c conda-forge -n {env_name} python={install['python']} -y"
                        logger_testbed.info(f"[Testbed] Creating environment {env_name}; Command: {cmd}")

                        # Install dependencies
                        subprocess.run(cmd, shell=True, **self.subprocess_args)
                        cmd = f"{exec_cmd} env update -f {path_to_reqs}"
                        logger_testbed.info(f"[Testbed] Installing dependencies for {env_name}; Command: {cmd}")

                        subprocess.run(cmd, shell=True, **self.subprocess_args)
                    else:
                        # `conda env create` based installation
                        cmd = f"{exec_cmd} env create --file {path_to_reqs}"
                        logger_testbed.info(f"[Testbed] Creating environment {env_name}; Command: {cmd}")
                        subprocess.run(cmd, shell=True, **self.subprocess_args)

                    # Remove environment.yml
                    os.remove(path_to_reqs)
                else:
                    # Create environment + install dependencies
                    cmd = f"{exec_cmd} create -n {env_name} python={install['python']} {pkgs} -y"
                    logger_testbed.info(f"[Testbed] Creating environment {env_name}; Command: {cmd}")
                    subprocess.run(cmd, shell=True, **self.subprocess_args)

                # Install additional packages if specified
                if "pip_packages" in install:
                    cmd = f"source {path_activate} {env_name} && pip install {install['pip_packages']}"
                    logger_testbed.info(f"[Testbed] Installing pip packages for {env_name}; Command: {cmd}")
                    subprocess.run(cmd, shell=True, **self.subprocess_args)

        return self

    def get_distributed_tasks(self) -> List:
        """
        Create task group (instances + keywords) for each repo/version

        Returns:
            list: List of task groups, each group containing task instances
                from the same repo with the same version
        """
        distributed_tasks = []
        for repo, map_version_to_instances in self.task_instances_grouped.items():
            repo_prefix = repo.replace("/", "__")
            for version, instances in map_version_to_instances.items():
                env_name = f"{repo_prefix}__{version}"
                task_set = {
                    "conda_path": self.path_conda,
                    "log_dir": self.log_dir,
                    "task_instances": instances,
                    "testbed": os.path.join(self.testbed, env_name),
                    "timeout": self.timeout,
                    "venv": env_name,
                    "version": version,
                    "verbose": self.verbose,
                }
                distributed_tasks.append(task_set)
        return distributed_tasks

    def _custom_restraints(self):
        """
        Custom restraints per repo
        """
        for repo, group in self.task_instances_grouped.items():
            if None in group:
                logger_testbed.info(f"[Testbed] Removed None version from repo {repo}")
                del group[None]
            versions = list(group.keys())
            for version in versions:
                if version not in MAP_VERSION_TO_INSTALL[repo]:
                    logger_testbed.info(f"[Testbed] Removed {version} version from repo {repo} (Install instructions not given)")
                    del group[version]

    def __exit__(self, exc_type, exc_value, exc_traceback):
        if self.temp_dir_work is not None:
            self.temp_dir_work.cleanup()
        if self.temp_dir_conda is not None:
            self.temp_dir_conda.cleanup()


logger_taskenv = logging.getLogger("taskenv_context_manager")


class TaskEnvContextManager:
    def __init__(
        self,
        instance: Dict,
        testbed: str,
        venv: str,
        log_dir: str,
        conda_path: str,
        verbose: bool = False,
        timeout: int = None,
        is_eval: bool = False,
    ):
        """
        Sets up execution context for a single task instance

        Args:
            instance (dict): Task instance
            testbed (str): Path to testbed directory
            venv (str): Name of conda environment (should exist in conda_path)
            log_dir (str): Path to log directory
            conda_path (str): Path to conda installation
            verbose (bool): Whether to show logs
            timeout (int): Timeout for actions
            is_eval (bool): Whether this is for evaluating a model on SWE Bench
                (Mainly for logging purposes)
        """
        logger_taskenv.propagate = verbose
        self.instance = instance
        self.testbed = testbed
        self.testbed_name = testbed.split("/")[-1]
        self.venv = venv
        self.conda_path = conda_path
        self.log_file = os.path.join(log_dir, f"{instance[KEY_INSTANCE_ID]}.log")
        self.is_eval = is_eval
        if is_eval:
            self.log_file = os.path.join(log_dir, f"{instance[KEY_INSTANCE_ID]}.{instance[KEY_MODEL]}.eval.log")
        self.cmd_activate = (f"source {os.path.join(self.conda_path, 'bin', 'activate')} {self.venv}")
        self.timeout = timeout
        self.cwd = os.getcwd()

    def __enter__(self):
        """
        Enter task environment, set up log file
        """
        os.chdir(self.testbed)
        with open(self.log_file, "w") as f:
            f.write(f"Task Metadata:\n\t- Instance ID: {self.instance[KEY_INSTANCE_ID]}\n\t- Testbed: {self.testbed}\n\t- Virtual Env.: {self.venv}\n")
            if self.is_eval:
                f.write(f"\t- Evaluation Model: {self.instance[KEY_MODEL]}\n")
        return self

    def reset_task_env(self, instance: Dict):
        """
        Reset task environment + testbed and checkout base commit of given task instance

        Args:
            instance (dict): Task instance
        Returns:
            bool: True if reset successful, False otherwise
        """
        subprocess_args = {"check": True, "shell": True, "stdout": subprocess.DEVNULL, "stderr": subprocess.DEVNULL}
        try:
            # Remove all paths in .gitignore
            if os.path.exists(".gitignore"):
                with open(".gitignore", "r") as f:
                    for line in f.readlines():
                        if line.startswith("#") or line.strip() == "":
                            continue
                        subprocess.run(f"rm -rf {line}", **subprocess_args)

            # Reset git repo + checkout base commit
            subprocess.run("git restore .", **subprocess_args)
            subprocess.run("git reset HEAD .", **subprocess_args)
            subprocess.run("git clean -fdx", **subprocess_args)
            subprocess.run(f"git -c advice.detachedHead=false checkout {instance['base_commit']}", **subprocess_args)
            logger_taskenv.info(f"[{self.testbed_name}] [{instance[KEY_INSTANCE_ID]}] Reset task environment to {instance['base_commit']}")
            return True
        except Exception as e:
            err_msg = (f"{RESET_FAILED}; Failed to reset task environment to {instance['base_commit']}: {e}")
            logger_taskenv.error(f"[{self.testbed_name}] {err_msg}")
            with open(self.log_file, "a") as f:
                f.write(err_msg)
            return False

    def run_install_task(self, instance: Dict) -> bool:
        """
        Run installation for task instance

        Args:
            instance (dict): Task instance
        Returns:
            bool: True if installation successful, False otherwise
        """
        # Get installation instructions by repo/version
        specifications = MAP_VERSION_TO_INSTALL[instance["repo"]][instance["version"]]

        # Run pre-install set up if provided
        if "pre_install" in specifications:
            for pre_install in specifications["pre_install"]:
                cmd_pre_install = f"{self.cmd_activate}; {pre_install}"
                logger_taskenv.info(f"[{self.testbed_name}] [{instance[KEY_INSTANCE_ID]}] Running pre-install setup command: {cmd_pre_install}")
                out_pre_install = subprocess.run(
                    cmd_pre_install,
                    shell=True,
                    text=True,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    timeout=self.timeout,
                )
                with open(self.log_file, "a") as f:
                    f.write(f"Pre-installation Command: {cmd_pre_install}\n")
                    f.write(f"Std. Output: {out_pre_install.stdout}\n")
                    f.write(f"Std. Error: {out_pre_install.stderr}\n")
                if out_pre_install.returncode != 0:
                    logger_taskenv.error(f"[{self.testbed_name}] [{instance[KEY_INSTANCE_ID]}] Pre-install setup failed")
                    with open(self.log_file, "a") as f:
                        f.write(f"\n{INSTALL_FAIL}\n")
                    return False

        cmd_install = f"{self.cmd_activate}; {specifications['install']}"
        logger_taskenv.info(f"[{self.testbed_name}] [{instance[KEY_INSTANCE_ID]}] Installing with command: {cmd_install}")
        try:
            # Run installation command
            out_install = subprocess.run(
                cmd_install,
                shell=True,
                text=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                timeout=self.timeout,
            )
            # Write installation logs to log file
            with open(self.log_file, "a") as f:
                f.write(f"Installation Command: {cmd_install}\n")
                f.write(f"Std. Output: {out_install.stdout}\n")
                f.write(f"Std. Error: {out_install.stderr}\n")

            if out_install.returncode != 0:
                # Installation failed
                logger_taskenv.error(f"[{self.testbed_name}] [{instance[KEY_INSTANCE_ID]}] Installation failed")
                with open(self.log_file, "a") as f:
                    f.write(f"\n{INSTALL_FAIL}\n")
                return False

            # Installation successful
            logger_taskenv.info(f"[{self.testbed_name}] [{instance[KEY_INSTANCE_ID]}] Installation successful")
            with open(self.log_file, "a") as f:
                f.write(f"\n{INSTALL_PASS}\n")
            return True
        except subprocess.TimeoutExpired:
            # Installation timed out
            logger_taskenv.error(f"[{self.testbed_name}] [{self.instance[KEY_INSTANCE_ID]}] Installation timed out")
            with open(self.log_file, "a") as f:
                f.write(f"\n{INSTALL_TIMEOUT}\n")
            return False

    def apply_patch(
        self, patch: str, patch_type: str = "", revert: bool = False
    ) -> bool:
        """
        Apply patch to task environment

        Args:
            patch (str): Plaintext of patch to apply
            patch_type (str): Type of patch (e.g. "eval", "test")
        Returns:
            bool: True if patch applied successfully, False otherwise
        """
        # If patch is `None`, indicate in log and skip
        if patch is None:
            logger_taskenv.error(f"[{self.testbed_name}] [{self.instance[KEY_INSTANCE_ID]}] Patch is `None` ({patch_type})")
            with open(self.log_file, "a") as f:
                f.write(f"{APPLY_PATCH_FAIL}; Prediction patch is `None`")
            return False

        # Write patch to temporary patch file in parent directory
        patch_path = os.path.join(
            os.path.dirname(self.testbed.rstrip("/")),
            f"temp_{self.instance[KEY_INSTANCE_ID]}_{patch_type}.patch",
        )
        with open(patch_path, "w") as f:
            f.write(patch)

        # Apply patch to testbed directory
        apply_cmd = f"git apply -v -R {patch_path}" if revert else f"git apply -v {patch_path}"
        out_patch = subprocess.run(
            apply_cmd,
            shell=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )
        os.remove(patch_path)

        log_cmd = "Revert" if revert else "Apply"
        if out_patch.returncode != 0:
            # Patch apply failed
            logger_taskenv.error(f"[{self.testbed_name}] [{self.instance[KEY_INSTANCE_ID]}] {log_cmd} patch failed ({patch_type})")
            with open(self.log_file, "a") as f:
                f.write(f"{APPLY_PATCH_FAIL}; ({patch_type})\nOutput:\n")
                f.write(out_patch.stdout)
                f.write(out_patch.stderr)
            return False

        # Patch apply succeeded
        logger_taskenv.info(f"[{self.testbed_name}] [{self.instance[KEY_INSTANCE_ID]}] {log_cmd} patch successful ({patch_type})")
        with open(self.log_file, "a") as f:
            f.write(f"{APPLY_PATCH_PASS} ({patch_type})\n")
        return True

    def run_tests_task(self, instance: Dict):
        """
        Run tests for task instance

        Args:
            instance (dict): Task instance
        Returns:
            bool: True if test script ran successfully, False otherwise
        """
        try:
            # Run test command for task instance
            test_cmd = f"{self.cmd_activate}; {instance['test_cmd']}"
            with open(self.log_file, "a") as f:
                f.write(f"Test Script: {test_cmd};\n")
            out_test = subprocess.run(test_cmd, shell=True, capture_output=True, timeout=self.timeout)

            # Write test results to log file
            with open(self.log_file, "a") as f:
                f.write(f"Output:\n")
                f.write(out_test.stdout.decode("utf-8"))
                f.write(out_test.stderr.decode("utf-8"))
                if out_test.returncode != 0:
                    f.write(f"\n{TESTS_FAILED}\n")
                else:
                    f.write(f"\n{TESTS_PASSED}\n")

            logger_taskenv.info(f"[{self.testbed_name}] [{instance[KEY_INSTANCE_ID]}] Test script run successful")
            return True
        except subprocess.TimeoutExpired:
            # Test command run timed out
            logger_taskenv.error(f"[{self.testbed_name}] [{instance[KEY_INSTANCE_ID]}] Test script run time out {self.timeout}")
            with open(self.log_file, "a") as f:
                f.write(f"{TESTS_TIMEOUT} after {self.timeout} seconds\n")
            return False
        except Exception as e:
            # Test command run failed
            logger_taskenv.error(f"[{self.testbed_name}] [{instance[KEY_INSTANCE_ID]}] Test script run failed")
            with open(self.log_file, "a") as f:
                f.write(f"{TESTS_ERROR}: {e}")
            return False

    def __exit__(self, exc_type, exc_value, exc_traceback):
        os.chdir(self.cwd)

harness/context_manager.py
princeton-nlp/SWE-bench
import argparse, os, re

from constants import (
    APPLY_PATCH_FAIL,
    KEY_INSTANCE_ID,
    KEY_MODEL,
    KEY_PREDICTION,
)
from context_manager import TestbedContextManager, TaskEnvContextManager
from engine_validation import setup_testbed
from multiprocessing import Pool, cpu_count
from tqdm.auto import tqdm
from typing import Dict
from utils import (
    extract_minimal_patch,
    get_instances,
    split_instances,
    DotDict
)


def overwrite_ablation(tcm: TaskEnvContextManager, task_instance: Dict):
    """
    Code for running ablation experiment to compare generating full files vs patches

    Args:
        tcm: TaskEnvContextManager
        task_instance: Dict containing task instance
    """
    # if full output is none, write to log and skip altogether
    if 'full_output' not in task_instance:
        print(f'[{task_instance[KEY_INSTANCE_ID]}] No `full_output` field, skipping')
        with open(tcm.log_file, 'a') as f_log:
            f_log.write(f'{APPLY_PATCH_FAIL}; No `full_output` field\n')
        return
    if task_instance['full_output'] is None:
        print(f'[{task_instance[KEY_INSTANCE_ID]}] `full_output` is None, skipping')
        with open(tcm.log_file, 'a') as f_log:
            f_log.write(f'{APPLY_PATCH_FAIL}; `full_output` is None\n')
        return

    # Attempt to set up environment with task + apply test patch
    if not tcm.reset_task_env(task_instance):
        return
    
    filename_pat = re.compile(r'\[start of ([\w\.\-\/]+)\]\n(.+?)\n\[end of \1\]', re.DOTALL)
    # Run installation
    if (
        not tcm.run_install_task(task_instance)
        or not tcm.apply_patch(task_instance["test_patch"], patch_type="test")
    ):
        return
    
    # overwrite files
    for filename, contents in filename_pat.findall(task_instance['full_output']):
        correct_filename = './' + filename.lstrip('/')
        correct_filename = os.path.abspath(correct_filename)
        if not correct_filename.startswith(os.getcwd()):
            print(f"[{task_instance[KEY_INSTANCE_ID]}] Generation attempted to create file outside of working directory")
            return

        # if os.path.exists(correct_filename):
        if not os.path.exists(correct_filename):
            folder = '/'.join(correct_filename.split('/')[:-1])
            if not os.path.exists(folder):
                os.makedirs(folder)
        with open(correct_filename, 'w') as f:
            f.write(contents)
            with open(tcm.log_file, 'a') as f_log:
                f_log.write(f'Overwrote {correct_filename}\n')
    
    # run testing script
    if not tcm.run_tests_task(task_instance):
        return
    
    return


def evaluate_predictions(data: Dict):
    """
    Sets up task environment context manager. Each prediction is then
    evaluated within the context manager.

    Args:
        data: Dict containing task instances and other data
            task_instances: List of [task instance, prediction] pairs to evalute
            + setup_testbed args
    """
    data_dict = DotDict(data)
    for task_instance in tqdm(
        data_dict.task_instances,
        disable=data_dict.verbose,
        desc=f"Evaluating predictions for {data_dict.log_dir}"
    ):
        with TaskEnvContextManager(
            task_instance,
            data_dict.testbed,
            data_dict.venv,
            data_dict.log_dir,
            data_dict.conda_path,
            verbose=data_dict.verbose,
            timeout=data_dict.timeout,
            is_eval=True,
        ) as tcm:
            # Attempt to set up environment with task instance
            if not tcm.reset_task_env(task_instance):
                continue

            # Attempt to apply prediction
            patch_type = "pred_try"
            if not tcm.apply_patch(task_instance[KEY_PREDICTION], patch_type=patch_type) \
                and task_instance[KEY_PREDICTION] is not None:
                task_instance[KEY_PREDICTION] = extract_minimal_patch(task_instance[KEY_PREDICTION])
                patch_type = "pred_minimal_try"
                if not tcm.apply_patch(task_instance[KEY_PREDICTION], patch_type=patch_type):
                    continue
            tcm.apply_patch(task_instance[KEY_PREDICTION], patch_type=patch_type, revert=True)
            patch_type = patch_type.replace("_try", "")

            # Run installation + testing script
            if (
                not tcm.run_install_task(task_instance)
                or not tcm.apply_patch(task_instance["test_patch"], patch_type="test")
                or not tcm.apply_patch(task_instance[KEY_PREDICTION], patch_type=patch_type)
                or not tcm.run_tests_task(task_instance)
            ):
                continue


def main(args):
    """
    Splits predictions into multiple groups if num_workers > 1. Each group is
    then evaluated in parallel.
    """
    if args.num_workers is None:
        args.num_workers = cpu_count()

    predictions = get_instances(args.predictions_path)

    # Remove predictions that have already been evaluated
    if args.skip_existing:
        predictions_filtered = []
        for p in predictions:
            path_log = os.path.join(args.log_dir, f"{p[KEY_INSTANCE_ID]}.{p[KEY_MODEL]}.eval.log")
            if not os.path.exists(path_log):
                predictions_filtered.append(p)
        if len(predictions_filtered) == 0:
            return
        else:
            predictions = predictions_filtered

    predictions_groups = split_instances(predictions, args.num_workers)

    data_groups = [
        {
            "task_instances": g,
            "func": evaluate_predictions,
            **vars(args),
        }
        for g in predictions_groups
    ]

    if args.num_workers == 1:
        setup_testbed(data_groups[0])
        return

    pool = Pool(processes=args.num_workers)
    pool.map(setup_testbed, data_groups)
    pool.close()
    pool.join()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--predictions_path", type=str, help="Path to predictions instances file", required=True)
    parser.add_argument("--log_dir", type=str, help="Path to log directory", required=True)
    parser.add_argument("--num_workers", type=int, default=1, help="(Optional) Number of workers")
    parser.add_argument("--path_conda", type=str, help="(Optional) Path to miniconda3 or anaconda installation")
    parser.add_argument("--skip_existing", action="store_true", help="(Optional) Skip existing logs")
    parser.add_argument("--testbed", type=str, help="(Optional) Path to testbed directory")
    parser.add_argument("--temp_dir", type=str, help="(Optional) Path to temporary directory for storing virtual envs")
    parser.add_argument("--timeout", type=int, default=None, help="(Optional) Timeout (seconds) for testing script execution")
    parser.add_argument("--verbose", action="store_true", help="(Optional) Verbose mode")
    args = parser.parse_args()
    main(args)

harness/engine_evaluation.py
princeton-nlp/SWE-bench
import argparse, os

from context_manager import TaskEnvContextManager, TestbedContextManager
from multiprocessing import Pool, cpu_count
from typing import Dict
from utils import get_instances, split_instances, DotDict


SKIP_INSTANCES = {"pytest-dev/pytest": ["6387", "7956", "3805"]}


def validate_args(args):
    """
    Validation for command line arguments
    """
    if not os.path.exists(args.instances_path):
        raise ValueError(f"Could not find instances file at {args.instances_path}")
    if not os.path.exists(args.log_dir):
        raise ValueError(f"Could not find log directory at {args.log_dir}")
    if args.path_conda is not None and not os.path.exists(args.path_conda):
        raise ValueError(f"Could not find conda installation at {args.path_conda}")
    if args.testbed is not None and not os.path.exists(args.testbed):
        raise ValueError(f"Could not find testbed at {args.testbed}")
    if args.temp_dir is not None and not os.path.exists(args.temp_dir):
        raise ValueError(f"Could not find temporary directory at {args.temp_dir}")
    if args.timeout is not None and args.timeout < 0:
        raise ValueError(f"Timeout must be a positive integer")
    if args.num_workers is not None and args.num_workers < 1:
        raise ValueError(f"Number of workers must be a positive integer")


def verify_task_instances(data: Dict):
    """
    Sets up task environment context manager. Each task instance is then
    installed and validated within the context manager.

    Args:
        data: Dict containing task instances and other data
            task_instances: List of task instances
            + setup_testbed args
    """
    data_dict = DotDict(data)
    for task_instance in data_dict.task_instances:
        with TaskEnvContextManager(
            task_instance,
            data_dict.testbed,
            data_dict.venv,
            data_dict.log_dir,
            data_dict.conda_path,
            verbose=data_dict.verbose,
            timeout=data_dict.timeout,
        ) as tcm:
            if (
                task_instance["repo"] in SKIP_INSTANCES
                and task_instance["pull_number"]
                in SKIP_INSTANCES[task_instance["repo"]]
            ):
                continue
            if tcm.log_file_exists:
                print(
                    f"Skipping {task_instance['instance_id']}, log file ({tcm.log_file}) already exists"
                )
                continue
            if (
                not tcm.reset_task_env(task_instance)
                or not tcm.run_install_task(task_instance)
                or not tcm.apply_patch(task_instance["test_patch"], patch_type="test")
                or not tcm.run_tests_task(task_instance)
                or not tcm.apply_patch(task_instance["patch"], patch_type="gold")
                or not tcm.run_tests_task(task_instance)
            ):
                continue


def setup_testbed(data: Dict):
    """
    Creates testbed context manager and runs verify_task_instances in parallel

    Args:
        data: Dict containing task instances and other data
            task_instances: List of task instances
            log_dir: Path to log directory
            path_conda: Path to miniconda3 or anaconda installation
            testbed: Path to testbed directory
            temp_dir: Path to temporary directory for storing virtual envs
            timeout: Timeout (seconds) for testing script execution
            verbose: Verbose mode
    """
    data_dict = DotDict(data)
    with TestbedContextManager(
        data_dict.task_instances,
        data_dict.log_dir,
        path_conda=data_dict.path_conda,
        testbed=data_dict.testbed,
        temp_dir=data_dict.temp_dir,
        timeout=data_dict.timeout,
        verbose=data_dict.verbose,
    ) as tcm:
        distributed_task_list = tcm.get_distributed_tasks()
        for task_list in distributed_task_list:
            print(
                f"{task_list['testbed']}: {len(task_list['task_instances'])} instances"
            )

        if len(distributed_task_list) == 1:
            data_dict.func(distributed_task_list[0])
            return

        pool = Pool(processes=len(distributed_task_list))
        pool.map(data_dict.func, distributed_task_list)
        pool.close()
        pool.join()


def main(args):
    """
    Splits task instances into multiple groups if num_workers > 1
    """
    if args.num_workers is None:
        args.num_workers = cpu_count()

    task_instances = get_instances(args.instances_path)
    task_instances_groups = split_instances(task_instances, args.num_workers)

    data_groups = [
        {
            "task_instances": g,
            "func": verify_task_instances,
            **vars(args),
        }
        for g in task_instances_groups
    ]

    for group in data_groups:
        del group["instances_path"]

    if args.num_workers == 1:
        setup_testbed(data_groups[0])
        return

    pool = Pool(processes=args.num_workers)
    pool.map(setup_testbed, data_groups)
    pool.close()
    pool.join()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--instances_path", type=str, help="Path to candidate task instances file", required=True)
    parser.add_argument("--log_dir", type=str, help="Path to log directory", required=True)
    parser.add_argument("--path_conda", type=str, help="(Optional) Path to miniconda3 or anaconda installation")
    parser.add_argument("--testbed", type=str, help="(Optional) Path to testbed directory")
    parser.add_argument("--temp_dir", type=str, help="(Optional) Path to temporary directory for storing virtual envs")
    parser.add_argument("--timeout", type=int, default=None, help="(Optional) Timeout (seconds) for testing script execution")
    parser.add_argument("--verbose", action="store_true", help="(Optional) Verbose mode")
    parser.add_argument("--num_workers", type=int, default=None, help="(Optional) Number of workers")
    args = parser.parse_args()
    validate_args(args)
    main(args)

harness/engine_validation.py
princeton-nlp/SWE-bench
import argparse, json, logging, os, shutil

from constants import (
    KEY_INSTANCE_ID,
    KEY_MODEL,
    KEY_PREDICTION,
)
from engine_evaluation import main as eval_engine
from multiprocessing import Pool
from utils import get_instances

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("run_evaluation")


def validate_predictions(predictions_path, tasks_ids):
    # Check that predictions file exists
    if not any([predictions_path.endswith(x) for x in [".json", ".jsonl"]]):
        raise ValueError("Predictions path must be .json or .jsonl file")
    predictions = get_instances(predictions_path)
    not_in_tasks = []
    # Check that predictions are correctly formatted
    for pred in predictions:
        if any([x not in pred for x in [KEY_INSTANCE_ID, KEY_MODEL, KEY_PREDICTION]]):
            raise ValueError("Every prediction must have instance_id, model, and patch fields")
        if pred[KEY_INSTANCE_ID] not in tasks_ids:
            not_in_tasks.append(pred[KEY_INSTANCE_ID])
    # Check that instance IDs specified by predictions exist
    if len(not_in_tasks) > 0:
        logger.warning(
            "Predictions for the following instance_ids were not "
            + "found in the tasks file and will not be considered: "
            + ", ".join(not_in_tasks)
        )


def main(
    predictions_path: str,
    swe_bench_tasks: str,
    log_dir: str,
    testbed: str,
    skip_existing: bool,
    timeout: int,
    verbose: bool
):
    """
    Runs evaluation on predictions for each model/repo/version combination.

    Args:
        predictions_path (str): Path to the predictions file.
        swe_bench_tasks (str): Path to the SWE-bench tasks file.
        log_dir (str): Path to the directory where logs will be saved.
        testbed (str): Path to the directory where testbeds will be saved.
        skip_existing (bool): Whether to skip evaluations for predictions that already have logs.
        timeout (int): Timeout for each evaluation.
        verbose (bool): Whether to print verbose output.

    Raises:
        ValueError: If log_dir is not a directory, testbed is not a directory, or swe_bench_tasks does not exist.
    """
    # Validate arguments
    if not os.path.isdir(log_dir):
        raise ValueError("--log_dir must be a directory")
    if not os.path.isdir(testbed):
        raise ValueError("--testbed must be a directory")
    if not os.path.exists(swe_bench_tasks):
        raise ValueError("--swe_bench_tasks does not exist")
    tasks = json.load(open(swe_bench_tasks))
    tasks_map = {t[KEY_INSTANCE_ID]: t for t in tasks}
    validate_predictions(predictions_path, [t[KEY_INSTANCE_ID] for t in tasks])

    # Group predictions by model
    predictions = json.load(open(predictions_path))
    map_model_to_predictions = {}
    for p in predictions:
        model = p[KEY_MODEL]
        if model not in map_model_to_predictions:
            map_model_to_predictions[model] = []
        map_model_to_predictions[model].append(p)
    logger.info(f"Found {len(predictions)} predictions across {len(map_model_to_predictions)} model(s) in predictions file")
    
    # For each model, split predictions by repo + save to folder
    eval_args = []
    temp_dirs = []
    for model, predictions in map_model_to_predictions.items():
        # Group predictions by repository, version
        map_repo_version_to_predictions = {}
        for p in predictions:
            repo = p[KEY_INSTANCE_ID].rsplit("-", 1)[0]
            if repo not in map_repo_version_to_predictions:
                map_repo_version_to_predictions[repo] = {}
            t = tasks_map[p[KEY_INSTANCE_ID]]
            p.update(t)
            version = t["version"]
            if version not in map_repo_version_to_predictions[repo]:
                map_repo_version_to_predictions[repo][version] = []
            map_repo_version_to_predictions[repo][version].append(p)
        
        # For each model/repo/version, create testbed folder and save predictions
        for repo in map_repo_version_to_predictions:
            for version in map_repo_version_to_predictions[repo]:
                # Create testbed folder + file for model/repo/version specific predictions
                testbed_model_repo_version_dir = os.path.join(testbed, model, repo, version)
                os.makedirs(testbed_model_repo_version_dir, exist_ok=True)
                file_name = f"{model}_{repo}_{version}_{predictions_path.split('/')[-1]}"
                file_path = os.path.join(testbed_model_repo_version_dir, file_name)
                with open(file_path, "w") as f:
                    args = argparse.Namespace()
                    args.log_dir = os.path.join(log_dir, model)
                    args.temp_dir = testbed_model_repo_version_dir
                    args.num_workers = 1
                    args.timeout = timeout
                    args.skip_existing = skip_existing
                    args.verbose = verbose

                    repo_version_predictions = map_repo_version_to_predictions[repo][version]
                    if skip_existing:
                        # Skip logs that already exist
                        predictions_filtered = []
                        for p in repo_version_predictions:
                            log_file = os.path.join(
                                args.log_dir,
                                f"{p[KEY_INSTANCE_ID]}.{p[KEY_MODEL]}.eval.log",
                            )
                            if not os.path.exists(log_file):
                                predictions_filtered.append(p)
                        if len(predictions_filtered) == 0:
                            logger.info(f"[{model}/{repo}/{version}] All predictions already exist, skipping")
                            continue
                        else:
                            logger.info(
                                f"[{model}/{repo}/{version}] # of predictions to evaluate: {len(predictions_filtered)} " + 
                                f"({len(repo_version_predictions) - len(predictions_filtered)} already evaluated)"
                            )
                            repo_version_predictions = predictions_filtered
                    else:
                        logger.info(f"[{model}/{repo}/{version}] # of predictions to evaluate: {len(repo_version_predictions)}")

                    json.dump(repo_version_predictions, f, indent=4)
                    args.predictions_path = file_path

                    eval_args.append(args)
                temp_dirs.append(testbed_model_repo_version_dir)
    
    # Run evaluation on each model/repo
    pool = Pool(processes=len(eval_args))
    pool.map(eval_engine, eval_args)
    pool.close()
    pool.join()

    # Clean up
    for temp_dir in temp_dirs:
        shutil.rmtree(temp_dir)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--predictions_path", type=str, help="Path to predictions file (must be .json)", required=True)
    parser.add_argument("--log_dir", type=str, help="Path to log directory", required=True)
    parser.add_argument("--swe_bench_tasks", type=str, help="Path to SWE-bench task instances file", required=True)
    parser.add_argument("--testbed", type=str, help="Path to testbed directory", required=True)
    parser.add_argument("--skip_existing", action="store_true", help="(Optional) Skip existing logs")
    parser.add_argument("--timeout", type=int, help="(Optional) Timeout in seconds (default: 900)", default=900)
    parser.add_argument("--verbose", action="store_true", help="(Optional) Verbose mode")
    args = parser.parse_args()
    logger.propagate = args.verbose
    main(**vars(args))

harness/run_evaluation.py
princeton-nlp/SWE-bench
import json
import os
import re
import requests
import subprocess

from constants import (
    MAP_REPO_TO_REQS_PATHS,
    MAP_REPO_TO_ENV_YML_PATHS,
    SWE_BENCH_URL_RAW,
    NON_TEST_EXTS,
)
from datetime import datetime
from dotenv import load_dotenv
from git import Repo
from typing import Dict, List


load_dotenv()


def get_conda_env_names(conda_source: str) -> List:
    """
    Get list of conda environment names for given conda path

    Args:
        conda_source (str): Path to conda executable
    Returns:
        env_names (list): List of conda environment names
    """
    # Get list of conda environments
    conda_envs = subprocess.run(
        f"{conda_source} env list", shell=True, check=True, capture_output=True
    )
    output = conda_envs.stdout.decode("utf-8")
    lines = output.split("\n")
    # Store environment names to list
    env_names = []
    for line in lines:
        if line.startswith("#"):
            continue
        if line.strip() == "":
            continue
        if " " in line:
            env_name = line.split(" ")[0]
            env_names.append(env_name)
    return env_names


def get_environment_yml(instance: Dict, env_name: str, save_path: str = None) -> str:
    """
    Get environment.yml for given task instance

    Args:
        instance (dict): SWE Bench Task instance
        env_name (str): Rename retrieved environment.yml to this name
        save_path (str): If provided, save environment.yml to this path
    Returns:
        environment.yml (str): If save_path given, returns path to saved environment.yml.
            Otherwise, returns environment.yml as string
    """
    # Attempt to find environment.yml at each path based on task instance's repo
    path_worked = False
    for req_path in MAP_REPO_TO_ENV_YML_PATHS[instance["repo"]]:
        reqs_url = os.path.join(
            SWE_BENCH_URL_RAW, instance["repo"], instance['environment_setup_commit'], req_path
        )
        reqs = requests.get(reqs_url)
        if reqs.status_code == 200:
            path_worked = True
            break
    if not path_worked:
        print(
            f"Could not find environment.yml at paths {MAP_REPO_TO_ENV_YML_PATHS[instance['repo']]}"
        )
        return None

    lines = reqs.text.split("\n")
    cleaned = []
    for line in lines:
        # Rename environment to given name
        if line.startswith("name:"):
            cleaned.append(f"name: {env_name}")
            continue
        cleaned.append(line)

    # Return environment.yml as string if no save path given
    if save_path is None:
        return "\n".join(cleaned)

    # Save environment.yml to given path and return path
    path_to_reqs = os.path.join(save_path, "environment.yml")
    with open(path_to_reqs, "w") as f:
        f.write("\n".join(cleaned))
    return path_to_reqs


def get_instances(instance_path: str) -> List:
    """
    Get task instances from given path

    Args:
        instance_path (str): Path to task instances
    Returns:
        task_instances (list): List of task instances
    """
    if any([instance_path.endswith(x) for x in [".jsonl", ".jsonl.all"]]):
        task_instances = list()
        with open(instance_path) as f:
            for line in f.readlines():
                task_instances.append(json.loads(line))
        return task_instances

    with open(instance_path) as f:
        task_instances = json.load(f)
    return task_instances


def get_requirements(instance: Dict, save_path: str = None):
    """
    Get requirements.txt for given task instance

    Args:
        instance (dict): task instance
        save_path (str): If provided, save requirements.txt to this path
    Returns:
        requirements.txt (str): If save_path given, returns path to saved requirements.txt.
            Otherwise, returns requirements.txt as string
    """
    # Attempt to find requirements.txt at each path based on task instance's repo
    path_worked = False
    for req_path in MAP_REPO_TO_REQS_PATHS[instance["repo"]]:
        reqs_url = os.path.join(
            SWE_BENCH_URL_RAW, instance["repo"], instance['environment_setup_commit'], req_path
        )
        reqs = requests.get(reqs_url)
        if reqs.status_code == 200:
            path_worked = True
            break
    if not path_worked:
        print(
            f"Could not find requirements.txt at paths {MAP_REPO_TO_REQS_PATHS[instance['repo']]}"
        )
        return None

    lines = reqs.text
    original_req = []
    additional_reqs = []
    req_dir = "/".join(req_path.split("/")[:-1])
    exclude_line = lambda line: any(
        [line.strip().startswith(x) for x in ["-e .", "#", ".[test"]]
    )

    for line in lines.split("\n"):
        if line.strip().startswith("-r"):
            # Handle recursive requirements
            file_name = line[len("-r") :].strip()
            reqs_url = os.path.join(
                SWE_BENCH_URL_RAW,
                instance["repo"],
                instance["base_commit"],
                req_dir,
                file_name,
            )
            reqs = requests.get(reqs_url)
            if reqs.status_code == 200:
                for line_extra in reqs.text.split("\n"):
                    if not exclude_line(line_extra):
                        additional_reqs.append(line_extra)
        else:
            if not exclude_line(line):
                original_req.append(line)

    # Combine all requirements into single text body
    additional_reqs.append("\n".join(original_req))
    all_reqs = "\n".join(additional_reqs)

    if save_path is None:
        return all_reqs

    path_to_reqs = os.path.join(save_path, "requirements.txt")
    with open(path_to_reqs, "w") as f:
        f.write(all_reqs)
    return path_to_reqs


def get_test_directives(instance: Dict) -> List:
    """
    Get test directives from the test_patch of a task instance

    Args:
        instance (dict): task instance
    Returns:
        directives (list): List of test directives
    """
    # Get test directives from test patch and remove non-test files
    diff_pat = r"diff --git a/.* b/(.*)"
    test_patch = instance["test_patch"]
    directives = re.findall(diff_pat, test_patch)
    directives = [
        d for d in directives if not any(d.endswith(ext) for ext in NON_TEST_EXTS)
    ]

    # For Django tests, remove extension + "tests/" prefix and convert slashes to dots (module referencing)
    if instance["repo"] == "django/django":
        directives_transformed = []
        for d in directives:
            d = d[: -len(".py")] if d.endswith(".py") else d
            d = d[len("tests/") :] if d.startswith("tests/") else d
            d = d.replace("/", ".")
            directives_transformed.append(d)
        directives = directives_transformed

    return directives


def clone_repo(repo_name: str, path: str, token: str = None) -> bool:
    """
    Wrapper for cloning repo from swe-bench organization

    Args:
        repo_name (str): Name of repo to clone
        path (str): Path to clone repo to
        token (str): GitHub token to use for cloning
    Returns:
        success (bool): True if repo cloned successfully, False otherwise
    """
    try:
        if token is None:
            token = os.environ.get("GITHUB_TOKEN", "git")
        repo_url = (
            f"https://{token}@github.com/swe-bench/"
            + repo_name.replace("/", "__")
            + ".git"
        )
        Repo.clone_from(repo_url, path)
        return True
    except Exception as e:
        print(e)
        return False


def split_instances(input_list: List, n: int) -> List:
    """
    Split a list into n approximately equal length sublists

    Args:
        input_list (list): List to split
        n (int): Number of sublists to split into
    Returns:
        result (list): List of sublists
    """
    avg_length = len(input_list) // n
    remainder = len(input_list) % n
    result, start = [], 0

    for i in range(n):
        length = avg_length + 1 if i < remainder else avg_length
        sublist = input_list[start : start + length]
        result.append(sublist)
        start += length

    return result


def find_python_by_date(target_date, date_format="%Y%m%d"):
    """
    Find python version closest to given date

    Args:
        target_date (str): Date to find python version for
        date_format (str): Format of target_date
    Returns:
        python_version (str): Python version closest to target_date
    """
    # Make web request to versions + date page
    url = f"https://www.python.org/doc/versions/"
    response = requests.get(url)

    # Look for all matches
    pattern = r"Python (.*)</a>, documentation released on (.*)\.</"
    matches = re.findall(pattern, response.text)

    # Convert NL dates to date time format
    def convert_to_yyyymmdd(input_date):
        # Parse the input date string using datetime
        date_obj = datetime.strptime(input_date, date_format)
        # Format the date object into YYYYMMDD format
        return date_obj.strftime("%Y%m%d")

    version_to_date = [(match[0], convert_to_yyyymmdd(match[1])) for match in matches]

    # Find Python
    for x in version_to_date:
        if target_date >= x[1]:
            return x[0]
    return None


class DotDict:
    """
    Wrapper class for accessing dictionary keys as attributes
    """

    def __init__(self, data):
        self.data = data

    def __getattr__(self, key):
        return self.data.get(key)


### MARK - Patch Correction
PATCH_PATTERN = re.compile(
    r"(?:diff[\w\_\.\ \/\-]+\n)?\-\-\-\s+a\/(?:.*?)\n\+\+\+\s+b\/(?:.*?)(?=diff\ |\-\-\-\ a\/|\Z)",
    re.DOTALL,
)
PATCH_FILE_PATTERN = re.compile(r"\-\-\-\s+a\/(?:.+)\n\+\+\+\s+b\/(?:.+)")
PATCH_HUNK_PATTERN = re.compile(
    r"\@\@\s+\-(\d+),(\d+)\s+\+(\d+),(\d+)\s+\@\@(.+?)(?=diff\ |\-\-\-\ a\/|\@\@\ \-|\Z)",
    re.DOTALL,
)


def get_first_idx(charlist):
    """Get index of first occurence of "-" or "+" in charlist"""
    first_min = charlist.index("-") if "-" in charlist else len(charlist)
    first_plus = charlist.index("+") if "+" in charlist else len(charlist)
    return min(first_min, first_plus)


def get_last_idx(charlist):
    """Get index of last occurence of "-" or "+" in charlist"""
    char_idx = get_first_idx(charlist[::-1])
    last_idx = len(charlist) - char_idx
    return last_idx + 1


def strip_content(hunk):
    """Remove trailing non +/- lines and trailing whitespace per line per hunk"""
    first_chars = list(map(lambda x: None if not len(x) else x[0], hunk.split("\n")))
    first_idx = get_first_idx(first_chars)
    last_idx = get_last_idx(first_chars)
    new_lines = list(map(lambda x: x.rstrip(), hunk.split("\n")[first_idx:last_idx]))
    new_hunk = "\n" + "\n".join(new_lines) + "\n"
    return new_hunk, first_idx - 1


def get_hunk_stats(pre_start, pre_len, post_start, post_len, hunk, total_delta):
    """Recalculate hunk start/end position and diff delta"""
    stats = {"context": 0, "added": 0, "subtracted": 0}
    hunk = hunk.split("\n", 1)[-1].strip("\n")
    for line in hunk.split("\n"):
        if line.startswith("-"):
            stats["subtracted"] += 1
        elif line.startswith("+"):
            stats["added"] += 1
        else:
            stats["context"] += 1
    context = stats["context"]
    added = stats["added"]
    subtracted = stats["subtracted"]
    pre_len = context + subtracted
    post_start = pre_start + total_delta
    post_len = context + added
    total_delta = total_delta + (post_len - pre_len)
    return pre_start, pre_len, post_start, post_len, total_delta


def extract_minimal_patch(model_patch):
    """
    Wrapper function that takes hunk and
    * Removes trailing non +/- lines and trailing whitespace per line per hunk
    * Recalculates hunk start/end position and diff delta
    * Returns new patch
    """
    model_patch = model_patch.lstrip("\n")
    new_patch = ""
    for patch in PATCH_PATTERN.findall(model_patch):
        total_delta = 0
        patch_header = PATCH_FILE_PATTERN.findall(patch)[0]
        if patch_header:
            new_patch += patch_header + "\n"
        for hunk in PATCH_HUNK_PATTERN.findall(patch):
            pre_start, pre_len, post_start, post_len, content = hunk
            pre_start, pre_len, post_start, post_len, content = list(
                map(lambda x: int(x) if x.isnumeric() else x, hunk)
            )
            content, adjust_pre_start = strip_content(content)
            pre_start += adjust_pre_start
            pre_start, pre_len, post_start, post_len, total_delta = get_hunk_stats(
                pre_start, pre_len, post_start, post_len, content, total_delta
            )
            new_patch += (
                f"@@ -{pre_start},{pre_len} +{post_start},{post_len} @@{content}"
            )
    return new_patch

harness/utils.py
princeton-nlp/SWE-bench
import json
import os
import ast
import jedi
import traceback
import subprocess
import time
import shutil
import numpy as np
from multiprocessing import Process, Lock, Value, Manager, set_start_method
from pyserini.search.lucene import LuceneSearcher
from git import Repo
from pathlib import Path
from .utils import ContextManager, is_test, list_files
from tqdm.auto import tqdm
from argparse import ArgumentParser
from tempfile import TemporaryDirectory
import logging

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)


def file_name_and_contents(filename, relative_path):
    text = relative_path + "\n"
    with open(filename, "r") as f:
        text += f.read()
    return text


def file_name_and_documentation(filename, relative_path):
    text = relative_path + "\n"
    try:
        with open(filename, "r") as f:
            node = ast.parse(f.read())

        # Get module docstring
        data = ast.get_docstring(node)
        if data:
            text += f"{data}"

        # Walk through all nodes in the AST
        for child_node in ast.walk(node):
            if isinstance(
                child_node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)
            ):
                data = ast.get_docstring(child_node)
                if data:
                    text += f"\n\n{child_node.name}\n{data}"
    except Exception as e:
        logger.error(e)
        logger.error(f"Failed to parse file {str(filename)}. Using simple filecontent.")
        with open(filename, "r") as f:
            text += f.read()
    return text


def file_name_and_docs_jedi(filename, relative_path):
    text = relative_path + "\n"
    with open(filename, "r") as f:
        source_code = f.read()
    try:
        script = jedi.Script(source_code, path=filename)
        module = script.get_context()
        docstring = module.docstring()
        text += f"{module.full_name}\n"
        if docstring:
            text += f"{docstring}\n\n"
        abspath = Path(filename).absolute()
        names = [
            name
            for name in script.get_names(
                all_scopes=True, definitions=True, references=False
            )
            if not name.in_builtin_module()
        ]
        for name in names:
            try:
                origin = name.goto(follow_imports=True)[0]
                if origin.module_name != module.full_name:
                    continue
                if name.parent().full_name != module.full_name:
                    if name.type in {"statement", "param"}:
                        continue
                full_name = name.full_name
                text += f"{full_name}\n"
                docstring = name.docstring()
                if docstring:
                    text += f"{docstring}\n\n"
            except KeyboardInterrupt:
                raise KeyboardInterrupt
            except:
                continue
    except KeyboardInterrupt:
        raise KeyboardInterrupt
    except Exception as e:
        logger.error(e)
        logger.error(f"Failed to parse file {str(filename)}. Using simple filecontent.")
        text = f"{relative_path}\n{source_code}"
        return text
    return text


DOCUMENT_ENCODING_FUNCTIONS = {
    "file_name_and_contents": file_name_and_contents,
    "file_name_and_documentation": file_name_and_documentation,
    "file_name_and_docs_jedi": file_name_and_docs_jedi,
}


def clone_repo(repo, root_dir, token, from_swebench, thread_id):
    thread_prefix = f'(thread #{thread_id}) (pid {os.getpid()}) '
    repo_dir = os.path.join(root_dir, repo.replace("/", "__"))
    if not os.path.exists(repo_dir):
        if from_swebench:
            repo_url = (
                f"https://{token}@github.com/swe-bench/"
                + repo.replace("/", "__")
                + ".git"
            )
        else:
            repo_url = f"https://{token}@github.com/{repo}.git"
        logger.info(thread_prefix + f"Cloning {repo}")
        Repo.clone_from(repo_url, repo_dir)
    return repo_dir


def build_documents(repo_dir, commit, document_encoding_func):
    documents = dict()
    with ContextManager(repo_dir, commit) as cm:
        filenames = list_files(cm.repo_path, include_tests=False)
        for relative_path in filenames:
            filename = Path(relative_path).absolute()
            text = document_encoding_func(filename, relative_path)
            documents[relative_path] = text
    return documents


def make_index(repo_dir, root_dir, commit, document_encoding_func, python, thread_id, instance_id):
    index_path = Path(root_dir, str(instance_id), "index")
    if index_path.exists():
        return index_path
    thread_prefix = f'(thread #{thread_id}) (pid {os.getpid()}) '
    documents_path = Path(root_dir, instance_id, "documents.jsonl")
    if not documents_path.parent.exists():
        documents_path.parent.mkdir(parents=True)
    documents = build_documents(repo_dir, commit, document_encoding_func)
    with open(documents_path, "w") as docfile:
        for relative_path, contents in documents.items():
            print(
                json.dumps({"id": relative_path, "contents": contents}),
                file=docfile,
                flush=True,
            )
    output = subprocess.run(
        [
            python,
            "-m",
            "pyserini.index",
            "--collection",
            "JsonCollection",
            "--generator",
            "DefaultLuceneDocumentGenerator",
            "--threads",
            "2",
            "--input",
            documents_path.parent.as_posix(),
            "--index",
            index_path.as_posix(),
            "--storePositions",
            "--storeDocvectors",
            "--storeRaw",
        ],
        check=False,
        capture_output=True,
    )
    if output.returncode == 130:
        raise KeyboardInterrupt
    if output.returncode != 0:
        logger.error(thread_prefix + output.stderr.decode("utf-8"))
        raise Exception(thread_prefix + f"Failed to build index for {instance_id}")
        return None
    logger.info(thread_prefix + f"Finished indexing {index_path}")
    return index_path


def make_indexes(
    root_dir,
    instances,
    document_encoding_func,
    python,
    from_swebench,
    token,
    counter,
    counter_lock,
    thread_id,
    index_paths,
):
    tmp_dir = TemporaryDirectory(dir=root_dir)
    thread_prefix = f'(thread #{thread_id}) (pid {os.getpid()}) '
    for instance in instances:
        repo = instance["repo"]
        commit = instance["base_commit"]
        instance_id = instance["instance_id"]
        try:
            repo_dir = clone_repo(repo, tmp_dir.name, token, from_swebench, thread_id)
            index_path = make_index(
                repo_dir, root_dir, commit, document_encoding_func, python, thread_id, instance_id
            )
            if index_path is None:
                continue
            index_paths[instance_id] = index_path
        except KeyboardInterrupt:
            logger.info(thread_prefix + "KeyboardInterrupt received. Terminating.")
            break
        except:
            logger.error(
                thread_prefix + f"Failed to process {repo}/{commit} (instance {instance_id})"
            )
            logger.error(thread_prefix + traceback.format_exc())
            continue
        with counter_lock:
            counter.value += 1
    tmp_dir.cleanup()
    return index_paths


def get_remaining_instances(instances, output_file):
    instance_ids = set()
    remaining_instances = list()
    if output_file.exists():
        with open(output_file, "r") as f:
            for line in f:
                instance = json.loads(line)
                instance_id = instance["instance_id"]
                instance_ids.add(instance_id)
        logger.warning(
            f"Found {len(instance_ids)} existing instances in {output_file}. Will skip them."
        )
    else:
        output_file.parent.mkdir(parents=True, exist_ok=True)
        return instances
    for instance in instances:
        instance_id = instance["instance_id"]
        if instance_id not in instance_ids:
            remaining_instances.append(instance)
    return remaining_instances


def chunk_list(l, n):
    for i in range(0, len(l), n):
        yield l[i : i + n]


def monitor(counter, lock, total):
    pid = os.getpid()
    # if this doesn't disappear, kill the process
    with tqdm(total=total, desc=f'Monitoring (pid {pid})') as pbar:
        while counter.value < total:
            with lock:
                pbar.n = counter.value
                pbar.refresh()
            time.sleep(2)


def search(instance, index_path):
    try:
        instance_id = instance["instance_id"]
        searcher = LuceneSearcher(index_path.as_posix())
        cutoff = len(instance["problem_statement"])
        while True:
            try:
                hits = searcher.search(
                    instance["problem_statement"][:cutoff],
                    k=20,
                    remove_dups=True,
                )
            except Exception as e:
                if "maxClauseCount" in str(e):
                    cutoff = int(round(cutoff * 0.8))
                    continue
                else:
                    raise e
            break
        results = {"instance_id": instance_id, "hits": []}
        for hit in hits:
            results["hits"].append({"docid": hit.docid, "score": hit.score})
        return results
    except Exception as e:
        logger.error(
            f"Failed to process {instance_id}"
        )
        logger.error(traceback.format_exc())
        return None
    

def search_indexes(remaining_instance, output_file, all_index_paths):
    with open(output_file, "a") as out_file:
        for instance in tqdm(remaining_instance, desc="Retrieving"):
            instance_id = instance["instance_id"]
            if instance_id not in all_index_paths:
                continue
            index_path = all_index_paths[instance_id]
            results = search(instance, index_path)
            if results is None:
                continue
            print(json.dumps(results), file=out_file, flush=True)


def get_missing_ids(instances, output_file):
    with open(output_file, "r") as f:
        written_ids = set()
        for line in f:
            instance = json.loads(line)
            instance_id = instance["instance_id"]
            written_ids.add(instance_id)
    missing_ids = set()
    for instance in instances:
        instance_id = instance["instance_id"]
        if instance_id not in written_ids:
            missing_ids.add(instance_id)
    return missing_ids


def get_index_paths(
        remaining_instances,
        root_dir_name,
        document_encoding_func,
        python,
        from_swebench,
        token,
        num_workers,
):
    instances_chunks = list(
        chunk_list(remaining_instances, int(np.ceil(len(remaining_instances) / num_workers)))
    )
    total_work = len(remaining_instances)
    with Manager() as manager:
        counter = manager.Value("i", 0)
        counter_lock = manager.Lock()
        all_index_paths = manager.dict()
        processes = list()
        for ix, instances_chunk in enumerate(instances_chunks):
            p = Process(
                target=make_indexes,
                args=(
                    root_dir_name,
                    instances_chunk,
                    document_encoding_func,
                    python,
                    from_swebench,
                    token,
                    counter,
                    counter_lock,
                    ix,
                    all_index_paths,
                ),
            )
            processes.append(p)
        monitor_process = Process(
                target=monitor,
                args=(counter, counter_lock, total_work),
            )
        for p in processes:
            p.start()
        monitor_process.start()
        for p in processes:
            p.join()
        if monitor_process.is_alive():
            monitor_process.terminate()
    return all_index_paths

def get_root_dir(use_tmp, output_dir, document_encoding_style):
    if use_tmp:
        tmp_root = "/scratch" if os.path.exists("/scratch") else "/tmp"
        root_dir = TemporaryDirectory(dir=tmp_root)
        root_dir_name = root_dir.name
    else:
        root_dir = Path(output_dir, document_encoding_style + "_indexes")
        if not root_dir.exists():
            root_dir.mkdir(parents=True, exist_ok=True)
        root_dir_name = root_dir
    return root_dir, root_dir_name


def main(instances_files, document_encoding_style, from_swebench, output_dir, use_tmp, num_workers):
    set_start_method('fork')
    document_encoding_func = DOCUMENT_ENCODING_FUNCTIONS[document_encoding_style]
    token = os.environ.get("GITHUB_TOKEN", "git")
    instances = list()
    for instances_file in instances_files:
        instances += [json.loads(line) for line in open(instances_file)]
    python = subprocess.run("which python", shell=True, capture_output=True)
    python = python.stdout.decode("utf-8").strip()
    output_file = Path(output_dir, document_encoding_style + ".retrieval.jsonl")
    remaining_instances = get_remaining_instances(instances, output_file)
    root_dir, root_dir_name = get_root_dir(use_tmp, output_dir, document_encoding_style)
    all_index_paths = get_index_paths(
        remaining_instances,
        root_dir_name,
        document_encoding_func,
        python,
        from_swebench,
        token,
        num_workers,
    )
    logger.info(f"Finished indexing {len(all_index_paths)} instances")
    search_indexes(remaining_instances, output_file, all_index_paths)
    missing_ids = get_missing_ids(instances, output_file)
    logger.info(f"Missing indexes for {len(missing_ids)} instances.\n{missing_ids}")
    logger.info(f"Saved retrieval results to {output_file}")
    if isinstance(root_dir, TemporaryDirectory):
        root_dir.cleanup()


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("--instances_files", nargs="+", required=True)
    parser.add_argument(
        "--document_encoding_style",
        required=True,
        choices=DOCUMENT_ENCODING_FUNCTIONS.keys(),
    )
    parser.add_argument("--from_swebench", action="store_true")
    parser.add_argument("--use_tmp", action="store_true")
    parser.add_argument("--output_dir", default="./retreival_results")
    parser.add_argument("--num_workers", type=int, default=4)
    args = parser.parse_args()
    main(**vars(args))

inference/make_datasets/bm25_retrieval.py
princeton-nlp/SWE-bench
"""Provided a source (raw) directory and the final (eval) directory, create a training split by removing all instances that are in the final directory from the source directory.
"""

import json
import logging
import random
from argparse import ArgumentParser
from pathlib import Path
from datasets import Dataset, DatasetDict, load_dataset
from tqdm.auto import tqdm

try:
    from create_instance import PROMPT_FUNCTIONS, add_text_inputs
    from tokenize_dataset import TOKENIZER_FUNCS
    from utils import string_to_bool
else:
    from .create_instance import PROMPT_FUNCTIONS, add_text_inputs
    from .tokenize_dataset import TOKENIZER_FUNCS
    from .utils import string_to_bool

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)


def load_jsonl_file(filename):
    if type(filename) == str:
        filename = Path(filename)
    if filename.name.endswith(".jsonl") or filename.name.endswith(".jsonl.all"):
        with open(filename) as f:
            return [json.loads(line) for line in f]
    elif filename.name.endswith(".json"):
        with open(filename) as f:
            return json.load(f)
    else:
        raise ValueError(f"Unknown file type {filename}")


def instances_generator(files):
    all_data = list()
    for file in tqdm(files, desc="Loading instance files"):
        all_data.extend(load_jsonl_file(file))
    return all_data


def get_training_and_eval_instances(raw_files, test_dataset):
    logger.info("Loading instances")
    raw_instances = list(instances_generator(raw_files))
    final_instances = list(test_dataset["test"])
    eval_repos = {x["repo"] for x in final_instances}
    train_instances = [x for x in raw_instances if x["repo"] not in eval_repos]
    train_instances = list(sorted(train_instances, key=lambda x: x["instance_id"]))
    eval_instances = list(sorted(final_instances, key=lambda x: x["instance_id"]))
    logger.info(f"Found {len(train_instances)} training ids")
    logger.info(f"Found {len(eval_instances)} eval ids")
    return train_instances, eval_instances


def extract_fields(instance):
    instance_id = instance["instance_id"]
    if instance["text_inputs"] is None or instance["patch"] is None:
        print(f"No text for {instance_id}")
        return None
    text_inputs = instance["text_inputs"].strip() + "\n\n"
    if text_inputs is None or instance["patch"] is None:
        print(f"No inputs for {instance_id}")
        return None
    patch = "\n".join([f"<patch>", instance["patch"], "</patch>"])
    return {**instance, "text": text_inputs, "patch": patch}


def main(
    raw_dir,
    test_dataset,
    output_dir,
    validation_ratio,
    github_token,
    retrieval_dir,
    prompt_style,
    file_source,
    k,
    python_only_patch,
    skip_train,
    max_context_len,
    tokenizer_name,
):
    if max_context_len is not None:
        assert tokenizer_name is not None
    if not Path(output_dir).exists():
        Path(output_dir).mkdir(parents=True)
    output_file = f"swe_prs_clean.{prompt_style}__po-{int(python_only_patch)}__st-{int(skip_train)}__fs-{file_source}"
    if k is not None:
        output_file += f"__k-{k}"
    if max_context_len is not None:
        output_file += f"__mcc-{max_context_len}-{tokenizer_name}"
    output_file = Path(output_dir, output_file)
    if output_file.exists():
        logger.info(f"Found {output_file.absolute().as_posix()}. Skipping")
        return
    test_dataset = load_dataset(test_dataset)
    training_instances, eval_instances = get_training_and_eval_instances(
        list(Path(raw_dir).glob("*.all")) if not skip_train else [],
        test_dataset,
    )
    if skip_train:
        training_instances = list()
    training_instances = {x["instance_id"]: x for x in training_instances}
    eval_instances = {x["instance_id"]: x for x in eval_instances[:20]}
    add_text_inputs(
        training_instances,
        retrieval_dir,
        k,
        github_token,
        prompt_style,
        file_source,
        python_only=python_only_patch,
        max_context_len=max_context_len,
        tokenizer_name=tokenizer_name,
    )
    add_text_inputs(
        eval_instances,
        retrieval_dir,
        k,
        github_token,
        prompt_style,
        file_source,
        python_only=python_only_patch,
        max_context_len=max_context_len,
        tokenizer_name=tokenizer_name,
    )
    breakpoint()
    columns = [
        "instance_id",
        "repo",
        "base_commit",
        "problem_statement",
        "hints_text",
        "created_at",
        "patch",
        "test_patch",
        "version",
        "FAIL_TO_PASS",
        "PASS_TO_PASS",
        "environment_setup_commit",
        "text_inputs",
    ]
    data = {key: list() for key in columns}
    test_data = {key: list() for key in columns}
    for instance_id, instance in tqdm(
        training_instances.items(),
        total=len(training_instances),
        desc="Creating training split",
    ):
        datum = extract_fields(instance)
        if datum is None:
            continue
        for key in columns:
            data[key].append(instance[key] if key in instance else '')
    for instance_id, instance in tqdm(
        eval_instances.items(), total=len(eval_instances), desc="Creating test split"
    ):
        datum = extract_fields(instance)
        if datum is None:
            continue
        for key in columns:
            test_data[key].append(instance[key] if key in instance else '')
    logger.info(f"Found {len(data['instance_id'])} training ids")
    logger.info(f"Found {len(test_data['instance_id'])} eval ids")
    trainval_dataset = Dataset.from_dict(data)
    testval_dataset = Dataset.from_dict(test_data)
    random.seed(42)
    validation_ids = random.sample(
        range(len(trainval_dataset)),
        int(round(validation_ratio * len(trainval_dataset))),
    )
    validation_2_ids = random.sample(
        range(len(testval_dataset)),
        int(round(0.1 * len(testval_dataset))),
    )
    train_ids = list(set(range(len(trainval_dataset))) - set(validation_ids))
    test_ids = list(set(range(len(testval_dataset))) - set(validation_2_ids))
    train_dataset = trainval_dataset.select(train_ids)
    validation_dataset = trainval_dataset.select(validation_ids)
    test_dataset = testval_dataset.select(test_ids)
    dataset = DatasetDict(
        {
            "train": train_dataset,
            "validation": validation_dataset,
            "test": test_dataset,
        }
    )
    logger.info(
        f"Saving {len(dataset['train'])} training instances and {len(dataset['validation'])} validation instances to {output_file} and {len(dataset['test'])} test instances to {output_file}"
    )
    dataset.save_to_disk(output_file)
    logger.info(f"Finsihed saving to {output_file}")


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("--raw_dir", type=str, required=False, help="Raw task instance outputs from collect.build_dataset.py. If you want to skip the training set, don't provide this.")
    parser.add_argument("--test_dataset", type=str, default="princeton-nlp/SWE-bench", help="Dataset to use for test set from HuggingFace Datasets.")
    parser.add_argument("--output_dir", type=str, required=True, help="Path to the output directory.")
    parser.add_argument("--validation_ratio", type=float, default=0.01, help="Ratio of training instances to use for validation.")
    parser.add_argument("--github_token", type=str, help="GitHub token.")
    parser.add_argument(
        "--skip_train", type=string_to_bool, nargs="?", const=True, default=False, help="Whether to skip the training set."
    )
    parser.add_argument(
        "--retrieval_dir",
        type=str,
        help="Path to the directory where the retrieval results are stored.",
    )
    parser.add_argument(
        "--prompt_style",
        type=str,
        default="style-3",
        choices=PROMPT_FUNCTIONS.keys(),
        help="Prompt style to use. See create_instance.PROMPT_FUNCTIONS for details.",
    )
    parser.add_argument(
        "--file_source",
        type=str,
        default="oracle",
        choices=["oracle", "bm25", "all", "oracle-edits-only"],
        help="Where to get the files from.",
    )
    parser.add_argument(
        "--k",
        type=int,
        default=None,
        help="Maximum number of files to use for retrieval.",
    )
    parser.add_argument(
        "--python_only_patch",
        type=string_to_bool,
        default=False,
        required=True,
        const=True,
        nargs="?",
        help="Path to the output file.",
    )
    parser.add_argument("--max_context_len", type=int, default=None)
    parser.add_argument(
        "--tokenizer_name",
        type=str,
        default=None,
        choices=TOKENIZER_FUNCS.keys(),
    )
    main(**vars(parser.parse_args()))

inference/make_datasets/create_base_dataset.py
princeton-nlp/SWE-bench
import json
import logging
import os
import traceback
from copy import deepcopy
from pathlib import Path
from tempfile import TemporaryDirectory
import unidiff
from tqdm.auto import tqdm

try:
    from tokenize_dataset import TOKENIZER_FUNCS
    from utils import AutoContextManager, ingest_directory_contents
except:
    from .tokenize_dataset import TOKENIZER_FUNCS
    from .utils import AutoContextManager, ingest_directory_contents

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)


PATCH_EXAMPLE = """--- a/file.py
+++ b/file.py
@@ -1,27 +1,35 @@
 def euclidean(a, b):
-    while b:
-        a, b = b, a % b
-    return a
+    if b == 0:
+        return a
+    return euclidean(b, a % b)
 
 
 def bresenham(x0, y0, x1, y1):
     points = []
     dx = abs(x1 - x0)
     dy = abs(y1 - y0)
-    sx = 1 if x0 < x1 else -1
-    sy = 1 if y0 < y1 else -1
-    err = dx - dy
+    x, y = x0, y0
+    sx = -1 if x0 > x1 else 1
+    sy = -1 if y0 > y1 else 1
 
-    while True:
-        points.append((x0, y0))
-        if x0 == x1 and y0 == y1:
-            break
-        e2 = 2 * err
-        if e2 > -dy:
+    if dx > dy:
+        err = dx / 2.0
+        while x != x1:
+            points.append((x, y))
             err -= dy
-            x0 += sx
-        if e2 < dx:
-            err += dx
-            y0 += sy
+            if err < 0:
+                y += sy
+                err += dx
+            x += sx
+    else:
+        err = dy / 2.0
+        while y != y1:
+            points.append((x, y))
+            err -= dx
+            if err < 0:
+                x += sx
+                err += dy
+            y += sy
 
+    points.append((x, y))
     return points"""


FULL_GENERATION_EXAMPLE = """[start of /src/this_file.py]
import os

def euclidean(a, b):
    if b == 0:
        return a
    return euclidean(b, a % b)
[end of /src/this_file.py]
[start of /src/another_file.py]
def bresenham(x0, y0, x1, y1):
    points = []
    dx = abs(x1 - x0)
    dy = abs(y1 - y0)
    x, y = x0, y0
    sx = -1 if x0 > x1 else 1
    sy = -1 if y0 > y1 else 1
    if dx > dy:
        err = dx / 2.0
        while x != x1:
            points.append((x, y))
            err -= dy
            if err < 0:
                y += sy
                err += dx
            x += sx
    else:
        err = dy / 2.0
        while y != y1:
            points.append((x
            err -= dx
            if err < 0:
                x += sx
                err += dy
            y += sy
    points.append((x, y))
    return points
[end of /src/another_file.py]"""


def add_lines_list(content):
    content_with_lines = list()
    for ix, line in enumerate(content.split("\n"), start=1):
        content_with_lines.append(f"{ix} {line}")
    return content_with_lines


def add_lines(content):
    return "\n".join(add_lines_list(content))


def make_code_text(files_dict, add_line_numbers=True):
    all_text = ""
    for filename, contents in sorted(files_dict.items()):
        all_text += f"[start of {filename}]\n"
        if add_line_numbers:
            all_text += add_lines(contents)
        else:
            all_text += contents
        all_text += f"\n[end of {filename}]\n"
    return all_text.strip("\n")


def make_code_text_edits_only(files_dict, patch, add_line_numbers=True):
    files = dict()
    patch = unidiff.PatchSet(patch)
    for patched_file in patch:
        source_file = patched_file.source_file.split("a/", 1)[-1]
        files[source_file] = list()
        for hunk in patched_file:
            start = hunk.source_start - 15
            end = start + hunk.source_length + 15
            files[source_file].append((start, end))
    all_text = ""
    for filename, content in files_dict.items():
        all_text += f"[start of {filename}]\n"
        content_with_lines = add_lines_list(content)
        for start, end in files[filename]:
            if start > 0:
                all_text += "...\n"
            all_text += "\n".join(content_with_lines[start:end])
            all_text += "\n"
            if end < len(content_with_lines):
                all_text += "...\n"
        all_text = all_text.strip("\n")
        all_text += f"\n[end of {filename}]\n"
    return all_text.strip("\n")


def prompt_style_1(instance):
    premise = "You will be provided with a partial code base and an issue statement explaining a problem to resolve."
    readmes_text = make_code_text(instance["readmes"])
    code_text = make_code_text(instance["file_contents"])
    instructions = (
        f"Now I need you to help solve this issue by generating a single patch file that I "
        + f"can apply directly to this repository using git apply. Please respond with a sing"
        + f"le patch file in the following format exactly."
    )
    problem_statement = instance["problem_statement"]
    final_text = [
        premise,
        "",
        "<issue>",
        problem_statement,
        "</issue>",
        "",
        "<code>",
        readmes_text,
        code_text,
        "</code>",
        "",
        instructions,
        "<patch>",
        PATCH_EXAMPLE,
        "</patch>",
    ]
    final_text = "\n".join(final_text)
    return final_text


def prompt_style_2(instance):
    premise = "You will be provided with a partial code base and an issue statement explaining a problem to resolve."
    readmes_text = make_code_text(instance["readmes"])
    code_text = make_code_text(instance["file_contents"])
    instructions = (
        f"I need you to solve this issue by generating a single patch file that I can apply "
        + f"directly to this repository using git apply. Please respond with a single patch "
        + f"file in the following format."
    )
    problem_statement = instance["problem_statement"]
    final_text = [
        premise,
        "<issue>",
        problem_statement,
        "</issue>",
        "<code>",
        readmes_text,
        code_text,
        "</code>",
        instructions,
        "<patch>",
        PATCH_EXAMPLE,
        "</patch>",
    ]
    final_text = "\n".join(final_text)
    return final_text


def prompt_style_2_edits_only(instance):
    premise = "You will be provided with a partial code base and an issue statement explaining a problem to resolve."
    readmes_text = make_code_text(instance["readmes"])
    code_text = make_code_text_edits_only(instance["file_contents"], instance["patch"])
    instructions = (
        f"I need you to solve this issue by generating a single patch file that I can apply "
        + f"directly to this repository using git apply. Please respond with a single patch "
        + f"file in the following format."
    )
    problem_statement = instance["problem_statement"]
    final_text = [
        premise,
        "<issue>",
        problem_statement,
        "</issue>",
        "<code>",
        readmes_text,
        code_text,
        "</code>",
        instructions,
        "<patch>",
        PATCH_EXAMPLE,
        "</patch>",
    ]
    final_text = "\n".join(final_text)
    return final_text


def prompt_style_3(instance):
    premise = "You will be provided with a partial code base and an issue statement explaining a problem to resolve."
    readmes_text = make_code_text(instance["readmes"])
    code_text = make_code_text(instance["file_contents"])
    example_explanation = (
        f"Here is an example of a patch file. It consists of changes to the code base. "
        + f"It specifies the file names, the line numbers of each change, and the removed and added lines. "
        + f"A single patch file can contain changes to multiple files."
    )
    final_instruction = (
        f"I need you to solve the provded issue by generating a single patch file that I can apply "
        + f"directly to this repository using git apply. Please respond with a single patch "
        + f"file in the format shown above."
    )
    problem_statement = instance["problem_statement"]
    final_text = [
        premise,
        "<issue>",
        problem_statement,
        "</issue>",
        "",
        "<code>",
        readmes_text,
        code_text,
        "</code>",
        "",
        example_explanation,
        "<patch>",
        PATCH_EXAMPLE,
        "</patch>",
        "",
        final_instruction,
        "Respond below:",
    ]
    final_text = "\n".join(final_text)
    return final_text


def prompt_style_4(instance):
    premise = "You will be provided with a partial code base and an issue statement explaining a problem to resolve."
    readmes_text = make_code_text(instance["readmes"], add_line_numbers=False)
    code_text = make_code_text(instance["file_contents"], add_line_numbers=False)
    instructions = (
        f"I need you to solve this issue by regenerating the full files in the code base that you would like to change. "
        + f"You can change as many files as you like. "
        + f"Please respond with a list of files and their revised contents in the following format."
    )
    problem_statement = instance["problem_statement"]
    final_text = [
        premise,
        "<issue>",
        problem_statement,
        "</issue>",
        "<code>",
        readmes_text,
        code_text,
        "</code>",
        instructions,
        "<example>",
        FULL_GENERATION_EXAMPLE,
        "</example>",
    ]
    final_text = "\n".join(final_text)
    return final_text


def ingest_files(filenames):
    files_dict = dict()
    for filename in filenames:
        with open(filename) as f:
            content = f.read()
        files_dict[filename] = content
    return files_dict


PROMPT_FUNCTIONS = {
    "style-1": prompt_style_1,
    "style-2": prompt_style_2,
    "style-3": prompt_style_3,
    "style-4": prompt_style_4,
    "style-2-edits-only": prompt_style_2_edits_only,
}


def add_retrieval_results(input_instances, retrieval_dir, k, file_source):
    retrieval_results = dict()
    for instance_id, instance in tqdm(
        input_instances.items(),
        total=len(input_instances),
        desc="Adding retrieval results",
    ):
        retrieval_results_path = Path(
            retrieval_dir,
            instance["repo"].split("/")[-1] + "-task-instances.retrieval.jsonl",
        )
        assert (
            retrieval_results_path.exists()
        ), f"Retrieval results not found at {retrieval_results_path}"
        if retrieval_results_path not in retrieval_results:
            d = [json.loads(line) for line in open(retrieval_results_path)]
            d = {x["instance_id"]: x["hits"] for x in d}
            retrieval_results[retrieval_results_path.as_posix()] = d
        instance["hits"] = retrieval_results[retrieval_results_path.as_posix()][
            instance_id
        ][:k]


def get_oracle_filenames(instance, python_only=True):
    source_files = {
        patch_file.source_file.split("a/", 1)[-1]
        for patch_file in unidiff.PatchSet(instance["patch"])
    }
    gold_docs = set()
    for source_file in source_files:
        if python_only and not source_file.endswith(".py"):
            continue
        gold_docs.add(source_file)
    return gold_docs


def add_text_inputs(
    input_instances,
    retrieval_dir,
    k,
    github_token,
    prompt_style,
    file_source,
    max_context_len=None,
    tokenizer_name=None,
    python_only=True,
    verbose=False,
):
    """Adds text inputs context for prediction in-place.

    Args:
    - input_instances: dictionary with unprocessed input instances.
    - retrieval_dir: if using retrieval method for file_contents, specify retrieval_dir to add retrieval results
    - k: if using retrieval, specifies the maximum number of files to included within context
    - github_token: github token to use for cloning private directories
    - prompt_style: specify the function to generate instructions and prompt provided an instance (from PROMPT_FUNCTIONS)
    - file_source: where to collect file_contents (e.g. oracle or bm25)
    - python_only: set to True for evaluation but false for training
    - verbose: set ContextManager verbose to True
    """
    if max_context_len is not None:
        assert (
            tokenizer_name is not None
        ), "Must specify tokenizer_name if using max_context_len"
        tokenizer, tokenizer_func = TOKENIZER_FUNCS[tokenizer_name]
    if github_token is None:
        github_token = os.environ.get("GITHUB_TOKEN", None)
    input_instances_copy = deepcopy(input_instances)
    if file_source in {"bm25"}:
        add_retrieval_results(input_instances_copy, retrieval_dir, k, file_source)
    orig_dir = os.getcwd()
    with TemporaryDirectory(
        dir="/scratch" if os.path.exists("/scratch") else "/tmp"
    ) as root_dir:
        for instance_id, instance in tqdm(
            input_instances_copy.items(),
            total=len(input_instances_copy),
            desc="Adding text inputs",
        ):
            try:
                with AutoContextManager(
                    instance, root_dir, token=github_token, verbose=verbose
                ) as cm:
                    readmes = cm.get_readme_files()
                    instance["readmes"] = ingest_files(readmes)
                    if max_context_len is not None:
                        instance["file_contents"] = dict()
                        base_text_inputs = PROMPT_FUNCTIONS[prompt_style](instance)
                        base_text_input_length = len(
                            tokenizer_func(base_text_inputs, tokenizer)
                        )
                    if file_source in {"oracle"}:
                        instance["file_contents"] = ingest_files(
                            get_oracle_filenames(instance, python_only=python_only)
                        )
                    elif file_source in {"bm25"}:
                        instance["file_contents"] = ingest_files(
                            [x["docid"] for x in instance["hits"]]
                        )
                    elif file_source in {"all"}:
                        instance["file_contents"] = ingest_directory_contents(
                            cm.repo_path
                        )
                    elif file_source in {"none"}:
                        instance["file_contents"] = dict()
                    else:
                        raise ValueError(f"Invalid file source {file_source}")
                    if max_context_len is not None:
                        cur_input_len = base_text_input_length
                        include_files = list()
                        for filename in [x["docid"] for x in instance["hits"]]:
                            content = make_code_text(
                                {filename: instance["file_contents"][filename]}
                            )
                            if tokenizer_name in {"llama"}:
                                tokens = tokenizer_func("\n" + content, tokenizer)
                                idx = tokens.index(13)
                                assert (
                                    idx <= 2
                                ), "Expected newline token id (13) to be one of the first three tokens"
                                tokens = tokens[idx + 1 :]  # remove newline tokens
                            else:
                                tokens = tokenizer_func(content, tokenizer)
                            if cur_input_len + len(tokens) < max_context_len:
                                include_files.append(filename)
                                cur_input_len += len(tokens)
                        instance["file_contents"] = {
                            filename: instance["file_contents"][filename]
                            for filename in include_files
                        }
                    input_instances[instance_id]["text_inputs"] = PROMPT_FUNCTIONS[
                        prompt_style
                    ](instance)
            except Exception as e:
                print(f"Failed on instance {instance_id}", e)
                traceback.print_exc()
                input_instances[instance_id]["text_inputs"] = None
            finally:
                # if AutoContextManager fails to exit properly future exits will return the wrong directory
                os.chdir(orig_dir)
    os.chdir(orig_dir)

inference/make_datasets/create_instance.py
princeton-nlp/SWE-bench
"""Provided a source (raw) directory and the final (eval) directory, create a training split by removing all instances that are in the final directory from the source directory.
"""

import logging
from argparse import ArgumentParser
from pathlib import Path

import tiktoken
from datasets import disable_caching, load_from_disk
from tqdm.auto import tqdm
from transformers import LlamaTokenizer

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)
logger.warning("Disabling caching")
disable_caching()


def cl100k(text, tokenizer):
    return tokenizer.encode(text, disallowed_special=())


def llama(text, tokenizer):
    return tokenizer(text, add_special_tokens=False, return_attention_mask=False)[
        "input_ids"
    ]


TOKENIZER_FUNCS = {
    "cl100k": (tiktoken.get_encoding("cl100k_base"), cl100k),
    "llama": (LlamaTokenizer.from_pretrained("togethercomputer/LLaMA-2-7B-32K"), llama),
}


def extract_fields(instance, tokenizer_name, tokenizer, tokenizer_func, eos_token):
    instance_id = instance["instance_id"]
    if instance["text"] is None or instance["patch"] is None:
        print(f"No text for {instance_id}")
        return {"input_ids": [], "labels": [], "text": "", "patch": ""}
    text_inputs = instance["text"].strip() + "\n"
    if text_inputs is None or instance["patch"] is None:
        print(f"No inputs for {instance_id}")
        return None
    patch = instance["patch"].strip()
    if len(eos_token) > 0:
        patch += f"\n{eos_token}"
    input_ids = tokenizer_func(text_inputs, tokenizer)
    if tokenizer_name in {"llama"}:
        label_ids = tokenizer_func(
            "\n" + patch, tokenizer
        )  # add newline to tokenize patch
        idx = label_ids.index(13)
        assert (
            idx <= 2
        ), "Expected newline token id (13) to be one of the first three tokens"
        label_ids = label_ids[idx + 1 :]  # remove newline tokens
    else:
        label_ids = tokenizer_func(patch, tokenizer)
    inputs = input_ids + label_ids[:-1]
    cond_len = len(input_ids) - 1
    labels = [-100] * cond_len + label_ids
    assert len(inputs) == len(labels)
    return {**instance, "input_ids": inputs, "labels": labels, "text": text_inputs, "patch": patch}


def extract_test_fields(instance, tokenizer_name, tokenizer, tokenizer_func, eos_token):
    instance_id = instance["instance_id"]
    if instance["text"] is None or instance["patch"] is None:
        print(f"No text for {instance_id}")
        return None
    text_inputs = instance["text"].strip() + "\n"
    if text_inputs is None or instance["patch"] is None:
        print(f"No inputs for {instance_id}")
        return None
    patch = instance["patch"].strip()
    if len(eos_token) > 0:
        patch += f"\n{eos_token}"
    input_ids = tokenizer_func(text_inputs, tokenizer)
    label_ids = tokenizer_func(patch, tokenizer)
    inputs = input_ids
    labels = label_ids
    return {**instance, "input_ids": inputs, "labels": labels, "text": text_inputs, "patch": patch}


def add_columns_from_dict(dataset, dict_columns):
    """dict_columns is a list of dicts with keys that are columns in dataset"""
    for column in dict_columns[0].keys():
        values = [d[column] for d in dict_columns]
        if column in dataset.column_names:
            dataset = dataset.remove_columns(column)
        dataset = dataset.add_column(column, values)
    return dataset


def main(
    dataset_path,
    output_dir,
    max_length,
    tokenizer_name,
    num_proc,
    no_map,
):
    if not Path(output_dir).exists():
        Path(output_dir).mkdir(parents=True)

    if tokenizer_name is not None:
        tokenizer, tokenizer_func = TOKENIZER_FUNCS[tokenizer_name]
        eos_token = getattr(tokenizer, "eos_token", "")

    dataset = load_from_disk(dataset_path)
    dataset = dataset.filter(lambda x: len(x["text"]) <= 5_000_000)
    for split in ["train", "validation"]:
        if split not in dataset:
            logger.warning(f"Split {split} not in dataset. Skipping")
            continue
        if not no_map:
            dataset[split] = dataset[split].map(
                lambda instance: extract_fields(
                    instance,
                    tokenizer_name,
                    tokenizer,
                    tokenizer_func,
                    eos_token,
                ),
                num_proc=num_proc,
                batched=False,
                desc=f"Tokenizing {split}",
            )
        elif len(dataset[split]) > 0:
            new_values = list(
                map(
                    lambda x: extract_fields(
                        x, tokenizer_name, tokenizer, tokenizer_func, eos_token
                    ),
                    tqdm(
                        dataset[split],
                        total=len(dataset[split]),
                        desc=f"Tokenizing {split}",
                    ),
                )
            )
            dataset[split] = add_columns_from_dict(dataset[split], new_values)
    for split in ["minitest", "test"]:
        if split not in dataset:
            logger.warning(f"Split {split} not in dataset. Skipping")
            continue
        if not no_map:
            dataset[split] = dataset[split].map(
                lambda instance: extract_test_fields(
                    instance,
                    tokenizer_name,
                    tokenizer,
                    tokenizer_func,
                    eos_token,
                ),
                num_proc=num_proc,
                batched=False,
                desc=f"Tokenizing {split}",
            )
        elif len(dataset[split]) > 0:
            new_values = list(
                map(
                    lambda x: extract_test_fields(
                        x, tokenizer_name, tokenizer, tokenizer_func, eos_token
                    ),
                    tqdm(
                        dataset[split],
                        total=len(dataset[split]),
                        desc=f"Tokenizing {split}",
                    ),
                )
            )
            dataset[split] = add_columns_from_dict(dataset[split], new_values)
    if max_length is not None:
        for split in ["train", "validation"]:
            if split not in dataset:
                logger.warning(f"Split {split} not in dataset. Skipping")
                continue
        logger.warning(f"Filtering {split} to max length {max_length}")
        dataset[split] = dataset[split].filter(
            lambda instance: len(instance["input_ids"]) <= max_length,
            num_proc=num_proc,
            batched=False,
        )
    output_file = Path(dataset_path).name + f"__tok-{tokenizer_name}"
    if max_length is not None:
        output_file += f"__max-{max_length}"
    output_file = Path(output_dir) / output_file
    logger.warning(f"Saving to {output_file}")
    dataset.save_to_disk(output_file)


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("--dataset_path", type=str, required=True)
    parser.add_argument("--output_dir", type=str, required=True)
    parser.add_argument("--max_length", type=int, default=None)
    parser.add_argument(
        "--tokenizer_name", type=str, required=True, choices=TOKENIZER_FUNCS.keys()
    )
    parser.add_argument("--num_proc", type=int, default=5)
    parser.add_argument("--no_map", type=bool, default=False)
    main(**vars(parser.parse_args()))

inference/make_datasets/tokenize_dataset.py
princeton-nlp/SWE-bench
import os
import re
import ast
import chardet
import subprocess
from argparse import ArgumentTypeError
from git import Repo
from pathlib import Path
from tempfile import TemporaryDirectory


DIFF_PATTERN = re.compile(r"^diff(?:.*)")
PATCH_PATTERN = re.compile(
    r"(?:diff[\w\_\.\ \/\-]+\n)?\-\-\-\s+a\/(?:.*?)\n\+\+\+\s+b\/(?:.*?)(?=diff\ |\-\-\-\ a\/|\Z)",
    re.DOTALL,
)
PATCH_FILE_PATTERN = re.compile(r"\-\-\-\s+a\/(?:.+)\n\+\+\+\s+b\/(?:.+)")
PATCH_HUNK_PATTERN = re.compile(
    r"\@\@\s+\-(\d+),(\d+)\s+\+(\d+),(\d+)\s+\@\@(.+?)(?=diff\ |\-\-\-\ a\/|\@\@\ \-|\Z)",
    re.DOTALL,
)


def get_first_idx(charlist):
    first_min = charlist.index('-')  if '-' in charlist else len(charlist)
    first_plus = charlist.index('+') if '+' in charlist else len(charlist)
    return min(first_min, first_plus)

def get_last_idx(charlist):
    char_idx = get_first_idx(charlist[::-1])
    last_idx = len(charlist) - char_idx
    return last_idx + 1

def strip_content(hunk):
    first_chars = list(map(lambda x: None if not len(x) else x[0], hunk.split('\n')))
    first_idx = get_first_idx(first_chars)
    last_idx = get_last_idx(first_chars)
    new_lines = list(map(lambda x: x.rstrip(), hunk.split('\n')[first_idx:last_idx]))
    new_hunk = '\n' + '\n'.join(new_lines) + '\n'
    return new_hunk, first_idx - 1


def get_hunk_stats(pre_start, pre_len, post_start, post_len, hunk, total_delta):
    stats = {"context": 0, "added": 0, "subtracted": 0}
    hunk = hunk.split("\n", 1)[-1].strip("\n")
    for line in hunk.split("\n"):
        if line.startswith("-"):
            stats["subtracted"] += 1
        elif line.startswith("+"):
            stats["added"] += 1
        else:
            stats["context"] += 1
    context = stats["context"]
    added = stats["added"]
    subtracted = stats["subtracted"]
    pre_len = context + subtracted
    post_start = pre_start + total_delta
    post_len = context + added
    total_delta = total_delta + (post_len - pre_len)
    return pre_start, pre_len, post_start, post_len, total_delta


def repair_patch(model_patch):
    if model_patch is None:
        return None
    model_patch = model_patch.lstrip("\n")
    new_patch = ""
    for patch in PATCH_PATTERN.findall(model_patch):
        total_delta = 0
        diff_header = DIFF_PATTERN.findall(patch)
        if diff_header:
            new_patch += diff_header[0] + "\n"
        patch_header = PATCH_FILE_PATTERN.findall(patch)[0]
        if patch_header:
            new_patch += patch_header + "\n"
        for hunk in PATCH_HUNK_PATTERN.findall(patch):
            pre_start, pre_len, post_start, post_len, content = hunk
            pre_start, pre_len, post_start, post_len, total_delta = get_hunk_stats(
                *list(map(lambda x: int(x) if x.isnumeric() else x, hunk)), total_delta
            )
            new_patch += (
                f"@@ -{pre_start},{pre_len} +{post_start},{post_len} @@{content}"
            )
    return new_patch


def extract_minimal_patch(model_patch):
    model_patch = model_patch.lstrip("\n")
    new_patch = ""
    for patch in PATCH_PATTERN.findall(model_patch):
        total_delta = 0
        diff_header = DIFF_PATTERN.findall(patch)
        patch_header = PATCH_FILE_PATTERN.findall(patch)[0]
        if patch_header:
            new_patch += patch_header + "\n"
        for hunk in PATCH_HUNK_PATTERN.findall(patch):
            pre_start, pre_len, post_start, post_len, content = hunk
            pre_start, pre_len, post_start, post_len, content = list(map(lambda x: int(x) if x.isnumeric() else x, hunk))
            content, adjust_pre_start = strip_content(content)
            pre_start += adjust_pre_start
            pre_start, pre_len, post_start, post_len, total_delta = get_hunk_stats(
                pre_start, pre_len, post_start, post_len, content, total_delta
            )
            new_patch += (
                f"@@ -{pre_start},{pre_len} +{post_start},{post_len} @@{content}"
            )
    return new_patch


def extract_diff(response):
    """
    Extracts the diff from a response formatted in different ways
    """
    if response is None:
        return None
    diff_matches = []
    other_matches = []
    pattern = re.compile(r"\<([\w-]+)\>(.*?)\<\/\1\>", re.DOTALL)
    for code, match in pattern.findall(response):
        if code in {"diff", "patch"}:
            diff_matches.append(match)
        else:
            other_matches.append(match)
    pattern = re.compile(r"```(\w+)?\n(.*?)```", re.DOTALL)
    for code, match in pattern.findall(response):
        if code in {"diff", "patch"}:
            diff_matches.append(match)
        else:
            other_matches.append(match)
    if diff_matches:
        return diff_matches[0]
    if other_matches:
        return other_matches[0]
    return response.split("</s>")[0]


def is_test(name, test_phrases=None):
    if test_phrases is None:
        test_phrases = ["test", "tests", "testing"]
    words = set(re.split(r" |_|\/|\.", name.lower()))
    return any(word in words for word in test_phrases)


class ContextManager:
    def __init__(self, repo_path, base_commit, verbose=False):
        self.repo_path = Path(repo_path).resolve().as_posix()
        self.old_dir = os.getcwd()
        self.base_commit = base_commit
        self.verbose = verbose

    def __enter__(self):
        os.chdir(self.repo_path)
        cmd = f"git reset --hard {self.base_commit} && git clean -fdxq"
        if self.verbose:
            subprocess.run(cmd, shell=True, check=True)
        else:
            subprocess.run(
                cmd,
                shell=True,
                check=True,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
        return self

    def get_environment(self):
        raise NotImplementedError()  # TODO: activate conda environment and return the environment file

    def get_readme_files(self):
        files = os.listdir(self.repo_path)
        files = list(filter(lambda x: os.path.isfile(x), files))
        files = list(filter(lambda x: x.lower().startswith("readme"), files))
        return files

    def __exit__(self, exc_type, exc_val, exc_tb):
        os.chdir(self.old_dir)


class AutoContextManager(ContextManager):
    """Automatically clones the repo if it doesn't exist"""

    def __init__(self, instance, root_dir=None, verbose=False, token="git"):
        self.tempdir = None
        if root_dir is None:
            self.tempdir = TemporaryDirectory()
            root_dir = self.tempdir.name
        self.root_dir = root_dir
        repo_dir = os.path.join(self.root_dir, instance["repo"].replace("/", "__"))
        if not os.path.exists(repo_dir):
            repo_url = (
                f"https://{token}@github.com/swe-bench/"
                + instance["repo"].replace("/", "__")
                + ".git"
            )
            if verbose:
                print(f"Cloning {instance['repo']} to {root_dir}")
            Repo.clone_from(repo_url, repo_dir)
        super().__init__(repo_dir, instance["base_commit"], verbose=verbose)
        self.instance = instance

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.tempdir is not None:
            self.tempdir.cleanup()
        return super().__exit__(exc_type, exc_val, exc_tb)


def get_imported_modules(filename):
    with open(filename, "rt") as file:
        tree = ast.parse(file.read(), filename)
    return [
        node
        for node in ast.iter_child_nodes(tree)
        if isinstance(node, (ast.Import, ast.ImportFrom))
    ]


def resolve_module_to_file(module, level, root_dir):
    components = module.split(".")
    if level > 0:
        components = components[:-level]
    for dirpath, dirnames, filenames in os.walk(root_dir):
        if dirpath.endswith(os.sep.join(components)):
            return [
                os.path.join(dirpath, filename)
                for filename in filenames
                if filename.endswith(".py")
            ]
    return []


def ingest_file_directory_contents(target_file, root_dir):
    imported_files = []
    files_to_check = [target_file]
    while files_to_check:
        current_file = files_to_check.pop()
        imported_files.append(current_file)
        imports = get_imported_modules(current_file)
        for node in imports:
            if isinstance(node, ast.Import):
                for alias in node.names:
                    files = resolve_module_to_file(alias.name, 0, root_dir)
                    for file in files:
                        if file not in imported_files and file not in files_to_check:
                            files_to_check.append(file)
            elif isinstance(node, ast.ImportFrom):
                files = resolve_module_to_file(node.module, node.level, root_dir)
                for file in files:
                    if file not in imported_files and file not in files_to_check:
                        files_to_check.append(file)
    return imported_files


def detect_encoding(filename):
    """
    Detect the encoding of a file
    """
    with open(filename, "rb") as file:
        rawdata = file.read()
    return chardet.detect(rawdata)["encoding"]


def list_files(root_dir, include_tests=False):
    files = []
    for filename in Path(root_dir).rglob("*.py"):
        if not include_tests and is_test(filename.as_posix()):
            continue
        files.append(filename.relative_to(root_dir).as_posix())
    return files


def ingest_directory_contents(root_dir, include_tests=False):
    files_content = {}
    for relative_path in list_files(root_dir, include_tests=include_tests):
        filename = os.path.join(root_dir, relative_path)
        encoding = detect_encoding(filename)
        if encoding is None:
            content = "[BINARY DATA FILE]"
        else:
            try:
                with open(filename, "r", encoding=encoding) as file:
                    content = file.read()
            except (UnicodeDecodeError, LookupError):
                content = "[BINARY DATA FILE]"
        files_content[relative_path] = content
    return files_content


def string_to_bool(v):
    if isinstance(v, bool):
        return v
    if v.lower() in ("yes", "true", "t", "y", "1"):
        return True
    elif v.lower() in ("no", "false", "f", "n", "0"):
        return False
    else:
        raise ArgumentTypeError(
            f"Truthy value expected: got {v} but expected one of yes/no, true/false, t/f, y/n, 1/0 (case insensitive)."
        )

inference/make_datasets/utils.py
princeton-nlp/SWE-bench
"""
This module contains functions for running a live inference session on a GitHub issue.
It clones the repository associated with the issue, builds a BM25 retrieval index, and
generates a prompt for the user to interact with the model. The output is saved to a
specified directory.
"""
import json
import shutil
import subprocess
from pathlib import Path
from ghapi.all import GhApi
import os
import re
import time
from datetime import datetime
from tqdm.auto import tqdm
from tempfile import TemporaryDirectory
from make_datasets.utils import ContextManager, string_to_bool, extract_diff, extract_minimal_patch
from make_datasets.bm25_retrieval import (
    make_index,
    clone_repo,
    search,
    DOCUMENT_ENCODING_FUNCTIONS,
)
from make_datasets.create_instance import (
    PROMPT_FUNCTIONS,
    TOKENIZER_FUNCS,
    make_code_text,
    ingest_files,
)
from run_model import call_chat, call_anthropic
import logging
from argparse import ArgumentParser

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


def get_problem_statement(owner, repo, issue_num, ghapi, include_comments=False):
    issue = ghapi.issues.get(owner, repo, issue_num)
    issue_text = "\n".join([issue.title, issue.body])
    # Solved issues may include comments that give answers away too much
    if include_comments:
        all_comments = list(ghapi.issues.list_comments(owner, repo, issue_num))
        comments = [comment.body for comment in all_comments]
        comment_text = "Comment: " if comments else "" + "\nComment:".join(comments)
        issue_text += "\n" + comment_text
    return issue_text


def get_readme_files(repo_path):
    files = list(Path(repo_path).iterdir())
    files = list(filter(lambda x: x.is_file(), files))
    files = list(filter(lambda x: x.name.lower().startswith("readme"), files))
    if files:
        files = sorted(files, key=lambda x: len(x.name))
        files = [files[0]]
    return [Path(file).relative_to(repo_path).as_posix() for file in files]


def make_instance(
    owner,
    repo,
    query,
    commit,
    root_dir,
    token,
    document_encoding_func,
    python,
    instance_id,
    tokenizer,
    tokenizer_func,
    prompt_style,
    max_context_len,
    include_readmes,
):
    """
    Creates an instance for a given query and repository.

    Args:
        owner (str): The owner of the repository.
        repo (str): The name of the repository.
        query (str): The query to search for.
        commit (str): The commit hash to use.
        root_dir (str): The root directory to clone the repository to.
        token (str): The GitHub token to use for authentication.
        document_encoding_func (function): The function to use for encoding documents.
        python (str): The path to the Python executable.
        instance_id (int): The ID of the instance.
        tokenizer (str): The name of the tokenizer to use.
        tokenizer_func (function): The function to use for tokenization.
        prompt_style (str): The style of prompt to use.
        max_context_len (int): The maximum length of the context.
        include_readmes (bool): Whether to include README files in the instance.

    Returns:
        dict: The instance.
    """
    thread_id = 0
    instance = {"instance_id": instance_id, "problem_statement": query}
    logger.info(f"Cloning repo {owner}/{repo}")
    repo_dir = clone_repo(f"{owner}/{repo}", root_dir, token, False, thread_id)
    if commit is None:
        commit = subprocess.check_output(
            ["git", "rev-parse", "HEAD"], cwd=repo_dir
        ).decode("utf-8").strip()
    logger.info(f"Buidling BM25 retrieval index for {owner}/{repo}@{commit}")
    index_dir = make_index(
        repo_dir,
        root_dir,
        commit,
        document_encoding_func,
        python,
        thread_id,
        instance_id,
    )
    results = search(instance, index_dir)
    hits = results["hits"]
    logger.info(f"Retrieved {len(hits)} documents")
    with ContextManager(repo_dir, commit) as cm:
        if include_readmes:
            readmes = get_readme_files(cm.repo_path)
        else:
            readmes = list()
        instance["readmes"] = ingest_files(readmes)
        for hit in hits:
            hit["file_contents"] = open(hit["docid"]).read()
        instance["file_contents"] = dict()
        base_text_inputs = PROMPT_FUNCTIONS[prompt_style](instance)
        base_text_input_length = len(tokenizer_func(base_text_inputs, tokenizer))
        instance["file_contents"] = {x["docid"]: x["file_contents"] for x in hits}
        cur_input_len = base_text_input_length
        include_files = list()
        for filename in [x["docid"] for x in hits]:
            content = make_code_text({filename: instance["file_contents"][filename]})
            tokens = tokenizer_func(content, tokenizer)
            if cur_input_len + len(tokens) < max_context_len:
                include_files.append(filename)
                cur_input_len += len(tokens)
        logger.info(
            f"Including {len(include_files)} files in context with {cur_input_len} tokens:\n"
            + "\n\t".join(sorted(include_files))
        )
        instance["file_contents"] = {
            filename: instance["file_contents"][filename] for filename in include_files
        }
        instance["text_inputs"] = PROMPT_FUNCTIONS[prompt_style](instance)
        return instance


def parse_issue_url(issue_url):
    issue_pat = re.compile(r"github\.com\/(.+?)\/(.+?)\/issues\/(\d+)")
    match = issue_pat.search(issue_url)
    if not match:
        raise ValueError(
            (
                f"issue_url ({issue_url}) does not seem to be a valid issue url."
                + "\nPlease use url like https://github.com/owner/repo/issues/12345"
            )
        )
    owner, repo, issue_num = match.groups()
    return owner, repo, issue_num


def main(
    model_name,
    prompt_style,
    issue_url,
    base_commit,
    max_context_length,
    document_encoding_func,
    output_dir,
    root_dir,
    include_readmes,
):
    if base_commit is not None and len(issue_url) != len(base_commit):
        raise ValueError(
            f"Must provide either no base commits or one base commit per issue url"
        )
    if base_commit is None:
        base_commit = [None] * len(issue_url)
    gh_token = os.environ.get("GITHUB_TOKEN", None)
    if gh_token is not None:
        logger.warning(f'Using GitHub token: {"*" * 8}{gh_token[-4:]}')
    gh = GhApi(token=gh_token)
    tokenizer, tokenizer_func = TOKENIZER_FUNCS["cl100k"]
    document_encoding_func = DOCUMENT_ENCODING_FUNCTIONS[document_encoding_func]
    python = subprocess.check_output(["which", "python"]).decode("utf-8").strip()
    outputs = list()
    for issue, commit in tqdm(zip(issue_url, base_commit), total=len(issue_url)):
        owner, repo, issue_num = parse_issue_url(issue)
        problem_statement = get_problem_statement(owner, repo, int(issue_num), gh)
        instance_id = f"{owner}__{repo}-{issue_num}"
        logger.info(f"Creating instance {instance_id}")
        instance = make_instance(
            owner,
            repo,
            problem_statement,
            commit,
            root_dir,
            gh_token,
            document_encoding_func,
            python,
            instance_id,
            tokenizer,
            tokenizer_func,
            prompt_style,
            max_context_length,
            include_readmes,
        )
        logger.info(f"Calling model {model_name}")
        start = time.time()
        if model_name.startswith("gpt"):
            import openai
            openai.api_key = os.environ.get("OPENAI_API_KEY", None)
            inputs = instance["text_inputs"]
            response, _ = call_chat(
                model_name, inputs, use_azure=False, temperature=0, top_p=1
            )
            completion = response.choices[0]["message"]["content"]
            logger.info(f'Generated {response.usage.completion_tokens} tokens in {(time.time() - start):.2f} seconds')
        else:
            from anthropic import Anthropic
            api_key = os.environ.get("ANTHROPIC_API_KEY", None)
            anthropic = Anthropic(api_key=api_key)
            response = call_anthropic(
                inputs, anthropic, model_name, temperature=0, top_p=1
            )
            completion = response.completion
        model_patch = extract_diff(completion)
        minimal_patch = extract_minimal_patch(model_patch)
        outputs.append(
            {
                "instance_id": instance_id,
                "response": completion,
                "problem_statement": problem_statement,
                "text_inputs": inputs,
                "model_patch": model_patch,
                "minimal_patch": minimal_patch,
            }
        )
    os.makedirs(output_dir, exist_ok=True)
    output_file = Path(
        output_dir,
        f'{model_name}__{prompt_style}__{datetime.now().strftime("%Y-%m-%d_%H-%M-%S")}.jsonl',
    )
    with open(output_file, "+a") as f:
        for output in outputs:
            print(json.dumps(output), file=f, flush=True)
    logger.info(f"Wrote output to {output_file}")


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("--model_name", type=str)
    parser.add_argument(
        "--prompt_style", type=str, choices=PROMPT_FUNCTIONS.keys(), default="style-3"
    )
    parser.add_argument("--issue_url", type=str, nargs="+")
    parser.add_argument("--base_commit", type=str, nargs="+")
    parser.add_argument("--max_context_length", type=int, default=16_000)
    parser.add_argument(
        "--document_encoding_func",
        type=str,
        choices=DOCUMENT_ENCODING_FUNCTIONS.keys(),
        default="file_name_and_contents",
    )
    parser.add_argument("--output_dir", type=str, default="./live_outputs")
    parser.add_argument("--root_dir", type=str, default="./run_live_data")
    parser.add_argument("--include_readmes", type=string_to_bool, default=False)
    args = parser.parse_args()
    main(**vars(args))

inference/run_live.py
princeton-nlp/SWE-bench
import json
import os
import time
import dotenv
import traceback
from pathlib import Path
from tqdm.auto import tqdm
import numpy as np
import openai
from anthropic import HUMAN_PROMPT, AI_PROMPT, Anthropic
from tenacity import (
    retry,
    stop_after_attempt,
    wait_random_exponential,
)
from datasets import load_dataset, concatenate_datasets
from make_datasets.utils import extract_diff
from argparse import ArgumentParser
import logging

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)
dotenv.load_dotenv()


MODEL_LIMITS = {
    'gpt-3.5-turbo-16k-0613': 16_000,
    'gpt-4-32k-0613': 31_000,
    'gpt-4-0613': 7_800,
}


MODEL_COST_PER_INPUT = {
    'gpt-3.5-turbo-16k-0613': 0.0000015,
    'gpt-4-0613': 0.00003,
    'gpt-4-32k-0613': 0.00006,
    'gpt-4': 0.00003,
    'gpt-35-turbo': 0.0000015,
    'gpt-4-32k': 0.00006,
}

 
MODEL_COST_PER_OUTPUT = {
    'gpt-3.5-turbo-16k-0613': 0.000002,
    'gpt-4-0613': 0.0006,
    'gpt-4-32k-0613': 0.00012,
    'gpt-4': 0.0006,
    'gpt-35-turbo': 0.000002,
    'gpt-4-32k': 0.00012,
}


ENGINES = {
    'gpt-3.5-turbo-16k-0613': 'gpt-35-turbo-16k',
    'gpt-4-0613': 'gpt-4',
    'gpt-4-32k-0613': 'gpt-4-32k',
}


def calc_cost(response):
    model_name = response.model
    input_tokens = response.usage.prompt_tokens
    output_tokens = response.usage.completion_tokens
    cost = (
        MODEL_COST_PER_INPUT[model_name] * input_tokens
        + MODEL_COST_PER_OUTPUT[model_name] * output_tokens
    )
    logger.info(f'input_tokens={input_tokens}, output_tokens={output_tokens}, cost={cost:.2f}')
    return cost


@retry(wait=wait_random_exponential(min=30, max=600), stop=stop_after_attempt(3))
def call_chat(model_name_or_path, inputs, use_azure, temperature, top_p):
    system_messages = inputs.split("\n", 1)[0]
    user_message = inputs.split("\n", 1)[1]
    try:
        if use_azure:
            response = openai.ChatCompletion.create(
                engine=ENGINES[model_name_or_path] if use_azure else None,
                messages=[
                    {"role": "system", "content": system_messages},
                    {"role": "user", "content": user_message},
                ],
                temperature=temperature,
                top_p=top_p,
            )
        else:
            response = openai.ChatCompletion.create(
                model=model_name_or_path,
                messages=[
                    {"role": "system", "content": system_messages},
                    {"role": "user", "content": user_message},
                ],
                temperature=temperature,
                top_p=top_p,
            )
        cost = calc_cost(response)
        return response, cost
    except openai.error.InvalidRequestError as e:
        if e.code == 'context_length_exceeded':
            print("Context length exceeded")
            return None
        raise e


def openai_inference(
    test_dataset,
    model_name_or_path,
    output_file,
    model_args,
    existing_ids,
    max_cost,
):
    openai_key = os.environ.get("OPENAI_API_KEY", None)
    openai.api_key = openai_key
    print(f"Using OpenAI key {'*' * max(0, len(openai_key)-5) + openai_key[-5:]}")

    use_azure = model_args.pop("use_azure", False)
    if use_azure:
        openai.api_type = "azure"
        openai.api_base = "https://pnlpopenai3.openai.azure.com/"
        openai.api_version = '2023-05-15'

    temperature = model_args.pop('temperature', 0.2)
    top_p = model_args.pop('top_p', 0.95 if temperature > 0 else 1)
    print(f'Using temperature={temperature}, top_p={top_p}')

    basic_args = {
        "model_name_or_path": model_name_or_path,
    }
    total_cost = 0
    test_dataset = test_dataset.filter(lambda x: len(x['input_ids']) <= MODEL_LIMITS[model_name_or_path])
    print(f"Filtered to {len(test_dataset)} instances")
    with open(output_file, "a+") as f:
        for datum in tqdm(test_dataset, desc=f"Inference for {model_name_or_path}"):
            instance_id = datum["instance_id"]
            if instance_id in existing_ids:
                continue
            output_dict = {"instance_id": instance_id}
            output_dict.update(basic_args)
            output_dict["text"] = f"{datum['text']}\n\n"
            if len(datum['input_ids']) > MODEL_LIMITS[model_name_or_path]:
                output_dict["full_output"] = None
                output_dict["model_patch"] = None
            else:
                if model_name_or_path == 'gpt-4-32k-0613' and len(datum['input_ids']) <= 6000:
                    response, cost = call_chat('gpt-4-0613', output_dict["text"], use_azure, temperature, top_p)
                    completion = response.choices[0]['message']['content']
                else:
                    response, cost = call_chat(
                        output_dict["model_name_or_path"], output_dict["text"], use_azure, temperature, top_p
                    )
                    completion = response.choices[0]['message']['content']
                total_cost += cost
                print(f"Total Cost: {total_cost:.2f}")
                output_dict["full_output"] = completion
                output_dict["model_patch"] = extract_diff(completion)
            print(json.dumps(output_dict), file=f, flush=True)
            if max_cost is not None and total_cost >= max_cost:
                print(f"Reached max cost {max_cost}, exiting")
                break


@retry(wait=wait_random_exponential(min=60, max=600), stop=stop_after_attempt(6))
def call_anthropic(inputs, anthropic, model_name_or_path, temperature, top_p):
    try:
        completion = anthropic.completions.create(
            model=model_name_or_path,
            max_tokens_to_sample=6000,
            prompt=inputs,
            temperature=temperature,
            top_p=top_p,
        )
        return completion
    except Exception as e:
        logger.error(e)
        logger.error(f"Inputs: {inputs}")
        traceback.print_exc()
        time.sleep(20)
        return None


def anthropic_inference(
    test_dataset,
    model_name_or_path,
    output_file,
    model_args,
    existing_ids,
    max_cost,
):
    api_key = model_args.pop(
        "anthropic_api_key", os.environ.get("ANTHROPIC_API_KEY", None)
    )
    anthropic = Anthropic(api_key=api_key)
    temperature = model_args.pop('temperature', 0.2)
    top_p = model_args.pop('top_p', 0.95 if temperature > 0 else 1)
    print(f'Using temperature={temperature}, top_p={top_p}')

    basic_args = {
        "model_name_or_path": model_name_or_path,
    }
    with open(output_file, "a+") as f:
        for datum in tqdm(test_dataset, desc=f"Inference for {model_name_or_path}"):
            instance_id = datum["instance_id"]
            if instance_id in existing_ids:
                continue
            output_dict = {"instance_id": instance_id}
            output_dict.update(basic_args)
            output_dict["text_inputs"] = f"{HUMAN_PROMPT} {datum['text']}\n\n{AI_PROMPT}"
            completion = call_anthropic(output_dict["text_inputs"], anthropic, model_name_or_path, temperature, top_p)
            output_dict["full_output"] = completion.completion
            output_dict["model_patch"] = extract_diff(completion.completion)
            print(json.dumps(output_dict), file=f, flush=True)


def parse_model_args(model_args):
    kwargs = dict()
    if model_args is not None:
        for arg in model_args.split(","):
            key, value = arg.split("=")
            # infer value type
            if value in {"True", "False"}:
                kwargs[key] = value == "True"
            elif value.isnumeric():
                kwargs[key] = int(value)
            elif value.replace(".", "", 1).isnumeric():
                kwargs[key] = float(value)
            elif value in {"None"}:
                kwargs[key] = None
            elif value in {"[]"}:
                kwargs[key] = []
            elif value in {"{}"}:
                kwargs[key] = {}
            elif value.startswith("'") and value.endswith("'"):
                kwargs[key] = value[1:-1]
            elif value.startswith('"') and value.endswith('"'):
                kwargs[key] = value[1:-1]
            else:
                kwargs[key] = value
    return kwargs


def main(
    dataset_name,
    model_name_or_path,
    shard_id,
    num_shards,
    output_dir,
    model_args,
    max_cost,
):
    if shard_id is None and num_shards is not None:
        logger.warning(f"Received num_shards={num_shards} but shard_id is None, ignoring")
    if shard_id is not None and num_shards is None:
        logger.warning(f"Received shard_id={shard_id} but num_shards is None, ignoring")
    model_args = parse_model_args(model_args)
    model_nickname = model_name_or_path
    if "checkpoint" in Path(model_name_or_path).name:
        model_nickname = Path(model_name_or_path).parent.name
    else:
        model_nickname = Path(model_name_or_path).name
    output_file = f"{model_nickname}__{dataset_name.split('/')[-1]}"
    if shard_id is not None and num_shards is not None:
        output_file += f"__shard-{shard_id}__num_shards-{num_shards}"
    output_file = Path(output_dir, output_file + ".jsonl")
    logger.info(f"Will write to {output_file}")
    existing_ids = set()
    if os.path.exists(output_file):
        with open(output_file, 'r') as f:
            for line in f:
                data = json.loads(line)
                instance_id = data["instance_id"]
                existing_ids.add(instance_id)
    logger.info(f"Read {len(existing_ids)} already completed ids")
    dataset = load_dataset(dataset_name)
    load_splits = [split for split in dataset.keys() if 'test' in split]
    dataset = concatenate_datasets([dataset[split] for split in load_splits])
    if 'input_ids' in dataset.features:
        lens = np.array(list(map(len, dataset['input_ids'])))
    else:
        lens = np.array(list(map(len, dataset['text'])))
    dataset = dataset.select(np.argsort(lens))
    if len(existing_ids) > 0:
        dataset = dataset.filter(lambda x: x['instance_id'] not in existing_ids, desc="Filtering existing ids")
    if shard_id is not None and num_shards is not None:
        dataset = dataset.shard(num_shards, shard_id, contiguous=True)
    inference_args = {
        "test_dataset": dataset,
        "model_name_or_path": model_name_or_path,
        "output_file": output_file,
        "model_args": model_args,
        "existing_ids": existing_ids,
        "max_cost": max_cost,
    }
    if model_name_or_path in {"claude-2"}:
        anthropic_inference(**inference_args)
    elif model_name_or_path.startswith("gpt"):
        openai_inference(**inference_args)
    else:
        raise ValueError(f"Invalid model name or path {model_name_or_path}")
    logger.info(f"Done!")


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument(
        "--dataset_name",
        type=str,
        required=True,
        help="HuggingFace dataset name, with pre-tokenized inputs",
    )
    parser.add_argument(
        "--model_name_or_path",
        type=str,
        help="Path to the directory containing a lora or model",
    )
    parser.add_argument(
        "--shard_id",
        type=int,
        default=None,
        help="Shard id to process. If None, process all shards.",
    )
    parser.add_argument(
        "--num_shards",
        type=int,
        default=None,
        help="Number of shards. If None, process all shards.",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default=None,
        required=True,
        help="Path to the output file.",
    )
    parser.add_argument(
        "--model_args",
        type=str,
        default=None,
        help="List of model arguments separated by commas. (e.g. 'top_p=0.95,temperature=0.70')",
    )
    parser.add_argument(
        "--max_cost",
        type=float,
        default=None,
        help="Maximum cost to spend on inference.",
    )
    args = parser.parse_args()
    main(**vars(args))

inference/run_model.py
princeton-nlp/SWE-bench
import json, os

from log_parsers import MAP_REPO_TO_PARSER, TestStatus
from typing import Dict, List
from getters import (
    get_file_name_from_lp,
    get_repo_from_lp,
    log_path_to_sms,
    FAIL_TO_PASS,
    FAIL_TO_FAIL,
    PASS_TO_PASS,
    PASS_TO_FAIL,
    test_failed,
    test_passed,
)


def convert_log_to_ground_truth(
    log_fp: List, save_dir: str = None, verbose: bool = False
) -> Dict:
    """
    Convert log file generated from check instances into ground truth dict

    Args:
        log_dir (str): path to log files
        save_dir (str): path to save results
        verbose (bool): whether or not to show logs
    Returns:
        dict: test case to test status mapping
    """
    inst_file_name = get_file_name_from_lp(log_fp)
    repo = get_repo_from_lp(log_fp)
    log_parser = MAP_REPO_TO_PARSER[repo]

    sms, found = log_path_to_sms(log_fp, log_parser)
    if not found:
        raise ValueError(
            "Log file could not be parsed properly (Before, After Logs not found)"
        )
    sm_before, sm_after = sms[0], sms[1]

    status_ground_truth = {
        FAIL_TO_PASS: [],
        FAIL_TO_FAIL: [],
        PASS_TO_PASS: [],
        PASS_TO_FAIL: [],
    }

    for test, status in sm_after.items():
        if status == TestStatus.PASSED.value:
            if test_passed(test, sm_before):
                status_ground_truth[PASS_TO_PASS].append(test)
            elif test_failed(test, sm_before):
                status_ground_truth[FAIL_TO_PASS].append(test)
        if status == TestStatus.FAILED.value:
            if test_passed(test, sm_before):
                status_ground_truth[PASS_TO_FAIL].append(test)
            elif test_failed(test, sm_before):
                status_ground_truth[FAIL_TO_FAIL].append(test)

    if save_dir is not None:
        results_file = f"{inst_file_name.split('.')[0]}.json"
        if verbose:
            print(f"Saving results to {os.path.join(save_dir, results_file)}")
        with open(os.path.join(save_dir, results_file), "w") as f:
            json.dump(status_ground_truth, f, indent=4)

    return status_ground_truth

metrics/conversion.py
princeton-nlp/SWE-bench
import re

from log_parsers import MAP_REPO_TO_PARSER, TestStatus
from typing import Dict, List


# Evaluation Log Constants
APPLY_PATCH_FAIL = ">>>>> Patch Apply Failed"
APPLY_PATCH_PASS = ">>>>> Applied Patch"
INSTALL_FAIL = ">>>>> Init Failed"
INSTALL_PASS = ">>>>> Init Succeeded"
RESET_FAILED = ">>>>> Reset Failed"
TESTS_TIMEOUT = ">>>>> Tests Timed Out"
TESTS_ERROR = ">>>>> Tests Errored"

# Result Categories
FAIL_TO_PASS = "FAIL_TO_PASS"
FAIL_TO_FAIL = "FAIL_TO_FAIL"
PASS_TO_PASS = "PASS_TO_PASS"
PASS_TO_FAIL = "PASS_TO_FAIL"


def get_diffs(sm_1: Dict, sm_2: Dict) -> Dict:
    """
    Get differences between two test status maps

    Args:
        sm_1 (dict): test case to test status mapping
        sm_2 (dict): test case to test status mapping
    Returns:
        dict: test case to test status mapping
    """
    set1 = set(sm_1.items())
    set2 = set(sm_2.items())
    diffs = set1 ^ set2

    diff_map = {}
    for diff in diffs:
        if diff[0] not in diff_map:
            diff_map[diff[0]] = []
        diff_map[diff[0]].append(diff[1])
    return diff_map


def get_logs_eval(log_fp: str) -> (Dict, bool):
    """
    Retrieve evaluation results for a task instance from its corresponding log file

    Args:
        log_fp (str): path to log file
    Returns:
        bool: whether the patch applied successfully
        dict: status map
    """
    repo = get_repo_from_lp(log_fp)
    log_parser = MAP_REPO_TO_PARSER[repo]

    with open(log_fp, "r") as f:
        content = f.read()
        if any([x in content for x in [APPLY_PATCH_FAIL, RESET_FAILED, TESTS_ERROR, TESTS_TIMEOUT, "Failed to reset task environment"]]) or APPLY_PATCH_PASS not in content:
            # Eval patch was not applied successfully
            return {}, False

        # Get status map of evaluation results
        content = content.split(f"{APPLY_PATCH_PASS} (pred)")[-1]
        return log_parser(content), True


def get_logs_gold(log_fp: str) -> (str, str):
    """
    Retrieve pre-patch, post-patch test logs from a log file

    Args:
        log_fp (str): path to log file
    Returns:
        str: pre-patch, post-patch test logs
    """
    with open(log_fp, "r") as f:
        content = f.read()
        if len(re.findall(APPLY_PATCH_PASS, content)) != 2:
            return None, None
        logs = content.split(APPLY_PATCH_PASS)
        log_before, log_after = logs[1], logs[2]
        return log_before, log_after


get_file_name_from_lp = lambda x: x.rsplit("/", 1)[-1]


get_id_from_lp = lambda x: get_file_name_from_lp(x).split(".")[0]


get_repo_from_lp = lambda x: get_id_from_lp(x).rsplit("-", 1)[0].replace("__", "/")


def log_path_to_sms(log_fp: str, log_parser) -> (List, bool):
    """
    Wrapper for getting log data from log_parser file

    Args:
        log_fp (str): path to log file
        log_parser (function): function to parse log file
    Returns:
        list: list of status maps
        bool: whether or not log file was parsed properly
    """
    log_before, log_after = get_logs_gold(log_fp)
    if log_before is None and log_after is None:
        # Skip if either one of test patch apply + patch apply failed
        return None, False

    try:
        sm_before = log_parser(log_before)
        sm_after = log_parser(log_after)
    except Exception as e:
        # Skip if log file was not parsed properly
        print(f"Error parsing log {log_fp}: {e}")
        sm_before, sm_after = None, None

    if sm_before is None or sm_after is None:
        # Skip if test patch or patch statuses are none
        return None, False

    return [sm_before, sm_after], True


test_passed = lambda case, sm: case in sm and sm[case] == TestStatus.PASSED.value

test_failed = lambda case, sm: case not in sm or any(
    [sm[case] == status for status in [TestStatus.FAILED.value, TestStatus.ERROR.value]]
)

metrics/getters.py
princeton-nlp/SWE-bench
import re

from enum import Enum
from typing import Dict


class TestStatus(Enum):
    FAILED = "FAILED"
    PASSED = "PASSED"
    SKIPPED = "SKIPPED"
    ERROR = "ERROR"


def parse_log_pytest(log: str) -> Dict:
    """
    Parser for test logs generated with PyTest framework

    Args:
        log (str): log content
    Returns:
        dict: test case to test status mapping
    """
    test_status_map = {}
    for line in log.split("\n"):
        if any([line.startswith(x.value) for x in TestStatus]):
            # Additional parsing for FAILED status
            if line.startswith(TestStatus.FAILED.value):
                line = line.replace(" - ", " ")
            test_case = line.split()
            if len(test_case) <= 1:
                continue
            test_status_map[test_case[1]] = test_case[0]
    return test_status_map


def parse_log_django(log: str) -> Dict:
    """
    Parser for test logs generated with Django tester framework

    Args:
        log (str): log content
    Returns:
        dict: test case to test status mapping
    """
    test_status_map = {}
    lines = log.split("\n")
    for line in lines:
        line = line.strip()
        if line.endswith(" ... ok"):
            test = line.split(" ... ok")[0]
            test_status_map[test] = TestStatus.PASSED.value
        if " ... skipped" in line:
            test = line.split(" ... skipped")[0]
            test_status_map[test] = TestStatus.SKIPPED.value
        if line.endswith(" ... FAIL"):
            test = line.split(" ... FAIL")[0]
            test_status_map[test] = TestStatus.FAILED.value
        if line.startswith("FAIL:"):
            test = line.split()[1].strip()
            test_status_map[test] = TestStatus.FAILED.value
        if line.endswith(" ... ERROR"):
            test = line.split(" ... ERROR")[0]
            test_status_map[test] = TestStatus.ERROR.value
        if line.startswith("ERROR:"):
            test = line.split()[1].strip()
            test_status_map[test] = TestStatus.ERROR.value
    return test_status_map


def parse_log_pytest_v2(log):
    """
    Parser for test logs generated with PyTest framework (Later Version)

    Args:
        log (str): log content
    Returns:
        dict: test case to test status mapping
    """
    test_status_map = {}
    escapes = "".join([chr(char) for char in range(1, 32)])
    for line in log.split("\n"):
        line = re.sub(r"\[(\d+)m", "", line)
        translator = str.maketrans("", "", escapes)
        line = line.translate(translator)
        if any([line.startswith(x.value) for x in TestStatus]):
            if line.startswith(TestStatus.FAILED.value):
                line = line.replace(" - ", " ")
            test_case = line.split()
            test_status_map[test_case[1]] = test_case[0]
    return test_status_map


def parse_log_seaborn(log):
    """
    Parser for test logs generated with seaborn testing framework

    Args:
        log (str): log content
    Returns:
        dict: test case to test status mapping
    """
    test_status_map = {}
    for line in log.split("\n"):
        if line.startswith(TestStatus.FAILED.value):
            test_case = line.split()[1]
            test_status_map[test_case] = TestStatus.FAILED.value
        elif f" {TestStatus.PASSED.value} " in line:
            parts = line.split()
            if parts[1] == TestStatus.PASSED.value:
                test_case = parts[0]
                test_status_map[test_case] = TestStatus.PASSED.value
    return test_status_map


def parse_log_sympy(log):
    """
    Parser for test logs generated with Sympy framework

    Args:
        log (str): log content
    Returns:
        dict: test case to test status mapping
    """
    test_status_map = {}
    pattern = r"(_*) (.*)\.py:(.*) (_*)"
    matches = re.findall(pattern, log)
    for match in matches:
        test_case = f"{match[1]}.py:{match[2]}"
        test_status_map[test_case] = TestStatus.FAILED.value
    for line in log.split("\n"):
        line = line.strip()
        if line.startswith("test_"):
            if line.endswith(" E"):
                test = line.split()[0]
                test_status_map[test] = TestStatus.ERROR.value
            if line.endswith(" F"):
                test = line.split()[0]
                test_status_map[test] = TestStatus.FAILED.value
            if line.endswith(" ok"):
                test = line.split()[0]
                test_status_map[test] = TestStatus.PASSED.value
    return test_status_map


parse_log_flask = parse_log_pytest
parse_log_matplotlib = parse_log_pytest
parse_log_pylint = parse_log_pytest
parse_log_requests = parse_log_pytest
parse_log_xarray = parse_log_pytest

parse_log_astropy = parse_log_pytest_v2
parse_log_scikit = parse_log_pytest_v2
parse_log_sphinx = parse_log_pytest_v2


MAP_REPO_TO_PARSER = {
    "astropy/astropy": parse_log_astropy,
    "django/django": parse_log_django,
    "matplotlib/matplotlib": parse_log_matplotlib,
    "mwaskom/seaborn": parse_log_seaborn,
    "pallets/flask": parse_log_flask,
    "psf/requests": parse_log_requests,
    "pydata/xarray": parse_log_xarray,
    "pylint-dev/pylint": parse_log_pylint,
    "pytest-dev/pytest": parse_log_pytest,
    "scikit-learn/scikit-learn": parse_log_scikit,
    "sphinx-doc/sphinx": parse_log_sphinx,
    "sympy/sympy": parse_log_sympy,
}

metrics/log_parsers.py
princeton-nlp/SWE-bench
from enum import Enum
from getters import FAIL_TO_FAIL, FAIL_TO_PASS, PASS_TO_FAIL, PASS_TO_PASS
from statistics import mean
from typing import List, Dict


class ResolvedStatus(Enum):
    NO = "RESOLVED_NO"
    PARTIAL = "RESOLVED_PARTIAL"
    FULL = "RESOLVED_FULL"


def compute_fail_to_pass(report: Dict) -> float:
    """
    Compute fail-to-pass metric. Accepts single report as argument.
    """
    total = len(report[FAIL_TO_PASS]["success"]) + len(report[FAIL_TO_PASS]["failure"])
    if total == 0:
        return 1
    return len(report[FAIL_TO_PASS]["success"]) / total


def compute_pass_to_pass(report: Dict) -> float:
    """
    Compute pass-to-pass metric. Accepts single report as argument.
    """
    total = len(report[PASS_TO_PASS]["success"]) + len(report[PASS_TO_PASS]["failure"])
    if total == 0:
        # TODO: Don't factor in p2p metrics
        return 1
    return len(report[PASS_TO_PASS]["success"]) / total


def compute_fail_to_pass_unweighted(reports: List[Dict]) -> float:
    """
    Compute unweighted fail-to-pass metric. Accepts list of reports as argument.
    """
    if len(reports) == 0:
        return 0
    return mean([compute_fail_to_pass(r) for r in reports])


def compute_pass_to_pass_unweighted(reports: List[Dict]) -> float:
    """
    Compute unweighted pass-to-pass metric. Accepts list of reports as argument.
    """
    if len(reports) == 0:
        return 0
    return mean([compute_pass_to_pass(r) for r in reports])


def compute_fail_to_pass_weighted(reports: List[Dict]) -> float:
    """
    Compute weighted fail-to-pass metric. Accepts list of reports as argument.
    """
    report_all = {
        FAIL_TO_PASS: {
            "success": [x for r in reports for x in r[FAIL_TO_PASS]["success"]],
            "failure": [x for r in reports for x in r[FAIL_TO_PASS]["failure"]],
        },
    }
    return compute_fail_to_pass(report_all)


def compute_pass_to_pass_weighted(reports: List[Dict]) -> float:
    """
    Compute weighted pass-to-pass metric. Accepts list of reports as argument.
    """
    report_all = {
        PASS_TO_PASS: {
            "success": [x for r in reports for x in r[PASS_TO_PASS]["success"]],
            "failure": [x for r in reports for x in r[PASS_TO_PASS]["failure"]],
        },
    }
    return compute_pass_to_pass(report_all)


def get_resolution_status(report: Dict) -> str:
    """
    Determine resolved status of an evaluation instance

    Criteria:
        - If fail-to-pass (Resolution) = 1 and pass-to-pass (Maintenance) = 1 -> FULL
        - If (fail-to-pass (Resolution) < 1 and > 0) and pass-to-pass (Maintenance) = 1 -> PARTIAL
        - Otherwise -> NO
    """
    f2p = compute_fail_to_pass(report)
    p2p = compute_pass_to_pass(report)

    if f2p == 1 and p2p == 1:
        return ResolvedStatus.FULL.value
    elif f2p < 1 and f2p > 0 and p2p == 1:
        return ResolvedStatus.PARTIAL.value
    else:
        return ResolvedStatus.NO.value
metrics/metrics.py
princeton-nlp/SWE-bench
import glob
import os

from typing import List

from log_parsers import MAP_REPO_TO_PARSER
from getters import (
    log_path_to_sms, get_diffs, get_repo_from_lp,
    APPLY_PATCH_FAIL, APPLY_PATCH_PASS, TESTS_TIMEOUT
)


def monitor_validation(
    path_to_logs: str, log_prefix: str = None
) -> (List, List, List, List):
    """
    Check log files generated from a `check_instances` run to see how many instances were successfully
    installed and/or tested.

    Args:
        path_to_logs (str): path to log files
        log_prefix (str): prefix of log files
    Returns:
        failed_install (list): list of log files where instances failed to install
        corrupt_test_patch (list): list of log files where test patch was corrupt
        corrupt_patch (list): list of log files where patch was corrupt
        success (list): list of log files where instances were successfully installed and tested
    """
    failed_install = list()
    corrupt_test_patch = list()
    corrupt_patch = list()
    timeout = list()
    success = list()
    count = 0

    # Iterate through all logs with prefix in path_to_logs
    file_format = f"{log_prefix}*.log" if log_prefix else "*.log"
    for x in glob.glob(os.path.join(path_to_logs, file_format)):
        with open(x, "r") as f:
            content = f.read()
            if any([x in content for x in [APPLY_PATCH_FAIL, APPLY_PATCH_PASS]]):
                # Count number of successful patch applies
                patch_applies = content.count(APPLY_PATCH_PASS)
                if TESTS_TIMEOUT in content:
                    # Running test timed out
                    timeout.append(x)
                elif patch_applies == 0:
                    # If none, then test patch was corrupt
                    corrupt_test_patch.append(x)
                elif patch_applies == 1:
                    # If one, then patch was corrupt
                    corrupt_patch.append(x)
                else:
                    # If both, then both test patch and patch were applied successfully
                    assert patch_applies == 2
                    success.append(x)
            else:
                # If no patch was applied, then installation failed
                failed_install.append(x)
            count += 1

    # Logging
    print(f"Total Attempts: {count}")
    print(f"Failed: {len(failed_install)}")
    print(f"Usable: {count - len(failed_install)}")
    print(f"Corrupt Test: {len(corrupt_test_patch)}")
    print(f"Corrupt Diff: {len(corrupt_patch)}")
    print(f"Test Script Timeout: {len(timeout)}")
    print(f"Success E2E: {len(success)}")

    # Check that all instances were accounted for
    assert (
        len(success)
        + len(corrupt_patch)
        + len(corrupt_test_patch)
        + len(failed_install)
        + len(timeout)
        == count
    )

    return failed_install, corrupt_test_patch, corrupt_patch, timeout, success


def monitor_logs_same_diff(log_dir: str, repo: str = None) -> (list, list):
    """
    Given a log directory and repo, return a list of logs where pre-test
    and post-test logs are same/different

    Args:
        log_dir (str): path to log files
        repo (str): repo name
    Returns:
        logs_same: list of logs where post test behavior is same
        logs_diff: list of logs where post test behavior is different
    """
    logs_same, logs_diff = [], []

    # Find all log files for a repo in the given log directory
    file_format = f"{repo.split('/')[0]}*.log" if repo else "*.log"
    for log_fp in glob.glob(os.path.join(log_dir, file_format)):
        if repo:
            log_parser = MAP_REPO_TO_PARSER[repo]
        else:
            # Get repo from log file path
            repo = get_repo_from_lp(log_fp)
            log_parser = MAP_REPO_TO_PARSER[repo]

        sms, found = log_path_to_sms(log_fp, log_parser)
        if not found:
            continue
        sm_before, sm_after = sms[0], sms[1]

        # Get differences between pre, post patch behavior for log file
        diffs = get_diffs(sm_before, sm_after)
        if len(diffs) == 0:
            logs_same.append(log_fp)
        else:
            logs_diff.append((log_fp, diffs))

    return logs_same, logs_diff

metrics/monitor.py
princeton-nlp/SWE-bench
import glob, json, os

from collections import Counter
from getters import (
    get_file_name_from_lp,
    get_logs_eval,
    get_id_from_lp,
    FAIL_TO_FAIL,
    FAIL_TO_PASS,
    PASS_TO_FAIL,
    PASS_TO_PASS,
    test_failed,
    test_passed,
)
from log_parsers import TestStatus
from metrics import (
    compute_fail_to_pass_unweighted,
    compute_fail_to_pass_weighted,
    compute_pass_to_pass_unweighted,
    compute_pass_to_pass_weighted,
    get_resolution_status,
)
from typing import Dict, List


### MARK - Eval Report Generation


def get_eval_report(
    eval_sm: Dict,
    gold_results: Dict,
    calculate_to_fail: bool = False
) -> Dict:
    """
    Create a report based on failure/pass change from gold results to eval results.

    Args:
        eval_sm (dict): evaluation status map
        gold_results (dict): gold results
        calculate_to_fail (bool): whether to calculate metrics for "x to fail" tests
    Returns:
        report (dict): report of metrics

    Metric Definitions (Gold Result Pair + Eval Result):
    - Fail-Pass (F2P) + P: Success (Resolution)
    - Pass-Pass (P2P) + P: Success (Maintenance)
    - Fail-Pass (F2P) + F: Failure
    - Pass-Pass (P2P) + F: Failure

    Miscellaneous Definitions
    - Fail-Fail (F2F) + F: Failure Maintenance
    - Pass-Fail (P2F) + F: Not considered
    - Fail-Fail (F2F) + P: Success (Extra Credit)
    - Pass-Fail (P2F) + P: Not considered
    """
    # Calculate resolution metrics
    f2p_success = []
    f2p_failure = []
    for test_case in gold_results[FAIL_TO_PASS]:
        if test_passed(test_case, eval_sm):
            # Assume silent success for now (test case not in eval_sm)
            f2p_success.append(test_case)
        elif test_failed(test_case, eval_sm):
            f2p_failure.append(test_case)

    # Calculate maintenance metrics
    p2p_success = []
    p2p_failure = []
    for test_case in gold_results[PASS_TO_PASS]:
        if test_passed(test_case, eval_sm):
            p2p_success.append(test_case)
        elif test_failed(test_case, eval_sm):
            p2p_failure.append(test_case)
    
    results = {
        FAIL_TO_PASS: {
            "success": f2p_success,
            "failure": f2p_failure,
        },
        PASS_TO_PASS: {
            "success": p2p_success,
            "failure": p2p_failure,
        }
    }

    f2f_success = []
    f2f_failure = []
    p2f_success = []
    p2f_failure = []
    if calculate_to_fail:
        # Calculate "extra credit" metrics
        for test_case in gold_results[FAIL_TO_FAIL]:
            if test_passed(test_case, eval_sm):
                f2f_success.append(test_case)
            elif test_failed(test_case, eval_sm):
                f2f_failure.append(test_case)

        # Calculate not considered metrics
        for test_case in gold_results[PASS_TO_FAIL]:
            if test_passed(test_case, eval_sm):
                p2f_success.append(test_case)
            elif test_failed(test_case, eval_sm):
                p2f_failure.append(test_case)

    results.update({
        FAIL_TO_FAIL: {
            "success": f2f_success,
            "failure": f2f_failure,
        },
        PASS_TO_FAIL: {
            "success": p2f_success,
            "failure": p2f_failure,
        }
    })
    return results


def get_eval_reports_for_logs(
    eval_logs: List,
    swe_bench_path: str,
    callback: callable = None,
    verbose: bool = False,
) -> (Dict, Dict):
    """
    Wrapper for getting eval report for a list of evaluation log paths.

    Args:
        eval_logs (list): list of paths to evaluation logs
        swe_bench_path (str): path to eval task instances (swe-bench.json)
        callback (callable): callback function for evaluation logs
        verbose (bool): whether to print verbose output
    Returns:
        reports_patch_success (dict): dict of eval reports for patch apply successes
        reports_patch_failure (dict): dict of eval reports for patch apply failures
    """
    reports_patch_success = {}
    reports_patch_failure = {}
    eval_refs = json.load(open(swe_bench_path, "r"))
    eval_refs = {t['instance_id']: t for t in eval_refs}

    for eval_log in eval_logs:
        # Remove task instances that do not satisfy callback
        if callback is not None and not callback(eval_log):
            continue

        # Get gold results
        instance_id = get_id_from_lp(eval_log)
        if instance_id not in eval_refs:
            if verbose:
                print(f"Gold results not found for {instance_id}")
            continue

        gold_results = eval_refs[instance_id]

        # Get eval logs
        eval_sm, has_report = get_logs_eval(eval_log)

        if not has_report:
            # If eval patch failed to apply, convert to report
            # format with tests as failures
            reports_patch_failure[get_file_name_from_lp(eval_log)] = {
                test_type: {"success": [], "failure": tests}
                for test_type, tests in gold_results.items()
            }
            continue

        # Compare eval status map and gold status map
        report = get_eval_report(eval_sm, gold_results)
        reports_patch_success[get_file_name_from_lp(eval_log)] = report

    return reports_patch_success, reports_patch_failure


def get_eval_reports_for_dir(
    eval_dir: str, swe_bench_path: str, callback: callable = None, verbose=False
) -> Dict:
    """
    Wrapper for getting eval report for a directory of evaluation logs.

    Args:
        eval_dir (str): path to directory of evaluation logs
        (See get_eval_reports_for_logs for other args)
    """
    if not os.path.exists(eval_dir):
        raise ValueError(f"Path {eval_dir} does not exist")
    logs_list = [x for x in glob.glob(os.path.join(eval_dir, "*.log"))]
    return get_eval_reports_for_logs(logs_list, swe_bench_path, callback, verbose)


### MARK - Model Evaluation Summary


def get_model_eval_summary(
    predicts_path: str,
    eval_dir: str,
    swe_bench_path: str,
    repo: str = None,
):
    """
    Generate a summary of model evaluation results.

    Args:
        predicts_path (str): path to predictions file
        eval_dir (str): path to directory of evaluation logs
        swe_bench_path (str): path to eval references (swe-bench-eval-refs.json)
        repo (str): if given, repo name to limit evaluation to
    """
    # Load Predictions
    preds = []
    with open(predicts_path, "r") as f:
        for line in f.readlines():
            preds.append(json.loads(line))

    # Filter by repo if provided
    criteria_eval_sm = None
    if repo is not None:
        criteria_pred = lambda pred: repo in pred["instance_id"]
        criteria_eval_sm = lambda eval_log: repo in eval_log
        preds = [x for x in preds if criteria_pred(x)]

    # Get reports
    reports_patch_success, reports_patch_failure = get_eval_reports_for_dir(
        eval_dir, swe_bench_path, callback=criteria_eval_sm, verbose=False
    )

    # Print reports for different granularities of patch success/failure
    summary = {
        "repo": repo if repo is not None else "all",
        "total_predictions": len(preds),
    }
    reports_by_patch_status = [
        ("Patch Apply Success", [reports_patch_success]),
        (
            "Patch Apply Success + Failure",
            [reports_patch_success, reports_patch_failure],
        ),
    ]
    format_dec = lambda x: round(x * 100, 2)
    for report_by_patch_status in reports_by_patch_status:
        r = [list(x.values()) for x in report_by_patch_status[1]]
        r = [item for sublist in r for item in sublist]

        resolutions = Counter([get_resolution_status(_r) for _r in r])
        summary[report_by_patch_status[0]] = {
            "f2p_weighted": format_dec(compute_fail_to_pass_weighted(r)),
            "p2p_weighted": format_dec(compute_pass_to_pass_weighted(r)),
            "f2p_unweighted": format_dec(compute_fail_to_pass_unweighted(r)),
            "p2p_unweighted": format_dec(compute_pass_to_pass_unweighted(r)),
            "cases": report_by_patch_status[1],
            "case_resolution_counts": dict(resolutions),
            "case_resolution_rates": {
                k: round(v / len(r) * 100, 2) for k, v in resolutions.items()
            },
        }

    return summary


def get_model_report(
    model: str, predict_path: str, swe_bench_path: str, log_dir_path: str
):
    """
    Generate a report of model evaluation results from predictions, task instances,
    and evaluation logs.

    Args:
        model (str): model name
        predict_path (str): path to predictions file
        swe_bench_path (str): path to eval references (swe-bench-eval-refs.json)
        log_dir_path (str): path to directory of evaluation logs
    Returns:
        report_map (dict): map of repo to report
    """
    eval_refs = json.load(open(swe_bench_path, "r"))
    eval_refs = [{key: t[key] for key in ["instance_id", "FAIL_TO_PASS", "PASS_TO_PASS"]} for t in eval_refs]
    eval_refs = {t['instance_id']: t for t in eval_refs}

    # Get predictions
    predictions = []
    if predict_path.endswith("jsonl"):
        with open(predict_path, "r") as f:
            for line in f.readlines():
                predictions.append(json.loads(line))
    else:
        predictions = json.load(open(predict_path, "r"))
    report_map = {}

    # Iterate through predictions
    for p in predictions:
        repo = p["instance_id"].split(".")[0].rsplit("-", 1)[0].replace("__", "/")
        if repo not in report_map:
            report_map[repo] = {
                "none": [],
                "generated": [],
                "with_logs": [],
                "applied": [],
                "resolved": [],
            }

        # Check if the model patch exists
        if p["model_patch"] == None:
            report_map[repo]["none"].append(p['instance_id'])
            continue
        report_map[repo]["generated"].append(p['instance_id'])

        # Get log file
        log_path = os.path.join(log_dir_path, f"{p['instance_id']}.{model}.eval.log")
        if not os.path.exists(log_path):
            continue
        report_map[repo]["with_logs"].append(p['instance_id'])

        # Get evaluation logs
        eval_sm, found = get_logs_eval(log_path)

        if not found:
            continue
        report_map[repo]["applied"].append(p['instance_id'])

        report = get_eval_report(eval_sm, eval_refs[p["instance_id"]])
        if get_resolution_status(report) == "RESOLVED_FULL":
            report_map[repo]["resolved"].append(p['instance_id'])

    return report_map

metrics/report.py
princeton-nlp/SWE-bench
import argparse, glob, json, logging, os, re, requests, subprocess, sys

from multiprocessing import Pool, Manager

sys.path.append("../harness")
from constants import (
    SWE_BENCH_URL_RAW,
    MAP_REPO_TO_VERSION_PATHS,
    MAP_REPO_TO_VERSION_PATTERNS,
)
from typing import Dict
from utils import get_instances, split_instances

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


INSTALL_CMD = {
    "pytest-dev/pytest": "pip install -e .",
    "matplotlib/matplotlib": "python -m pip install -e .",
    "pydata/xarray": "pip install -e .",
}


def _find_version_in_text(text: str, instance: Dict) -> str:
    """
    Helper function for applying regex patterns to look for versions in text

    Args:
        text (str): Text to search
        instance (dict): Instance to find version for
    Returns:
        str: Version text, if found
    """
    for pattern in MAP_REPO_TO_VERSION_PATTERNS[instance["repo"]]:
        matches = re.search(pattern, text)
        if matches is not None:
            return str(matches.group(1)).replace(" ", "")


def get_version(instance, is_build=False, path_repo=None):
    """
    Function for looking up the version of a task instance.

    If is_build is True, then the version is looked up by 1. building the repo
    at the instance's base commit, 2. activating the conda environment, and 3.
    looking for the version according to a predefined list of paths.

    Otherwise, the version is looked up by searching GitHub at the instance's
    base commit for the version according to a predefined list of paths.

    Args:
        instance (dict): Instance to find version for
        is_build (bool): Whether to build the repo and look for the version
        path_repo (str): Path to repo to build
    Returns:
        str: Version text, if found
    """
    keep_major_minor = lambda x, sep: ".".join(x.strip().split(sep)[:2])
    paths_to_version = MAP_REPO_TO_VERSION_PATHS[instance["repo"]]
    version = None
    for path_to_version in paths_to_version:
        init_text = None
        if is_build and path_repo is not None:
            version_path_abs = os.path.join(path_repo, path_to_version)
            if os.path.exists(version_path_abs):
                logger.info(f"Found version file at {path_to_version}")
                with open(path_to_version, "r") as f:
                    init_text = f.read()
        else:
            url = os.path.join(
                SWE_BENCH_URL_RAW,
                instance["repo"],
                instance["base_commit"],
                path_to_version,
            )
            init_text = requests.get(url).text
        version = _find_version_in_text(init_text, instance)
        if version is not None:
            if "." in version:
                version = keep_major_minor(version, ".")
            if "," in version:
                version = keep_major_minor(version, ",")
            version = re.sub("[^0-9\.]", "", version)
            return version
    return version


def map_version_to_task_instances(task_instances: list) -> Dict:
    """
    Create a map of version key to list of task instances

    Args:
        task_instances (list): List of task instances
    Returns:
        dict: Map of version key to list of task instances
    """
    return_map = {}
    if "version" in task_instances[0]:
        for instance in task_instances:
            version = instance["version"]
            if version not in return_map:
                return_map[version] = []
            return_map[version].append(instance)
        return return_map
    for instance in task_instances:
        version = get_version(instance)
        if version not in return_map:
            return_map[version] = []
        return_map[version].append(instance)
    return return_map


def get_versions_from_build(data: Dict):
    """
    Logic for looking up versions by building the repo at the instance's base
    commit and looking for the version according to repo-specific paths.

    Args:
        data (dict): Dictionary of data for building a repo for any task instance
            in a given list.
    """
    data_tasks, path_repo, conda_env, path_conda, save_path = (
        data["data_tasks"],
        data["path_repo"],
        data["conda_env"],
        data["path_conda"],
        data["save_path"],
    )
    # Activate conda environment and set installation command
    cmd_activate = f"source {os.path.join(path_conda, 'bin/activate')}"
    cmd_source = f"source {os.path.join(path_conda, 'etc/profile.d/conda.sh')}"
    cmd_install = INSTALL_CMD[data_tasks[0]["repo"]]

    # Change directory to repo testbed
    cwd = os.getcwd()
    os.chdir(path_repo)

    for instance in data_tasks[::-1]:
        # Reset repo to base commit
        subprocess.run(
            "git restore .", check=True, shell=True, stdout=subprocess.DEVNULL
        )
        subprocess.run(
            "git reset HEAD .", check=True, shell=True, stdout=subprocess.DEVNULL
        )
        subprocess.run(
            "git clean -fd", shell=True, check=True, stdout=subprocess.DEVNULL
        )
        out_check = subprocess.run(
            f"git -c advice.detachedHead=false checkout {instance['base_commit']}",
            shell=True,
            stdout=subprocess.DEVNULL,
        )
        if out_check.returncode != 0:
            logger.error(f"[{instance['instance_id']}] Checkout failed")
            continue

        # Run installation command in repo
        out_install = subprocess.run(
            f"{cmd_source}; {cmd_activate} {conda_env}; {cmd_install}",
            shell=True,
            stdout=subprocess.DEVNULL,
        )
        if out_install.returncode != 0:
            logger.error(f"[{instance['instance_id']}] Installation failed")
            continue

        # Look up version according to repo-specific paths
        version = get_version(instance, is_build=True, path_repo=path_repo)
        instance["version"] = version
        logger.info(f'For instance {instance["instance_id"]}, version is {version}')

    # Save results
    with open(save_path, "w") as f:
        json.dump(data_tasks, fp=f)
    os.chdir(cwd)


def get_versions_from_web(data: Dict):
    """
    Logic for looking up versions by searching GitHub at the instance's base
    commit and looking for the version according to repo-specific paths.

    Args:
        data (dict): Dictionary of data for searching GitHub for any task instance
            in a given list.
    """
    data_tasks, save_path = data["data_tasks"], data["save_path"]
    version_not_found = data["not_found_list"]
    for instance in data_tasks:
        version = get_version(instance)
        if version is not None:
            instance["version"] = version
            logger.info(f'For instance {instance["instance_id"]}, version is {version}')
        elif version_not_found is not None:
            logger.info(f'[{instance["instance_id"]}]: version not found')
            version_not_found.append(instance)
    with open(save_path, "w") as f:
        json.dump(data_tasks, fp=f)


def merge_results(instances_path: str, repo_prefix: str, output_dir: str = None) -> int:
    """
    Helper function for merging JSON result files generated from multiple threads.

    Args:
        instances_path (str): Path to original task instances without versions
        repo_prefix (str): Prefix of result files (repo name)
        output_dir (str): Path to save merged results to
    Returns:
        int: Number of instances in merged results
    """
    # Merge values from result JSON files into a single list
    merged = []
    for task_with_version_path in glob.glob(f"{repo_prefix}_versions_*.json"):
        with open(task_with_version_path, "r") as f:
            task_with_version = json.load(f)
            merged.extend(task_with_version)
        os.remove(task_with_version_path)

    # Save merged results to original task instances file's path with `_versions` suffix
    old_path_file = instances_path.split("/")[-1]
    instances_path_new = f"{old_path_file.split('.')[0]}_versions.json"
    if output_dir is not None:
        instances_path_new = os.path.join(output_dir, instances_path_new)
    with open(f"{instances_path_new}", "w") as f:
        json.dump(merged, fp=f)
    logger.info(f"Saved merged results to {instances_path_new} ({len(merged)} instances)")
    return len(merged)


def main(args):
    """
    Main function for looking up versions for task instances.
    """
    # Get task instances + split into groups for each thread
    data_tasks = get_instances(args.instances_path)
    data_task_lists = split_instances(data_tasks, args.num_workers)
    repo_prefix = data_tasks[0]["repo"].replace("/", "__")

    logger.info(
        f"Getting versions for {len(data_tasks)} instances for {data_tasks[0]['repo']}"
    )
    logger.info(
        f"Split instances into {len(data_task_lists)} groups with lengths {[len(x) for x in data_task_lists]}"
    )

    # If retrieval method includes GitHub, then search GitHub for versions via parallel call
    if any([x == args.retrieval_method for x in ["github", "mix"]]):
        manager = Manager()
        shared_result_list = manager.list()
        pool = Pool(processes=args.num_workers)
        pool.map(
            get_versions_from_web,
            [
                {
                    "data_tasks": data_task_list,
                    "save_path": f"{repo_prefix}_versions_{i}.json"
                    if args.retrieval_method == "github"
                    else f"{repo_prefix}_versions_{i}_web.json",
                    "not_found_list": shared_result_list
                    if args.retrieval_method == "mix"
                    else None,
                }
                for i, data_task_list in enumerate(data_task_lists)
            ],
        )
        pool.close()
        pool.join()

        if args.retrieval_method == "github":
            # If retrieval method is just GitHub, then merge results and return
            assert len(data_tasks) == merge_results(
                args.instances_path, repo_prefix, args.output_dir
            )
            return
        elif args.retrieval_method == "mix":
            # Otherwise, remove instances that were found via GitHub from the list
            shared_result_list = list(shared_result_list)
            total_web = len(data_tasks) - len(shared_result_list)
            logger.info(f"Retrieved {total_web} versions from web")
            data_task_lists = split_instances(shared_result_list, args.num_workers)
            logger.info(
                f"Split instances into {len(data_task_lists)} groups with lengths {[len(x) for x in data_task_lists]} for build"
            )

    # Check that all required arguments for installing task instances are present
    assert any([x == args.retrieval_method for x in ["build", "mix"]])
    assert all([x in args for x in ["testbed", "path_conda", "conda_env"]])
    conda_exec = os.path.join(args.path_conda, "bin/conda")

    cwd = os.getcwd()
    os.chdir(args.testbed)
    for x in range(0, args.num_workers):
        # Clone git repo per thread
        testbed_repo_name = f"{repo_prefix}__{x}"
        if not os.path.exists(testbed_repo_name):
            logger.info(
                f"Creating clone of {data_tasks[0]['repo']} at {testbed_repo_name}"
            )
            cmd_clone = (
                f"git clone git@github.com:swe-bench/{repo_prefix} {testbed_repo_name}"
            )
            subprocess.run(cmd_clone, shell=True, check=True, stdout=subprocess.DEVNULL)
        else:
            logger.info(
                f"Repo for {data_tasks[0]['repo']} exists: {testbed_repo_name}; skipping..."
            )
        # Clone conda environment per thread
        conda_env_name = f"{args.conda_env}_clone_{x}"
        if not os.path.exists(os.path.join(args.path_conda, "envs", conda_env_name)):
            logger.info(f"Creating clone of {args.conda_env} at {conda_env_name}")
            cmd_clone_env = f"{conda_exec} create --name {conda_env_name} --clone {args.conda_env} -y"
            subprocess.run(
                cmd_clone_env, shell=True, check=True, stdout=subprocess.DEVNULL
            )
        else:
            logger.info(
                f"Conda clone for thread {x} exists: {conda_env_name}; skipping..."
            )
    os.chdir(cwd)

    # Create pool tasks
    pool_tasks = []
    for i in range(0, args.num_workers):
        testbed_repo_name = f"{repo_prefix}__{i}"
        pool_tasks.append(
            {
                "data_tasks": data_task_lists[i],
                "path_repo": os.path.join(args.testbed, testbed_repo_name),
                "conda_env": f"{args.conda_env}_clone_{i}",
                "path_conda": args.path_conda,
                "save_path": os.path.join(cwd, f"{repo_prefix}_versions_{i}.json"),
            }
        )

    # Parallelized call
    pool = Pool(processes=args.num_workers)
    pool.map(get_versions_from_build, pool_tasks)
    pool.close()
    pool.join()

    # Check that correct number of instances were versioned
    if args.retrieval_method == "mix":
        assert (
            len(data_tasks)
            == merge_results(args.instances_path, repo_prefix, args.output_dir) + total_web
        )
    elif args.retrieval_method == "build":
        assert len(data_tasks) == merge_results(
            args.instances_path, repo_prefix, args.output_dir
        )

    # Remove testbed repo and conda environments
    if args.cleanup:
        cwd = os.getcwd()
        os.chdir(args.testbed)
        for x in range(0, args.num_workers):
            # Remove git repo
            testbed_repo_name = f"{repo_prefix}__{x}"
            subprocess.run(f"rm -rf {testbed_repo_name}", shell=True, check=True)

            # Remove conda environment
            cmd_rm_env = (
                f"{conda_exec} remove --name {args.conda_env}_clone_{x} --all -y"
            )
            subprocess.run(cmd_rm_env, shell=True, check=True)
        os.chdir(cwd)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--instances_path", required=True, type=str, default=None, help="Path to task instances")
    parser.add_argument("--retrieval_method", required=True, choices=["build", "mix", "github"], default="github", help="Method to retrieve versions")
    parser.add_argument("--cleanup", action="store_true", help="Remove testbed repo and conda environments")
    parser.add_argument("--conda_env", type=str, default=None, help="Conda environment to use")
    parser.add_argument("--path_conda", type=str, default=None, help="Path to conda")
    parser.add_argument("--num_workers", type=int, default=1, help="Number of threads to use")
    parser.add_argument("--output_dir", type=str, default=None, help="Path to save results")
    parser.add_argument("--testbed", type=str, default=None, help="Path to testbed repo")
    args = parser.parse_args()
    main(args)

versioning/get_versions.py
princeton-nlp/SWE-bench
import json
import os
import re
import requests
import sys

from datetime import datetime

sys.path.append("../harness")
from utils import get_instances

PATH_TASKS_ASTROPY = "<path to astropy task instances>"

# Get raw astropy dataset
data_tasks = get_instances(PATH_TASKS_ASTROPY)

# Get version to date from astropy homepage
resp = requests.get("https://docs.astropy.org/en/latest/changelog.html")
pattern = (
    r'<a class="reference internal nav-link" href="#version-(.*)">Version (.*)</a>'
)
matches = re.findall(pattern, resp.text)
matches = list(set(matches))

# Get (date, version) pairs
date_format = "%Y-%m-%d"
keep_major_minor = lambda x, sep: ".".join(x.strip().split(sep)[:2])

# Iterate through matches, construct (version, date) pairs
times = []
for match in matches:
    match_parts = match[1].split(" ")
    version, date = match_parts[0], match_parts[1].strip(")").strip("(")
    version = keep_major_minor(version, ".")
    date_obj = datetime.strptime(date, date_format)
    times.append((date_obj.strftime("%Y-%m-%d"), version))

# Group times by major/minor version
map_version_to_times = {}
for time in times:
    if time[1] not in map_version_to_times:
        map_version_to_times[time[1]] = []
    map_version_to_times[time[1]].append(time[0])

# Pick the most recent time as the version cut off date
version_to_time = [(k, max(v)) for k, v in map_version_to_times.items()]
version_to_time = sorted(version_to_time, key=lambda x: x[0])[::-1]

# Assign version to each task instance
for task in data_tasks:
    created_at = task["created_at"].split("T")[0]
    for t in version_to_time:
        found = False
        if t[1] < created_at:
            task["version"] = t[0]
            found = True
            break
    if not found:
        task["version"] = version_to_time[-1][0]

# Construct map of versions to task instances
map_v_to_t = {}
for task in data_tasks:
    if task["version"] not in map_v_to_t:
        map_v_to_t[task["version"]] = []
    map_v_to_t[task["version"]].append(t)

# Save matplotlib versioned data to repository
with open(
    os.path.join(PATH_TASKS_ASTROPY, "astropy-task-instances_versions.json"),
    "w",
) as f:
    json.dump(data_tasks, fp=f)

versioning/get_versions_astropy.py
princeton-nlp/SWE-bench
import json
import os
import re
import requests
import sys

from datetime import datetime

sys.path.append("../harness")
from utils import get_instances

PATH_TASKS_MATPLOTLIB = "<path to matplotlib task instances>"

# Get raw matplotlib dataset
data_tasks = get_instances(PATH_TASKS_MATPLOTLIB)

# Get version to date from matplotlib home page
resp = requests.get("https://matplotlib.org/stable/users/release_notes#past-versions")
pattern = r'<a class="reference internal" href="prev_whats_new/whats_new_(.*).html">What\'s new in Matplotlib (.*)</a>'
matches = re.findall(pattern, resp.text)
matches = list(set(matches))

# Get (date, version) pairs
date_format = "%b %d, %Y"
keep_major_minor = lambda x, sep: ".".join(x.strip().split(sep)[:2])

times = []
for match in matches:
    version, s = match[0], match[1]
    if "(" not in s:
        continue
    version = keep_major_minor(version, ".")
    date_string = s[s.find("(") + 1 : s.find(")")]
    date_obj = datetime.strptime(date_string, date_format)
    times.append((date_obj.strftime("%Y-%m-%d"), version))
times = sorted(times, key=lambda x: x[0])[::-1]

for task in data_tasks:
    created_at = task["created_at"].split("T")[0]
    for t in times:
        if t[0] < created_at:
            task["version"] = t[1]
            break

# Construct map of versions to task instances
map_v_to_t = {}
for t in data_tasks:
    if t["version"] not in map_v_to_t:
        map_v_to_t[t["version"]] = []
    map_v_to_t[t["version"]].append(t)

# Save matplotlib versioned data to repository
with open(
    os.path.join(PATH_TASKS_MATPLOTLIB, "matplotlib-task-instances_versions.json"),
    "w",
) as f:
    json.dump(data_tasks, fp=f)

versioning/get_versions_matplotlib.py
princeton-nlp/SWE-bench
import json
import os
import re
import requests
import sys

from datetime import datetime

sys.path.append("../harness")
from utils import get_instances

PATH_TASKS_XARRAY = "<path to xarray task instances>"

# Get raw xarray dataset
data_tasks = get_instances(PATH_TASKS_XARRAY)

# Get version to date from xarray home page
resp = requests.get("https://docs.xarray.dev/en/stable/whats-new.html")
pattern = (
    r'<a class="reference internal nav-link( active)?" href="#v(.*)">v(.*) \((.*)\)</a>'
)
matches = re.findall(pattern, resp.text)
matches = list(set(matches))
matches = [x[1:] for x in matches]

# Get (date, version) pairs
date_formats = ["%B %d %Y", "%d %B %Y"]
keep_major_minor = lambda x, sep: ".".join(x.strip().split(sep)[:2])

times = []
for match in matches:
    parts = match[0].split("-")
    version = keep_major_minor(".".join(parts[0:3]), ".")
    date_str = " ".join(parts[3:])

    for f_ in date_formats:
        try:
            date_obj = datetime.strptime(date_str, f_)
            times.append((date_obj.strftime("%Y-%m-%d"), version))
        except:
            continue
        break

times = sorted(times, key=lambda x: x[0])[::-1]

for task in data_tasks:
    created_at = task["created_at"].split("T")[0]
    found = False
    for t in times:
        if t[0] < created_at:
            task["version"] = t[1]
            found = True
            break
    if not found:
        task["version"] = None

# Save xarray versioned data to repository
with open(
    os.path.join(PATH_TASKS_XARRAY, "xarray-task-instances_versions.json"),
    "w",
) as f:
    json.dump(data_tasks, fp=f)

versioning/get_versions_xarray.py
princeton-nlp/SWE-bench