Function test_eval_token_interval

tests/trainer/test_trainer_eval.py:231–300

Tests that the trainer evaluates the model at the correct intervals when using token-based intervals.

(
    tiny_bert_tokenizer,
    eval_interval: str,
    batch_size: int,
    sequence_length: int,
    tmp_path: pathlib.Path,
)

@pytest.mark.parametrize('batch_size', [1, 4, 5])
@pytest.mark.parametrize('sequence_length', [1, 16])
def test_eval_token_interval(
    tiny_bert_tokenizer,
    eval_interval: str,
    batch_size: int,
    sequence_length: int,
    tmp_path: pathlib.Path,
):
    """Tests that the trainer evaluates the model at the correct intervals when using token-based intervals."""
    tokens_per_batch = batch_size * sequence_length
    max_duration_time = Time.from_timestring('5ba')
    eval_interval_time = Time.from_timestring(eval_interval)
    max_duration_tokens = max_duration_time.value * tokens_per_batch

    # calculate the expected number of evals
    last_token_iter = 0
    next_multiple = eval_interval_time.value
    expected_num_evals = 0
    last_multiple_added = -1
    for token_iter in range(0, max_duration_tokens + tokens_per_batch, tokens_per_batch):
        if last_token_iter < next_multiple <= token_iter:
            last_multiple_added = next_multiple
            expected_num_evals += 1
        last_token_iter = token_iter
        while next_multiple <= last_token_iter:
            next_multiple += eval_interval_time.value

    if last_multiple_added + tokens_per_batch <= max_duration_tokens:
        expected_num_evals += 1

    num_eval_batches = 2
    expected_batch_evals = expected_num_evals * num_eval_batches

    transformers = pytest.importorskip('transformers')
    model = SimpleTransformerMaskedLM(vocab_size=tiny_bert_tokenizer.vocab_size)
    pretraining_train_dataset = RandomTextLMDataset(
        size=100,
        vocab_size=tiny_bert_tokenizer.vocab_size,
        sequence_length=sequence_length,
        use_keys=True,
    )

    collator = transformers.DataCollatorForLanguageModeling(tokenizer=tiny_bert_tokenizer, mlm_probability=0.15)
    dataloader = DataLoader(
        pretraining_train_dataset,
        batch_size=batch_size,
        sampler=dist.get_sampler(pretraining_train_dataset),
        collate_fn=collator,
    )
    eval_dataloader = DataLoader(
        pretraining_train_dataset,
        batch_size=batch_size,
        sampler=dist.get_sampler(pretraining_train_dataset),
        collate_fn=collator,
    )

    event_counter_callback = EventCounterCallback()
    trainer = Trainer(
        model=model,