Function process_fn

SwissArmyTransformer/tests/test_jsonlds.py:15–32

(src, seq_len)

'--iterable-dataset', '--split','1', '--num-workers', '0'])

def process_fn(src, seq_len):
    tokenizer = AutoTokenizer.from_pretrained("/mnt/shared/official_pretrains/hf_home/chatglm2-6b", trust_remote_code=True, local_files_only=True)

    buffered_token_ids = None
    eos_id = torch.tensor([tokenizer.eos_token_id])

    for x in src:
        type_id = 0 if 'code' in x else 1
        txt = x['code'] if 'code' in x else x['content']
        tokenized_ids = torch.tensor(tokenizer.encode(txt))
        if buffered_token_ids is None:
            buffered_token_ids = tokenized_ids
        else:
            buffered_token_ids = torch.cat((buffered_token_ids, eos_id, tokenized_ids), dim=0)
        # yield per seq_len
        while buffered_token_ids.shape[0] >= seq_len:
            yield {'txt': buffered_token_ids[:seq_len], 'type_id': type_id}
            buffered_token_ids = buffered_token_ids[seq_len:]
        
def create_func(path, args):
    return JsonlIterableDataset(path, partial(process_fn, seq_len=512), seed=1, shuffle_buffer=1)