get hash of content field return hash insecurehashlib md5re subpattern examplecontent encodeutf8 hexdigest def linestatsexample calculates mean and max line length of file alphafrac np meanc isalnum for c in examplecontent return alphafrac alphafrac def checkuniquesexample uniques check if file is autogenerated by looking for keywords in the first few lines of the file keywords autogenerated autogenerated automatically generated lines examplecontent splitlines for line in ziprangescanwidth lines for keyword in keywords if keyword in line lower return autogenerated true else return autogenerated false def isconfigortestexample scanwidth5 coeff0 05 keywords unit tests test file configuration file lines examplecontent splitlines countconfig 0 counttest 0 first test for line in ziprangescanwidth lines for keyword in keywords if keyword in line lower return configortest true second test nlines examplecontent countn threshold intcoeff nlines for line in lines countconfig line lower countconfig counttest line lower counttest if countconfig threshold or counttest threshold return configortest true return configortest false def hasnokeywordsexample check if file uses symbol less than minimum times lines examplecontent splitlines counter 0 for line in lines counter line lower count if counter minimum return hasfewassignments false return hasfewassignments true def chartokenratioexample chain all preprocessing steps into one function to not fill cache results results updategethashexample results updatelinestatsexample results updatealphastatsexample results updatechartokenratioexample results updateisautogeneratedexample results updateisconfigortestexample results updatehasnokeywordsexample results updatehasfewassignmentsexample return results def filterexample uniques args compress a file with gzip with openfilepath rb as fin with gzip openstrfilepath gz wb compresslevel6 as fout shutil copyfileobjfin fout os unlinkfilepath settings parser hfargumentparserpreprocessingarguments args parser parseargs if args numworkers is none args numworkers multiprocessing cpucount tokenizer autotokenizer frompretrainedargs tokenizerdir load dataset tstart time time ds loaddatasetargs datasetname splittrain printftime to load dataset time timetstart 2f run preprocessing tstart time time ds ds mappreprocess numprocargs numworkers printftime to preprocess dataset time timetstart 2f deduplicate hashes uniques setds uniquehash frac lenuniques lends printffraction of duplicates 1frac 2 deduplicate data and apply heuristics tstart time time dsfilter ds filterfilter fnkwargsuniques uniques args args printftime to filter dataset time timetstart 2f printfsize of filtered dataset lendsfilter deduplicate with minhash and jaccard similarity if args neardeduplication tstart time time dsfilter duplicateclusters deduplicatedatasetdsfilter args jaccardthreshold printftime to deduplicate dataset time timetstart 2f printfsize of deduplicate dataset lendsfilter save data in batches of samplesperfile outputdir pathargs outputdir outputdir mkdirexistoktrue save duplicateclusters in the outputdir as artifacts not sure it is the right place the save it if args neardeduplication with openoutputdir duplicateclusters json w as f json dumpduplicateclusters f datadir outputdir data datadir mkdirexistoktrue tstart time time for filenumber index in enumeraterange0 lendsfilter args samplesperfile filepath strdatadir ffilefilenumber1 012 json endindex minlendsfilter index args samplesperfile dsfilter selectlistrangeindex endindex tojsonfilepath compressfilefilepath printftime to save dataset time timetstart 2f get hash of content field calculates mean and max line length of file calculates mean and max line length of file check if current hash is still in set of unique hashes and remove if true check if file is autogenerated by looking for keywords in the first few lines of the file check if file is a configuration file or a unit test by 1 looking for keywords in the first few lines of the file 2 counting number of occurence of the words config and test with respect to number of lines first test second test check if a python file has none of the keywords for funcion class for loop while loop check if file uses symbol less than minimum times compute character token ratio of the file with tokenizer chain all preprocessing steps into one function to not fill cache filter dataset with heuristics config test and has_no_keywords files are removed with a given probability compress a file with g zip settings load dataset run preprocessing deduplicate hashes deduplicate data and apply heuristics deduplicate with minhash and jaccard similarity save data in batches of samples_per_file save duplicate_clusters in the output_dir as artifacts not sure it is the right place the save it
import gzip
import json
import multiprocessing
import os
import re
import shutil
import time
from pathlib import Path

import numpy as np
from arguments import PreprocessingArguments
from datasets import load_dataset
from huggingface_hub.utils import insecure_hashlib
from minhash_deduplication import deduplicate_dataset

from transformers import AutoTokenizer, HfArgumentParser


PATTERN = re.compile(r"\s+")


def get_hash(example):
    
    return {"hash": insecure_hashlib.md5(re.sub(PATTERN, "", example["content"]).encode("utf-8")).hexdigest()}


def line_stats(example):
    
    line_lengths = [len(line) for line in example["content"].splitlines()]
    return {"line_mean": np.mean(line_lengths), "line_max": max(line_lengths)}


def alpha_stats(example):
    
    alpha_frac = np.mean([c.isalnum() for c in example["content"]])
    return {"alpha_frac": alpha_frac}


