Function create_fixed_short_dataset

efficiency/multiprocess_bench.py:16–22

(tokenizer, num_samples=8192)

from multiprocessing import Process, Queue

def create_fixed_short_dataset(tokenizer, num_samples=8192):
    tokens = torch.randint(100, 16000, (num_samples, 512))
    mask = torch.ones(num_samples, 512)
    return {
        'input_ids': tokens.long(),
        'attention_mask': mask.float()
    }

def create_fixed_long_dataset(tokenizer, num_samples=8192):
    tokens = torch.randint(100, 16000, (num_samples, 8192))