bigcode/commitpackmeta

1from concurrent.futures import ThreadPoolExecutor2import os3import random4import subprocess5import timeit6 7import datasets8 9"""Example10git init11git remote add origin https://github.com/huggingface/evaluate.git12git fetch --depth 2 origin 9b056cdd5eb95459ae80142014865263e7dd75b813# Get file after change14git checkout FETCH_HEAD -- README.md15# Get file before change16git checkout FETCH_HEAD^ -- README.md17"""18 19# In the multiprocessing case, the below leads to each process creating the same directory20# random = random.Random(42)  # make it reproducible21 22NUM_THREADS = 6423NUM_PROC = 6424# DEBUG_SIZE = 102425 26CWD = os.getcwd()27 28# Shell utils29def run_in_shell(cmd: str, cwd=None, timeout=60):30    return subprocess.run([cmd], capture_output=True, shell=True, cwd=cwd, timeout=timeout)31 32def get_file_contents(commit, old_file, new_file, repo, cwd=None):33    completed = run_in_shell("git init", cwd=cwd)34    completed = run_in_shell("git remote add origin " + repo, cwd=cwd)35    completed = run_in_shell("git fetch --depth 2 origin " + commit, cwd=cwd)36     # If it requires authentication37    if completed.returncode != 0:38        #print("ERRORC1", completed)39        return ("", "", completed.returncode, completed.stderr.decode(errors='ignore'))40    # Optionally do git diff at the same time (Saving code needs to be added)41    # git_diff = run_in_shell(f"git diff {commit}^ {commit}", cwd=cwd).stdout.decode(errors='ignore')42    completed = run_in_shell("git checkout FETCH_HEAD -- " + new_file, cwd=cwd)43    new_contents = run_in_shell("cat " + new_file, cwd=cwd).stdout.decode(errors='ignore')44    completed = run_in_shell("git checkout FETCH_HEAD^ -- " + old_file, cwd=cwd)45    # If there's only a new file, but no old file46    if completed.returncode != 0:47        #print("ERRORC2", completed)48        return (new_contents, "", completed.returncode, completed.stderr.decode(errors='ignore'))49    old_contents = run_in_shell("cat " + old_file, cwd=cwd).stdout.decode(errors='ignore')50    return (new_contents, old_contents, completed.returncode, completed.stderr.decode(errors='ignore'))51 52def get_diff(ex):53    commit_id = ex["commit"]54    repos = list(set(ex["repos"].split(",")))55    old_file = ex["old_file"]56    new_file = ex["new_file"]57    # Initialize58    returncode = 059    stderr = "unknown"60 61    for i, repo in enumerate(repos):62        repo = "https://xxx:xxx@github.com/" + repo + ".git"63        # Create a random directory to store the repo64        random_dir = CWD + "/" + str(random.randint(0, 1000000))65        # Can take very long when running many processes66        run_in_shell("mkdir " + random_dir, timeout=300)67        try:68            new_contents, old_contents, returncode, stderr = get_file_contents(commit_id, old_file, new_file, repo, cwd=random_dir)69        except Exception as e:70            #print("ERROR", commit_id, old_file, new_file, repo, str(random_dir), e)71            # Break in case of many repos that all lead us nowhere72            if i > 10:73                break74            continue75        finally:76            run_in_shell("rm -rf " + random_dir) # clean up again77        ex["new_contents"] = new_contents78        ex["old_contents"] = old_contents79        ex["returncode"] = returncode80        ex["stderr"] = stderr81        return ex82    # If no repo worked83    ex["new_contents"] = ""84    ex["old_contents"] = ""85    ex["returncode"] = returncode86    ex["stderr"] = stderr87    return ex88 89def get_diff_multi_threaded_processed(batch):90    with ThreadPoolExecutor(max_workers=NUM_THREADS) as executor:91        # Convert dict of lists to list of dicts then map to threads92        results = list(executor.map(get_diff, [dict(zip(batch,t)) for t in zip(*batch.values())]))93        # Convert list of dicts to dict of lists94        return {k: [dic[k] for dic in results] for k in results[0]}95 96if __name__ == "__main__":97    # git clone https://huggingface.co/datasets/bigcode/github-commits98    ds = datasets.load_dataset("./github-commits", use_auth_token=True)["train"]99 100    ### OPTIONAL FILTERING ###101    #"""102    java = [".java"]103    javascript = [104        ".js",105        "._js",106        ".bones",107        ".es6",108        ".jake",109        ".jsb",110        ".jscad",111        ".jsfl",112        ".jsm",113        ".jss",114        ".njs",115        ".pac",116        ".sjs",117        ".ssjs",118        ".xsjs",119        ".xsjslib"120    ]121    python = [122        ".py",123        ".bzl",124        ".gyp",125        ".lmi",126        ".pyde",127        ".pyp",128        ".pyt",129        ".pyw",130        ".tac",131        ".wsgi",132        ".xpy"133    ]134 135    import json136    with open("programming-languages.json", "r") as f:137        extensions = json.load(f)138    suffices = tuple([suffix for suffices in extensions.values() for suffix in suffices])139    def filter_extension(ex):140        return ex["new_file"].endswith(suffices)141 142    def filter_extension_python(ex):143        return ex["new_file"].endswith(python)144 145    def filter_update(ex):146        return ex["message"] != "Update " + ex["old_file"]147 148    filter_msg = ["initial commit", "please\n", "please", "lalala"]149 150    def filter_misc(ex):151        return ex["message"] not in filter_msg152 153    # Removes ~10M154    ds = ds.filter(filter_extension, num_proc=NUM_PROC)155    print("After Extension filter", len(ds))156    # Removes ~1M157    ds = ds.filter(filter_update, num_proc=NUM_PROC)158    print("After Update filter", len(ds))159    #ds = ds.filter(filter_extension_python, num_proc=NUM_PROC)160    #print("After Python filter", len(ds))161    ds = ds.filter(filter_misc, num_proc=NUM_PROC)162    print("After Misc filter", len(ds))163    #ds = ds.select(range(DEBUG_SIZE))164    START = 0 # Modify for each instance (0 - 7)165    samples_per_instance = 64 * 64 * 64 * 32 # 8_388_608166    select_start = START * samples_per_instance167    select_end = START * samples_per_instance + samples_per_instance168    ds = ds.select(range(select_start, select_end))169    print(f"Going from {select_start} till {select_end}")170 171    #"""172    ### END FILTERING ###173 174    175 176    ### ALTERNATIVELY LOAD EXISTING SPLIT ###177    """178    path = "github-commits-diff/data/diffs_50331648_58720256.jsonl"179    ds = datasets.load_dataset("json", data_files=path)180    sub_ds = ds.filter(lambda x: x['stderr'].startswith("fatal: unable to acces"))181    skipped_ds = ds.filter(lambda x: not(x['stderr'].startswith("fatal")))182    183    datasets.concatenate_datasets((184        skipped_ds,185        sub_ds.map(get_diff_multi_threaded_processed, num_proc=NUM_PROC, batch_size=NUM_THREADS, batched=True),186    )).to_json(path.replace(".", "_new."), num_proc=NUM_PROC)187    exit()188    """189    ### END LOAD EXISTING ###190 191    def run_multi_processing_threading():192        ds.map(get_diff_multi_threaded_processed, num_proc=NUM_PROC, batch_size=NUM_THREADS, batched=True).to_json(f"diffs_{select_start}_{select_end}.jsonl", num_proc=NUM_PROC)193 194    # Benchmarking195    #NUM_TRIALS = 1196    #print(f"Timing multithreading + multiprocessing using {NUM_THREADS} threads and {NUM_PROC} processes")197    #time = timeit.timeit(stmt=run_multi_processing_threading, number=NUM_TRIALS)198    #print("Time:", time)199    #with open("mpt.txt", "w") as f:200    #    f.write(str(time))201 202    # Running203    run_multi_processing_threading()