def check_uniques(example, uniques):
    
    if example["hash"] in uniques:
        uniques.remove(example["hash"])
        return True
    else:
        return False


def is_autogenerated(example, scan_width=5):
    
    keywords = ["auto-generated", "autogenerated", "automatically generated"]
    lines = example["content"].splitlines()
    for _, line in zip(range(scan_width), lines):
        for keyword in keywords:
            if keyword in line.lower():
                return {"autogenerated": True}
    else:
        return {"autogenerated": False}


def is_config_or_test(example, scan_width=5, coeff=0.05):
    

    keywords = ["unit tests", "test file", "configuration file"]
    lines = example["content"].splitlines()
    count_config = 0
    count_test = 0
    
    for _, line in zip(range(scan_width), lines):
        for keyword in keywords:
            if keyword in line.lower():
                return {"config_or_test": True}
    
    nlines = example["content"].count("\n")
    threshold = int(coeff * nlines)
    for line in lines:
        count_config += line.lower().count("config")
        count_test += line.lower().count("test")
        if count_config > threshold or count_test > threshold:
            return {"config_or_test": True}
    return {"config_or_test": False}


def has_no_keywords(example):
    
    keywords = ["def ", "class ", "for ", "while "]
    lines = example["content"].splitlines()
    for line in lines:
        for keyword in keywords:
            if keyword in line.lower():
                return {"has_no_keywords": False}
    return {"has_no_keywords": True}


def has_few_assignments(example, minimum=4):
    
    lines = example["content"].splitlines()
    counter = 0
    for line in lines:
        counter += line.lower().count("=")
        if counter > minimum:
            return {"has_few_assignments": False}
    return {"has_few_assignments": True}


def char_token_ratio(example):
    
    input_ids = tokenizer(example["content"], truncation=False)["input_ids"]
    ratio = len(example["content"]) / len(input_ids)
    return {"ratio": ratio}


def preprocess(example):
    
    results = {}
    results.update(get_hash(example))
    results.update(line_stats(example))
    results.update(alpha_stats(example))
    results.update(char_token_ratio(example))
    results.update(is_autogenerated(example))
    results.update(is_config_or_test(example))
    results.update(has_no_keywords(example))
    results.update(has_few_assignments(example))
    return results


def filter(example, uniques, args):
    
    if not check_uniques(example, uniques):
        return False
    elif example["autogenerated"]:
        return False
    elif example["line_max"] > args.line_max:
        return False
    elif example["line_mean"] > args.line_mean:
        return False
    elif example["alpha_frac"] < args.alpha_frac:
        return False
    elif example["ratio"] < args.min_token_ratio:
        return False
    elif example["config_or_test"] and np.random.rand() <= args.filter_proba:
        return False
    elif example["has_no_keywords"] and np.random.rand() <= args.filter_proba:
        return False
    elif example["has_few_assignments"]:
        return False
    else:
        return True


def compress_file(file_path):
    
    with open(file_path, "rb") as f_in:
        with gzip.open(str(file_path) + ".gz", "wb", compresslevel=6) as f_out:
            shutil.copyfileobj(f_in, f_out)
    os.unlink(file_path)



parser = HfArgumentParser(PreprocessingArguments)
args = parser.parse_args()
if args.num_workers is None:
    args.num_workers = multiprocessing.cpu_count()
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_dir)


t_start = time.time()
ds = load_dataset(args.dataset_name, split="train")
print(f"Time to load dataset: {time.time()-t_start:.2f}")


t_start = time.time()
ds = ds.map(preprocess, num_proc=args.num_workers)
print(f"Time to preprocess dataset: {time.time()-t_start:.2f}")


uniques = set(ds.unique("hash"))
frac = len(uniques) / len(ds)
print(f"Fraction of duplicates: {1-frac:.2%}")


t_start = time.time()
ds_filter = ds.filter(filter, fn_kwargs={"uniques": uniques, "args": args})
print(f"Time to filter dataset: {time.time()-t_start:.2f}")
print(f"Size of filtered dataset: {len(ds_filter)}")


if args.near_deduplication:
    t_start = time.time()
    ds_filter, duplicate_clusters = deduplicate_dataset(ds_filter, args.jaccard_threshold)
    print(f"Time to deduplicate dataset: {time.time()-t_start:.2f}")
    print(f"Size of deduplicate dataset: {len(ds_filter)}")


output_dir = Path(args.output_dir)
output_dir.mkdir(exist_ok=True)



if args.near_deduplication:
    with open(output_dir / "duplicate_clusters.json", "w") as f:
        json.dump(duplicate_clusters, f)

data_dir = output_dir / "data"
data_dir.mkdir(exist_ok=True)

t_start = time.time()
for file_number, index in enumerate(range(0, len(ds_filter), args.samples_per_file)):
    file_path = str(data_dir / f"file-{file_number+1:012}.json")
    end_index = min(len(ds_filter), index + args.samples_per_file)
    ds_filter.select(list(range(index, end_index))).to_json(file_path)
    compress_file(file_path)
print(f"Time to save dataset: {time.time()-t_start:.2f}")