INNER CODE UNIT · Python

take

Guitaricet/relora · pretokenize.py:60

        def take(ds, n):
            return Dataset.from_generator(lambda: (yield from ds.take(n)))
        dataset_dict = {k: take(v, args.take) for k, v in dataset.items()}
        dataset = DatasetDict(dataset_dict)

    logger.info("Tokenizing and chunking the dataset")
    _time = time.time()
    dataset = tokenize_and_chunk(
        tokenizer=tokenizer,
        dataset=dataset,
        text_field=args.text_field,
        sequence_length=args.sequence_length,
        num_cpu=args.num_cpu,
    )
    _hours = (time.time() - _time) / 3600
    logger.info(f"Tokenization and chunking took {_hours:.2f} hours")

    dataset.save_to_disk(save_path)

View source record →