test_confidence.py

from __future__ import annotations

from typing import ClassVar

import torch

from thinker_r.confidence import find_subsequence, score_generated_choice


class TinyTokenizer:
    tokens: ClassVar[dict[str, list[int]]] = {
        "<answer>": [1],
        "<answer": [1],
        "A": [2],
        "B": [3],
        "C": [4],
        "D": [5],
        " A": [6],
        " B": [7],
        " C": [8],
        " D": [9],
        ">A": [12],
        ">B": [13],
        ">C": [14],
        ">D": [15],
    }
    decoded: ClassVar[dict[int, str]] = {
        1: "<answer>",
        2: "A",
        3: "B",
        4: "C",
        5: "D",
        10: "\n",
        14: ">C",
    }

    def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
        del add_special_tokens
        return self.tokens.get(text, [99, 100])

    def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str:
        del skip_special_tokens
        return "".join(self.decoded.get(token, "?") for token in ids)


def test_find_subsequence() -> None:
    assert find_subsequence([4, 1, 2, 3], [1, 2]) == 1
    assert find_subsequence([1, 2], [3]) is None


def test_score_generated_choice_skips_whitespace_and_scores_labels() -> None:
    tokenizer = TinyTokenizer()
    generated = [20, 1, 10, 4]
    scores = [torch.zeros(1, 20) for _ in generated]
    scores[3][0, 4] = 8.0
    result = score_generated_choice(tokenizer, generated, scores)
    assert result is not None
    assert result["position"] == 3
    assert result["prediction"] == "C"
    assert result["confidence"] > 0.99
    assert result["margin"] > 0.98


def test_score_generated_choice_requires_answer_tag() -> None:
    tokenizer = TinyTokenizer()
    assert score_generated_choice(tokenizer, [4], [torch.zeros(1, 20)]) is None


def test_score_generated_choice_handles_tag_letter_token_merge() -> None:
    tokenizer = TinyTokenizer()
    generated = [20, 1, 14]
    scores = [torch.zeros(1, 20) for _ in generated]
    scores[2][0, 14] = 8.0
    result = score_generated_choice(tokenizer, generated, scores)
    assert result is not None
    assert result["position"] == 2
    assert result["prediction"] == "C"
    assert result["confidence"] > 0.99