from inteprelens.utils.ds import CHGDataset


def test_chg_dataset_from_texts_builds_masks(llama_tokenizer):
    prompt = "What is 2+2?"
    target = "4"
    dataset = CHGDataset.from_texts(
        texts=[prompt],
        targets=[target],
        tokenizer=llama_tokenizer,
        device="cpu",
    )

    input_ids = dataset["input_ids"]
    loss_masks = dataset["loss_masks"]

    target_with_space = target if target.startswith(" ") else f" {target}"
    expected_target_tokens = len(
        llama_tokenizer(target_with_space, add_special_tokens=False)["input_ids"]
    )

    assert input_ids.shape == loss_masks.shape
    assert input_ids.shape[0] == 1
    assert loss_masks.sum().item() == expected_target_tokens
    assert loss_masks[0, -1].item() is False
    assert any("<|TARGET|>" in tok for tok in llama_tokenizer.get_vocab())


def test_chg_dataset_from_texts_respects_max_length(llama_tokenizer):
    dataset = CHGDataset.from_texts(
        texts=["A B C D"],
        targets=["E"],
        tokenizer=llama_tokenizer,
        max_length=3,
        device="cpu",
    )

    assert dataset["input_ids"].shape[1] == 3
    assert dataset["loss_masks"].shape[1] == 3
