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 hash_func(text):31    return hashlib.md5(re.sub(PATTERN, '', text).encode("utf-8")).hexdigest()32 33def get_hash(example):34    """Get hash of content field."""35    return {"hash": hash_func(example["content"])}36 37 38def line_stats(example):39    """Calculates mean and max line length of file."""40    line_lengths = [len(line) for line in example["content"].splitlines()]41    return {"line_mean": np.mean(line_lengths), "line_max": max(line_lengths)}42 43 44def alpha_stats(example):45    """Calculates mean and max line length of file."""46    alpha_frac = np.mean([c.isalnum() for c in example["content"]])47    return {"alpha_frac": alpha_frac}48 49 50def check_uniques(example, uniques):51    """Check if current hash is still in set of unique hashes and remove if true."""52    if example["hash"] in uniques:53        uniques.remove(example["hash"])54        return True55    else:56        return False57 58 59def is_autogenerated(example, scan_width=5):60    """Check if file is autogenerated by looking for keywords in the first few lines of the file."""61    keywords = ["auto-generated", "autogenerated", "automatically generated"]62    lines = example["content"].splitlines()63    for _, line in zip(range(scan_width), lines):64        for keyword in keywords:65            if keyword in line.lower():66                return {"autogenerated": True}67    else:68        return {"autogenerated": False}69 70 71def preprocess(example):72    """Chain all preprocessing steps into one function to not fill cache."""73    results = dict()74    results.update(get_hash(example))75    results.update(line_stats(example))76    return results77 78 79def filter(example, uniques, args):80    """Filter dataset with heuristics."""81    if not check_uniques(example, uniques):82        return False83    elif example["line_max"] > args.line_max:84        return False85    else:86        return True87    88def save_shard(shard_tuple):89    """Save shard"""90    filename, shard = shard_tuple91    shard.to_parquet(filename)92 93# Load dataset94t_start = time.time()95ds = load_dataset(args.dataset_name, split="train", chunksize=40<<20)96print(f"Time to load dataset: {time.time()-t_start:.2f}")97 98# Run preprocessing99t_start = time.time()100ds = ds.map(preprocess, num_proc=args.num_workers)101print(f"Time to preprocess dataset: {time.time()-t_start:.2f}")102print(ds)103 104# Deduplicate hashes105uniques = set(ds.unique("hash"))106frac = len(uniques) / len(ds)107print(f"Fraction of duplicates: {1-frac:.2%}")108 109# Deduplicate data and apply heuristics110t_start = time.time()111ds = ds.filter(filter, fn_kwargs={"uniques": uniques, "args": args})112ds = ds.remove_columns(["line_mean", "line_max", "copies", "hash"])113print(f"Time to filter dataset: {time.time()-t_start:.2f}")114print(f"Size of filtered dataset: {len(ds)}")115 116 117# Save dataset in repo118repo = Repository(119        local_dir=args.out_path,120        clone_from=args.org + "/" + args.repo_name,121        repo_type="dataset",122        private=True,123        use_auth_token=True,124        git_user="lvwerra",125        git_email="leandro.vonwerra@gmail.com",126        )127 128os.mkdir(args.out_path + "/data")129 130if ds._indices is not None:131    dataset_nbytes = ds.data.nbytes * len(ds._indices) / len(ds.data)132else:133    dataset_nbytes = ds.data.nbytes134 135num_shards = int(dataset_nbytes / args.shard_size) + 1136print(f"Number of shards: {num_shards}")137 138t_start = time.time()139shards = (ds.shard(num_shards=num_shards, index=i, contiguous=True) for i in range(num_shards))140filenames = (f"{args.out_path}/data/train-{index:05d}-of-{num_shards:05d}.parquet" for index in range(num_shards))141 142with Pool(16) as p:143    list(tqdm(p.imap_unordered(save_shard, zip(filenames, shards), chunksize=4), total=num_shards))144print(f"Time to save dataset: {time.time()-t_start:.2f}")145 146# To push to hub run `git add/commit/push` inside dataset repo folder