tests/test_chat_format.py

"""Unit tests for qwip_atlas.chat_format.

These pin the two correctness-critical behaviours without needing a real model:
  1. When chat templating is OFF, prompts pass through untouched and special
     tokens are added by the tokenizer as usual (one BOS).
  2. When chat templating is ON, each prompt is wrapped via apply_chat_template
     (whose output already carries BOS), so encode_prompts must tokenize with
     add_special_tokens=False -> exactly one BOS, never two.
"""
from qwip_atlas.chat_format import (
    chat_template_enabled,
    encode_prompts,
    format_prompt_text,
    set_chat_template,
)


class FakeTokenizer:
    """Minimal stand-in. apply_chat_template prepends a BOS marker exactly like
    a Llama-3 template does; the call records the add_special_tokens it was given."""

    BOS = "<BOS>"

    def __init__(self):
        self.last_call = None

    def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=False):
        assert tokenize is False
        content = messages[-1]["content"]
        tail = "<ASSISTANT>" if add_generation_prompt else ""
        return f"{self.BOS}<USER>{content}{tail}"

    def __call__(self, prompts, add_special_tokens=True, **kwargs):
        # Emulate a tokenizer: prepend BOS iff add_special_tokens, and only if the
        # text doesn't already start with one. A correct caller never doubles it.
        texts = [prompts] if isinstance(prompts, str) else list(prompts)
        ids = []
        for t in texts:
            n_bos = 1 if add_special_tokens else 0
            n_bos += t.count(self.BOS)
            ids.append({"text": t, "n_bos": n_bos})
        self.last_call = {"texts": texts, "add_special_tokens": add_special_tokens}
        return {"input_ids": ids}


def test_flag_resolution_default_off():
    tok = FakeTokenizer()
    assert chat_template_enabled(tok) is False


def test_set_and_read_flag():
    tok = FakeTokenizer()
    set_chat_template(tok, True)
    assert chat_template_enabled(tok) is True
    assert chat_template_enabled(tok, override=False) is False  # explicit wins


def test_off_passes_through_with_one_bos():
    tok = FakeTokenizer()
    out = encode_prompts(tok, ["hello", "world"], return_tensors="pt")
    assert tok.last_call["add_special_tokens"] is True
    assert all("<USER>" not in r["text"] for r in out["input_ids"])  # not wrapped
    assert all(r["n_bos"] == 1 for r in out["input_ids"])  # exactly one BOS


def test_on_wraps_and_avoids_double_bos():
    tok = FakeTokenizer()
    set_chat_template(tok, True)
    out = encode_prompts(tok, ["hello", "world"], return_tensors="pt")
    # add_special_tokens forced off because the template already emits BOS
    assert tok.last_call["add_special_tokens"] is False
    for r in out["input_ids"]:
        assert "<USER>" in r["text"] and r["text"].endswith("<ASSISTANT>")
        assert r["n_bos"] == 1  # the critical assertion: single BOS, not two


def test_explicit_add_special_tokens_respected_when_on():
    tok = FakeTokenizer()
    set_chat_template(tok, True)
    # A caller that explicitly asks for special tokens is honoured (would double).
    encode_prompts(tok, ["x"], add_special_tokens=True)
    assert tok.last_call["add_special_tokens"] is True


def test_format_prompt_text_single():
    tok = FakeTokenizer()
    assert format_prompt_text(tok, "hi") == "hi"  # off
    set_chat_template(tok, True)
    assert format_prompt_text(tok, "hi") == "<BOS><USER>hi<ASSISTANT>"