import torch

from inteprelens.utils.tokenisation import PromptTokenizer


def test_prompt_tokenizer(llama_tokenizer):
    prompt_tokenizer = PromptTokenizer(llama_tokenizer)
    input_ids, loss_masks = prompt_tokenizer.tokenize_batch(
        prompts=["What is 2+2?"],
        targets=["4"],
    )

    assert prompt_tokenizer.marker.startswith(" ")
    assert prompt_tokenizer.marker.endswith(" ")
    assert prompt_tokenizer.marker in llama_tokenizer.get_vocab()

    expected_marker_id = llama_tokenizer(
        prompt_tokenizer.marker, add_special_tokens=False
    )["input_ids"][0]
    assert prompt_tokenizer.marker_id == expected_marker_id

    assert isinstance(input_ids, torch.Tensor)
    assert isinstance(loss_masks, torch.Tensor)
    assert input_ids.shape == loss_masks.shape
    assert input_ids.ndim == 2
    assert loss_masks.dtype == torch.bool
    assert loss_masks.any().item()
