<s> import numpy <EOL> import itertools <EOL> import operator <EOL> from tm_dataset import PytablesBitextIterator <EOL> class HomogenousData ( PytablesBitextIterator ) : <EOL> def __init__ ( self , * args , ** kwargs ) : <EOL> PytablesBitextIterator . __init__ ( self , * args , ** kwargs ) <EOL> self . batch_iter = None <EOL> def get_homogenous_batch_iter ( self ) : <EOL> end_of_iter = False <EOL> while True : <EOL> k_batches = <NUM_LIT:10> <EOL> batch_size = self . batch_size <EOL> x = [ ] <EOL> y = [ ] <EOL> for k in xrange ( k_batches ) : <EOL> try : <EOL> dx , dy = PytablesBitextIterator . next ( self ) <EOL> except StopIteration : <EOL> end_of_iter = True <EOL> break <EOL> if dx == None or dy == None : <EOL> break <EOL> x += dx <EOL> y += dy <EOL> if len ( x ) <= <NUM_LIT:0> or len ( y ) <= <NUM_LIT:0> : <EOL> raise StopIteration <EOL> lens = numpy . asarray ( [ map ( len , x ) , map ( len , y ) ] ) <EOL> order = numpy . argsort ( lens . max ( axis = <NUM_LIT:0> ) ) if k_batches > <NUM_LIT:1> else numpy . arange ( len ( x ) ) <EOL> for k in range ( k_batches ) : <EOL> if k * batch_size > len ( order ) : <EOL> break <EOL> indices = order [ k * batch_size : ( k + <NUM_LIT:1> ) * batch_size ] <EOL> yield [ [ x [ ii ] for ii in indices ] , [ y [ ii ] for ii in indices ] ] <EOL> if end_of_iter : <EOL> raise StopIteration <EOL> def next ( self , peek = False ) : <EOL> if not self . batch_iter : <EOL> self . batch_iter = self . get_homogenous_batch_iter ( ) <EOL> if not self . batch_iter : <EOL> raise StopIteration <EOL> try : <EOL> batch = next ( self . batch_iter ) <EOL> except StopIteration : <EOL> self . batch_iter = None <EOL> raise StopIteration <EOL> return batch [ <NUM_LIT:0> ] , batch [ <NUM_LIT:1> ] </s>
