<commit_before>from __future__ import print_function

import multiprocessing
import os
import time
from operator import mul
from os.path import join
import warnings

import h5py
import numpy as np

import PyGreentea as pygt

''' where this will be used:
* train()
  * getting unordered batches of a dataset, specified with offset, input size and output size
  * behavior should be like a queue. at the end of each training iteration, train() specifies a replacement
* process()
  * getting batches from a dataset, specified with offset and input size
'''


def update_shared_dataset(index_of_shared, index_of_which_dataset, input_slice, output_slice, transform=True):
    # print("in process id {}".format(os.getpid()))
    shared_dataset = shared_datasets[index_of_shared]
    original_dataset = datasets[index_of_which_dataset]
    # print("original_dataset: ", [
    #     (key,
    #      type(original_dataset[key]),
    #      original_dataset[key])
    #     for key in original_dataset.keys()])
    dataset_numpy = dict()
    # load inputs
    # print("original_dataset['data']: ", original_dataset['data'].shape, original_dataset['data'].dtype)
    data_slice = np.array(original_dataset['data'][input_slice], dtype=np.float32) / (2. ** 8)
    if transform:
        if 'transform' in original_dataset:
            # print('Pre:',(data_slice.min(),data_slice.mean(),data_slice.max()))
            lo, hi = original_dataset['transform']['scale']
            data_slice = 0.5 + (data_slice-0.5)*np.random.uniform(low=lo,high=hi)
            lo, hi = original_dataset['transform']['shift']
            data_slice = data_slice + np.random.uniform(low=lo,high=hi)
        else:
            print("WARNING: source data doesn't have 'transform' attribute.")
    dataset_numpy['data'] = data_slice
    # load outputs if desired
    if output_slice is not None:
        dataset_numpy['components'] = np.array(original_dataset['components'][output_slice])
        if 'label' in original_dataset:
            dataset_numpy['label'] = np.array(original_dataset['label'][output_slice])
        else:
            # compute affinities from components
            dataset_numpy['label'] = pygt.malis.seg_to_affgraph(dataset_numpy['components'], original_dataset['nhood'])
            warnings.warn("Computing affinity labels because 'label' wasn't provided in data source.", UserWarning)
        if 'mask' in original_dataset:
            dataset_numpy['mask'] = np.array(original_dataset['mask'][output_slice], dtype=np.uint8)
        else:
            # assume no masking
            dataset_numpy['mask'] = np.ones_like(dataset_numpy['components'], dtype=np.uint8)
            warnings.warn("No mask provided. Setting to 1 everywhere.", UserWarning)
    for key in shared_dataset:
        source_array = dataset_numpy[key].astype(dtypes[key])
        target_mp_array = shared_dataset[key].get_obj()
        if pygt.DEBUG:
            print("dataset_numpy['{0}']: dtype {1} and shape {2}".format(key, source_array.dtype, source_array.shape))<commit_msg>Fix slicing for labels when loading from HDF5<commit_after>from __future__ import print_function

import multiprocessing
import os
import time
from operator import mul
from os.path import join
import warnings

import h5py
import numpy as np

import PyGreentea as pygt

''' where this will be used:
* train()
  * getting unordered batches of a dataset, specified with offset, input size and output size
  * behavior should be like a queue. at the end of each training iteration, train() specifies a replacement
* process()
  * getting batches from a dataset, specified with offset and input size
'''


def update_shared_dataset(index_of_shared, index_of_which_dataset, input_slice, output_slice, transform=True):
    # print("in process id {}".format(os.getpid()))
    shared_dataset = shared_datasets[index_of_shared]
    original_dataset = datasets[index_of_which_dataset]
    # print("original_dataset: ", [
    #     (key,
    #      type(original_dataset[key]),
    #      original_dataset[key])
    #     for key in original_dataset.keys()])
    dataset_numpy = dict()
    # load inputs
    # print("original_dataset['data']: ", original_dataset['data'].shape, original_dataset['data'].dtype)
    data_slice = np.array(original_dataset['data'][input_slice], dtype=np.float32) / (2. ** 8)
    if transform:
        if 'transform' in original_dataset:
            # print('Pre:',(data_slice.min(),data_slice.mean(),data_slice.max()))
            lo, hi = original_dataset['transform']['scale']
            data_slice = 0.5 + (data_slice-0.5)*np.random.uniform(low=lo,high=hi)
            lo, hi = original_dataset['transform']['shift']
            data_slice = data_slice + np.random.uniform(low=lo,high=hi)
        else:
            print("WARNING: source data doesn't have 'transform' attribute.")
    dataset_numpy['data'] = data_slice
    # load outputs if desired
    if output_slice is not None:
        dataset_numpy['components'] = np.array(original_dataset['components'][output_slice])
        if 'label' in original_dataset:
            label_slice = (slice(None),) + output_slice
            dataset_numpy['label'] = np.array(original_dataset['label'][label_slice])
        else:
            # compute affinities from components
            dataset_numpy['label'] = pygt.malis.seg_to_affgraph(dataset_numpy['components'], original_dataset['nhood'])
            warnings.warn("Computing affinity labels because 'label' wasn't provided in data source.", UserWarning)
        if 'mask' in original_dataset:
            dataset_numpy['mask'] = np.array(original_dataset['mask'][output_slice], dtype=np.uint8)
        else:
            # assume no masking
            dataset_numpy['mask'] = np.ones_like(dataset_numpy['components'], dtype=np.uint8)
            warnings.warn("No mask provided. Setting to 1 everywhere.", UserWarning)
    for key in shared_dataset:
        source_array = dataset_numpy[key].astype(dtypes[key])
        target_mp_array = shared_dataset[key].get_obj()
        if pygt.DEBUG:
            print("dataset_numpy['{0}']: dtype {1} and shape {2}".format(key, source_array.dtype, source_array.shape))
{'fraction_non_alphanumeric': 0.116000647144475, 'fraction_numerical': 0.005015369681281346, 'mean_word_length': 3.609992542878449, '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': '5952460', 'n_tokens_mistral': 1781, 'n_tokens_neox': 1768, 'n_words': 445}