<commit_before>import h5py
import numpy as np


def get_paths(group, path):
    out = set()
    for key in group.keys():
        if isinstance(group[key], h5py.Group):
            out |= get_paths(group[key], path + "/" + key)
        if isinstance(group[key], h5py.Dataset):
            out.add(path + "/" + key)
    return out


def paths_agree(file1, file2, path, dim):
    try:
        return (path in file1) and (path in file2) and file1[path].shape[dim] == file2[path].shape[dim]
    except IndexError as e:
        #1D arrays do weird things
        return len(file1[path].shape) == len(file2[path].shape) == dim == 1


def h5_merge(infilenames, outfilename, orientation="horiz", reserved_paths=[]):
    files = [h5py.File(filename, "r", driver="core", libver="latest") for filename in infilenames]
    # collect all common paths between the files
    paths = set()
    for file in files:
        paths |= get_paths(file, "")
    dim_ind = 0 if orientation == "horiz" else 1
    merge_paths = [path for path in paths if all([path in file and paths_agree(file, files[0], path, dim_ind) for file in files])]
    outfile = h5py.File(outfilename, "w", driver="core", libver="latest")
    for path in merge_paths:
        concat_axis = 1 if orientation == "horiz" else 0
        if path in reserved_paths:
            outfile.create_dataset(path, data=files[0][path])
        else:
            outfile.create_dataset(path, data=np.concatenate([file[path] for file in files], axis=concat_axis))
    outfile.close()
    for file in files:
        file.close()<commit_msg>h5merge now also writes base sample id<commit_after>import h5py
import os
def get_paths(group, path):
    out = set()
    for key in group.keys():
        if isinstance(group[key], h5py.Group):
            out |= get_paths(group[key], path + "/" + key)
        if isinstance(group[key], h5py.Dataset):
            out.add(path + "/" + key)
    return out


def paths_agree(file1, file2, path, dim):
    try:
        return (path in file1) and (path in file2) and file1[path].shape[dim] == file2[path].shape[dim]
    except IndexError as e:
        #1D arrays do weird things
        return len(file1[path].shape) == len(file2[path].shape) == dim == 1


def h5_merge(infilenames, outfilename, orientation="horiz", reserved_paths=[]):
    files = [h5py.File(filename, "r", driver="core", libver="latest") for filename in infilenames]
    # collect all common paths between the files
    paths = set()
    for file in files:
        paths |= get_paths(file, "")
    dim_ind = 0 if orientation == "horiz" else 1
    merge_paths = [path for path in paths if all([path in file and paths_agree(file, files[0], path, dim_ind) for file in files])]
    outfile = h5py.File(outfilename, "w", driver="core", libver="latest")
    for path in merge_paths:
        concat_axis = 1 if orientation == "horiz" else 0
        if path in reserved_paths:
            outfile.create_dataset(path, data=files[0][path])
        else:
            outfile.create_dataset(path, data=np.concatenate([file[path] for file in files], axis=concat_axis))
    # create a dataset which stores sample ids
    outfile.create_dataset("sampleId", (1, len(infilenames)))
    for i in range(0, len(infilenames)):
        outfile["baseSampleId"][i] = int(os.path.basename(os.path.splitext(infilenames[i])[0])) # not the most elegant, but works for now
    outfile.close()
    for file in files:
        file.close()
{'fraction_non_alphanumeric': 0.1292102206736353, 'fraction_numerical': 0.013937282229965157, 'mean_word_length': 3.1808252427184467, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6178125', 'n_tokens_mistral': 1119, 'n_tokens_neox': 1092, 'n_words': 330}