<commit_before>import os
import sys
import shutil
import numpy
import parsable
import loom.config
import loom.runner
import loom.format
import loom.cFormat
import loom.test.util
from loom.util import parallel_map
from distributions.io.stream import json_load, protobuf_stream_load
parsable = parsable.Parsable()

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DATA = os.path.join(ROOT, 'data')
DATASETS = os.path.join(DATA, 'datasets')
RESULTS = os.path.join(DATA, 'results')
MODEL = os.path.join(DATASETS, '{}/model.pb.gz')
GROUPS = os.path.join(DATASETS, '{}/groups')
ASSIGN = os.path.join(DATASETS, '{}/assign.pbs.gz')
ROWS = os.path.join(DATASETS, '{}/rows.pbs.gz')


def mkdir_p(dirname):
    'like mkdir -p'
    if not os.path.exists(dirname):
        os.makedirs(dirname)


def list_options_and_exit(*required):
    print 'try one of:'
    for name in loom.test.util.list_datasets():
        if all(os.path.exists(patt.format(name)) for patt in required):
            print '  {}'.format(name)
    sys.exit(1)


parsable.command(loom.runner.profilers)


@parsable.command
def load(name=None, debug=False):
    '''
    Import a datasets, list available datasets, or load 'all' datasets.
    '''
    if name is None:
        list_options_and_exit()

    names = [name]

    if names == ['all']:
        names = [
            n
            for n in loom.test.util.list_datasets()
            if not os.path.exists(ROWS.format(n))
        ]

    args = [(n, debug) for n in names]
    parallel_map(_load, args)


def _load((name, debug)):
    print 'loading', name
    data_path = os.path.join(DATASETS, name)
    mkdir_p(data_path)
    model = MODEL.format(name)
    groups = GROUPS.format(name)
    assign = ASSIGN.format(name)
    rows = ROWS.format(name)

    dataset = loom.test.util.get_dataset(name)
    meta = dataset['meta']
    data = dataset['data']
    mask = dataset['mask']
    latent = dataset['latent']
    tardis_conf = dataset['tardis_conf']

    loom.format.import_latent(
        meta_in=meta,
        latent_in=latent,
        tardis_conf_in=tardis_conf,
        model_out=model,
        groups_out=groups,
        assign_out=assign)
    loom.format.import_data(
        meta_in=meta,
        data_in=data,
        mask_in=mask,
        rows_out=rows,
        validate=debug)
    loom.runner.shuffle(rows_in=rows, rows_out=rows)

    meta = json_load(meta)
    object_count = len(meta['object_pos'])
    feature_count = len(meta['feature_pos'])
    print '{}: {} rows x {} cols'.format(name, object_count, feature_count)


@parsable.command
def info(name=None, debug=False):
    '''
    Get information about a dataset, or list available datasets.
    '''
    if name is None:
        list_options_and_exit(ROWS)

    if debug:
        pos = 0
        dumped = 'None'
        sizes = []
        try:
            rows = loom.cFormat.protobuf_stream_load(ROWS.format(name))
            for pos, row in enumerate(rows):
                dumped = row.dump()
                sizes.append(row.ByteSize())
        except:
            print 'error after row {} with data:\n{}'.format(pos, dumped)
            raise<commit_msg>Sketch checkpoint construction in benchmark.load<commit_after>import os
import sys
import shutil
import numpy
import parsable
import loom.config
import loom.runner
import loom.format
import loom.cFormat
import loom.schema_pb2
import loom.test.util
from loom.util import parallel_map
from distributions.io.stream import json_load, protobuf_stream_load
from distributions.fileutil import tempdir
parsable = parsable.Parsable()

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DATA = os.path.join(ROOT, 'data')
DATASETS = os.path.join(DATA, 'datasets')
RESULTS = os.path.join(DATA, 'results')
ROWS = os.path.join(DATASETS, '{}/rows.pbs.gz')
MODEL = os.path.join(DATASETS, '{}/model.pb.gz')
GROUPS = os.path.join(DATASETS, '{}/groups')
ASSIGN = os.path.join(DATASETS, '{}/assign.pbs.gz')
CHECKPOINT = os.path.join(DATASETS, '{}/checkpoint.pb.gz')


