INNER CODE UNIT · Python
pack_examples
CogStack/OpenGPT · opengpt/dataset_utils.py:200
def pack_examples(examples, block_size, packing_type='partial'):
r''' Used with a prepared HF dataset, will pack/group examples. Use with care, can mess up many things
if the input is not formated properly (requires the <|eod|> token).
packing_type: partial/full/no
'''
# Concatenate all texts.
if packing_type == 'partial':
result = {k:[] for k in examples.keys()}
_key = list(examples.keys())[0] # Take whichever key
new_example = {k:[] for k in examples.keys()}
for ind in range(len(examples[_key])):
# Trim long sequences to block_size, this is required for partial packing
example = {k:v[ind][0:block_size] for k,v in examples.items()}
if len(new_example[_key]) + len(example[_key]) > block_size:
result = {k:result[k] + [v] for k,v in new_example.items()}
new_example = example