import torch

from antipasto.eval import append_choice_newline


class NewlineTokenizer:
    def __call__(self, text, add_special_tokens):
        assert text == "\n"
        assert not add_special_tokens
        return {"input_ids": [107]}


def test_append_choice_newline_appends_before_scoring():
    input_ids = torch.tensor([[1, 2], [3, 4]])
    attention_mask = torch.tensor([[1, 1], [1, 0]])

    input_ids, attention_mask = append_choice_newline(
        input_ids, NewlineTokenizer(), attention_mask
    )

    assert input_ids.tolist() == [[1, 2, 107], [3, 4, 107]]
    assert attention_mask.tolist() == [[1, 1, 1], [1, 0, 1]]