def checkpoint_files(path):
    return {
        'model': os.path.join(path, '{}/model.pb.gz'),
        'groups': os.path.join(path, '{}/groups'),
        'assign': os.path.join(path, '{}/assign.pbs.gz'),
        'checkpoint': os.path.join(path, '{}/checkpoint.pb.gz'),
    }


def mkdir_p(dirname):
    'like mkdir -p'
    if not os.path.exists(dirname):
        os.makedirs(dirname)


def list_options_and_exit(*required):
    print 'try one of:'
    for name in loom.test.util.list_datasets():
        if all(os.path.exists(patt.format(name)) for patt in required):
            print '  {}'.format(name)
    sys.exit(1)


parsable.command(loom.runner.profilers)


@parsable.command
def load(name=None, debug=False):
    '''
    Import a datasets, list available datasets, or load 'all' datasets.
    '''
    if name is None:
        list_options_and_exit()

    names = [name]

    if names == ['all']:
        names = [
            n
            for n in loom.test.util.list_datasets()
            if not os.path.exists(ROWS.format(n))
        ]

    args = [(n, debug) for n in names]
    parallel_map(_load, args)


def _load((name, debug)):
    print 'loading', name
    data_path = os.path.join(DATASETS, name)
    mkdir_p(data_path)
    model = MODEL.format(name)
    groups = GROUPS.format(name)
    assign = ASSIGN.format(name)
    rows = ROWS.format(name)

    dataset = loom.test.util.get_dataset(name)
    meta = dataset['meta']
    data = dataset['data']
    mask = dataset['mask']
    latent = dataset['latent']
    tardis_conf = dataset['tardis_conf']

    loom.format.import_latent(
        meta_in=meta,
        latent_in=latent,
        tardis_conf_in=tardis_conf,
        model_out=model,
        groups_out=groups,
        assign_out=assign)
    loom.format.import_data(
        meta_in=meta,
        data_in=data,
        mask_in=mask,
        rows_out=rows,
        validate=debug)
    loom.runner.shuffle(rows_in=rows, rows_out=rows)

    TODO = False
    if TODO:
        with tempdir():
            config_in = os.path.abspath('config.pb.gz')
            loom.config.config_dump({}, config_in)
            model_in = os.path.abspath('model_in.pb.gz')
            groups_in = os.path.abspath('groups_in')
            assign_in = os.path.abspath('assign_in.pbs.gz')
            checkpoint_in = os.path.abspath('checkpoint_in.pb.gz')
            model_out = os.path.abspath('model_out.pb.gz')
            groups_out = os.path.abspath('groups_out')
            assign_out = os.path.abspath('assign_out.pbs.gz')
            checkpoint_out = os.path.abspath('checkpoint_out.pb.gz')
            loom.runner.infer(
                config_in=config_in,
                rows_in=rows,
                model_in=model_in,
                groups_in=groups_in,
                assign_in=assign_in,
                checkpoint_in=checkpoint_in,
                model_out=model_out,
                groups_out=groups_out,
                assign_out=assign_out,
                checkpoint_out=checkpoint_out)
            checkpoint = loom.schema_pb2.Checkpoint()
            with open(checkpoint_out) as f:
                checkpoint.ParseFromString(f.read())
            max_reject_iters =\
                loom.config.DEFAULTS['schedule']['max_reject_iters']
            if checkpoint.schedule.reject_iters <= max_reject_iters:
                raise NotImplementedError('TODO save infiles')

    meta = json_load(meta)
    object_count = len(meta['object_pos'])
    feature_count = len(meta['feature_pos'])
    print '{}: {} rows x {} cols'.format(name, object_count, feature_count)


@parsable.command
def info(name=None, debug=False):
    '''
    Get information about a dataset, or list available datasets.
    '''
    if name is None:
        list_options_and_exit(ROWS)

    if debug:
        pos = 0
        dumped = 'None'
        sizes = []
        try:
            rows = loom.cFormat.protobuf_stream_load(ROWS.format(name))
            for pos, row in enumerate(rows):
                dumped = row.dump()
                sizes.append(row.ByteSize())
        except:
            print 'error after row {} with data:\n{}'.format(pos, dumped)
            raise
{'fraction_non_alphanumeric': 0.12679834186783712, 'fraction_numerical': 0.000731528895391368, 'mean_word_length': 2.585227272727273, 'pattern_counts': {'":': 0, '<': 4, '<?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': '5853723', 'n_tokens_mistral': 2778, 'n_tokens_neox': 2718, 'n_words': 534}