#!/usr/local/bin/pypy

import sys

from conll import open2, dump_corpus

__author__ = 'husnusensoy'

def filename(basefilename, idx, count):

return "%s.%d.%d" % (basefilename[:-6], idx, count)

if __name__ == "__main__":

import argparse

parser = argparse.ArgumentParser(description='CoNLL file to be splited')

parser.add_argument('file', help='CoNLL file to be split')

parser.add_argument('--tail', action='store_true', default=False,

help='tail/head -<n> file for the first piece. Default: head')

parser.add_argument('--n', default=1000, help="Number of rows. Default: 1000")

parser.add_argument('--first', help='File name for the first piece. Default: <file>.1.(count)')

parser.add_argument('--second', help='File name for the second piece. Default: <file>.2.(count)')

args = parser.parse_args()

assert (args.first != args.second or (not args.first and not args.second))

with open2(args.file) as fp:

corpus = [conll for conll in fp]

if args.tail:

corpus1 = corpus[:-args.n]

corpus2 = corpus[-args.n:]

else:

corpus1 = corpus[:args.n]

corpus2 = corpus[args.n:]

file1 = args.first if args.first else filename(args.file, 1, len(corpus1))

file2 = args.second if args.second else filename(args.file, 2, len(corpus2))

print >> sys.stderr, "%s (%d) is splitted as \n\t%s (%d)\n\t%s (%d)" % (

args.file, len(corpus), file1, len(corpus1), file2, len(corpus2))

dump_corpus(file1, corpus1)

dump_corpus(file2, corpus2)