dli-testsplit2

#!/usr/bin/python

import argparse
import hashlib
import imp
import pickle

import dlinputs as dli
from dlinputs import filters, gopen, tarrecords, utils

parser = argparse.ArgumentParser("Split a dataset into training and test sets.")
parser.add_argument("-n", "--nohash", action="store_true")
parser.add_argument("-s", "--split", type=int, default=10)
parser.add_argument("-m", "--mod", type=int, default=0)
parser.add_argument("-M", "--maxcount", type=int, default=1e6)
parser.add_argument("-S", "--maxsize", type=int, default=1e9)
parser.add_argument("-k", "--key", default="__key__")
parser.add_argument("input")
parser.add_argument("training")
parser.add_argument("testing")
args = parser.parse_args()

source = gopen.sharditerator_once(args.input, decode=False)

def hash(x):
    assert x is not None
    if isinstance(x, buffer):
        x = str(buffer)
    if not isinstance(x, str):
        x = pickle.dumps(x, -1)
    h = hashlib.md5()
    h.update(x)
    return int(h.hexdigest()[-7:], 16)

def openout(fname):
    if "%" in fname:
        return tarrecords.ShardWriter(fname, maxcount=args.maxcount, maxsize=args.maxsize, encode=False)
    else:
        return tarrecords.TarWriter(fname, encode=False)

training = openout(args.training)
testing = openout(args.testing)

num_training = 0
num_testing = 0

for i, sample in enumerate(source):
    if i%1000==0:
        print i
    key = sample.get(args.key)
    assert key is not None
    h = i if args.nohash else hash(key)
    match = (h % args.split == args.mod)
    if match:
        testing.write(sample)
        num_testing += 1
    else:
        training.write(sample)
        num_training += 1

testing.close()
training.close()

print "#testing", num_testing, "#training", num_training, "total", i