codeparrot/github-code

1import gzip2import multiprocessing3import os4import shutil5import time6from argparse import Namespace7from collections import Counter8import numpy as np9from datasets import load_dataset, utils10import re11from huggingface_hub import Repository12from multiprocessing import Pool13from tqdm import tqdm14 15# Settings16config = {17    "dataset_name": "./data/github",18    "num_workers": 96,19    "line_max": 1000,20    "out_path": "./data/github-code",21    "repo_name": "github-code",22    "org": "lvwerra",23    "shard_size": 1000 << 20}24 25args = Namespace(**config)26 27PATTERN = re.compile(r'\s+')28 29 30def get_hash(example):31    """Get hash of content field."""32    return {"hash": hash(re.sub(PATTERN, '', example["content"]))}33 34 35def line_stats(example):36    """Calculates mean and max line length of file."""37    line_lengths = [len(line) for line in example["content"].splitlines()]38    return {"line_mean": np.mean(line_lengths), "line_max": max(line_lengths)}39 40 41def alpha_stats(example):42    """Calculates mean and max line length of file."""43    alpha_frac = np.mean([c.isalnum() for c in example["content"]])44    return {"alpha_frac": alpha_frac}45 46 47def check_uniques(example, uniques):48    """Check if current hash is still in set of unique hashes and remove if true."""49    if example["hash"] in uniques:50        uniques.remove(example["hash"])51        return True52    else:53        return False54 55 56def is_autogenerated(example, scan_width=5):57    """Check if file is autogenerated by looking for keywords in the first few lines of the file."""58    keywords = ["auto-generated", "autogenerated", "automatically generated"]59    lines = example["content"].splitlines()60    for _, line in zip(range(scan_width), lines):61        for keyword in keywords:62            if keyword in line.lower():63                return {"autogenerated": True}64    else:65        return {"autogenerated": False}66 67 68def preprocess(example):69    """Chain all preprocessing steps into one function to not fill cache."""70    results = dict()71    results.update(get_hash(example))72    results.update(line_stats(example))73    return results74 75 76def filter(example, uniques, args):77    """Filter dataset with heuristics."""78    if not check_uniques(example, uniques):79        return False80    elif example["line_max"] > args.line_max:81        return False82    else:83        return True84    85def save_shard(shard_tuple):86    """Save shard"""87    filename, shard = shard_tuple88    shard.to_parquet(filename)89 90# Load dataset91t_start = time.time()92ds = load_dataset(args.dataset_name, split="train", chunksize=40<<20)93print(f"Time to load dataset: {time.time()-t_start:.2f}")94 95# Run preprocessing96t_start = time.time()97ds = ds.map(preprocess, num_proc=args.num_workers)98print(f"Time to preprocess dataset: {time.time()-t_start:.2f}")99print(ds)100 101# Deduplicate hashes102uniques = set(ds.unique("hash"))103frac = len(uniques) / len(ds)104print(f"Fraction of duplicates: {1-frac:.2%}")105 106# Deduplicate data and apply heuristics107t_start = time.time()108ds = ds.filter(filter, fn_kwargs={"uniques": uniques, "args": args})109ds = ds.remove_columns(["line_mean", "line_max", "copies", "hash"])110print(f"Time to filter dataset: {time.time()-t_start:.2f}")111print(f"Size of filtered dataset: {len(ds)}")112 113 114# Save dataset in repo115repo = Repository(116        local_dir=args.out_path,117        clone_from=args.org + "/" + args.repo_name,118        repo_type="dataset",119        private=True,120        use_auth_token=True,121        git_user="lvwerra",122        git_email="leandro.vonwerra@gmail.com",123        )124 125os.mkdir(args.out_path + "/data")126 127if ds._indices is not None:128    dataset_nbytes = ds.data.nbytes * len(ds._indices) / len(ds.data)129else:130    dataset_nbytes = ds.data.nbytes131 132num_shards = int(dataset_nbytes / args.shard_size) + 1133print(f"Number of shards: {num_shards}")134 135t_start = time.time()136shards = (ds.shard(num_shards=num_shards, index=i, contiguous=True) for i in range(num_shards))137filenames = (f"{args.out_path}/data/train-{index:05d}-of-{num_shards:05d}.parquet" for index in range(num_shards))138 139with Pool(16) as p:140    list(tqdm(p.imap_unordered(save_shard, zip(filenames, shards), chunksize=4), total=num_shards))141print(f"Time to save dataset: {time.time()-t_start:.2f}")142 143# To push to hub run `git add` and `git push` inside dataset repo folder