tests/test_bpe_expert.py

"""Tests for the global BPE tokenizer integration in ENSEMBLE."""
from __future__ import annotations

import os

import pytest

from ensemble import Expert, Brain
from palimseste.bpe import BPETokenizer


CORPUS = (
    "the mitochondria is the powerhouse of the cell. "
    "what is the capital of france. the capital of france is paris. "
    "what is pi. pi is approximately three point one four. "
    "two plus two equals four. three times three equals nine. "
) * 6

MATH_QA = [
    ("what is pi", "pi is approximately three point one four"),
    ("what is two plus two", "two plus two equals four"),
    ("what is three times three", "three times three equals nine"),
] * 4

GEO_QA = [
    ("what is the capital of france", "the capital of france is paris"),
    ("what is the capital of japan", "the capital of japan is tokyo"),
] * 4


@pytest.fixture
def bpe():
    return Expert.build_bpe(CORPUS, vocab_size=400, D=3000)


class TestBuildBPE:
    def test_returns_bpe_tokenizer(self, bpe):
        assert isinstance(bpe, BPETokenizer)
        assert bpe.vocab_size_actual > 100  # learned some merges

    def test_bpe_is_deterministic(self):
        a = Expert.build_bpe(CORPUS, vocab_size=400, D=3000, seed=0)
        b = Expert.build_bpe(CORPUS, vocab_size=400, D=3000, seed=0)
        assert a.id2token == b.id2token
        assert a.merges == b.merges


class TestBPEExpert:
    def test_from_text_with_bpe(self, bpe):
        e = Expert.from_text(CORPUS, domain="bio", D=3000, tokenizer=bpe)
        assert e._tokenizer_type == "bpe"
        assert e.n_traces > 0
        assert e.model.tokenizer is bpe

    def test_from_qa_pairs_with_bpe(self, bpe):
        e = Expert.from_qa_pairs(MATH_QA, domain="math", D=3000, tokenizer=bpe)
        assert e._tokenizer_type == "bpe"
        assert e.n_traces > 0

    def test_default_is_char(self):
        e = Expert.from_text("hello world. " * 20, D=2000)
        assert e._tokenizer_type == "char"

    def test_bpe_answer_is_string(self, bpe):
        e = Expert.from_qa_pairs(MATH_QA, domain="math", D=3000, tokenizer=bpe)
        a = e.answer("what is pi")
        assert isinstance(a, str)

    def test_bpe_answer_better_than_char(self, bpe):
        """BPE should produce a more complete answer than char-level on the same data."""
        char_expert = Expert.from_qa_pairs(MATH_QA, domain="m", D=3000)
        bpe_expert = Expert.from_qa_pairs(MATH_QA, domain="m", D=3000, tokenizer=bpe)
        char_ans = char_expert.answer("what is two plus two")
        bpe_ans = bpe_expert.answer("what is two plus two")
        expected = "two plus two equals four"
        # BPE answer should be at least as long / complete as the char one
        assert len(bpe_ans) >= len(char_ans)
        # BPE should contain the key correct substring
        assert "two plus two" in bpe_ans.lower()


class TestBPESaveLoad:
    def test_round_trip_with_shared_bpe(self, bpe, tmp_path):
        e = Expert.from_qa_pairs(MATH_QA, domain="math", D=3000, tokenizer=bpe)
        result = e.save(tmp_path / "math.exp")
        e2 = Expert.load(result.path, tokenizer=bpe)
        assert e2._tokenizer_type == "bpe"
        assert e2.n_traces == e.n_traces
        assert e2.signature_hv == e.signature_hv
        assert e2.model._self_hv == e.model._self_hv

    def test_bpe_expert_smaller_than_source(self, bpe, tmp_path):
        e = Expert.from_qa_pairs(MATH_QA, domain="math", D=3000, tokenizer=bpe)
        result = e.save(tmp_path / "math.exp")
        assert result.expert_size_bytes < result.source_size_bytes

    def test_load_bpe_without_tokenizer_raises(self, bpe, tmp_path):
        e = Expert.from_qa_pairs(MATH_QA, domain="math", D=3000, tokenizer=bpe)
        result = e.save(tmp_path / "math.exp")
        with pytest.raises(ValueError, match="BPE"):
            Expert.load(result.path)

    def test_no_vocab_json_for_bpe(self, bpe, tmp_path):
        """BPE experts must NOT serialize a per-expert vocab (it's global)."""
        e = Expert.from_qa_pairs(MATH_QA, domain="math", D=3000, tokenizer=bpe)
        result = e.save(tmp_path / "math.exp")
        assert not os.path.exists(os.path.join(result.path, "vocab.json"))


class TestBPEBrain:
    def test_brain_with_bpe_experts(self, bpe):
        math = Expert.from_qa_pairs(MATH_QA, domain="math", D=3000, tokenizer=bpe)
        geo = Expert.from_qa_pairs(GEO_QA, domain="geography", D=3000, tokenizer=bpe)
        brain = Brain()
        brain.add_expert(math)
        brain.add_expert(geo)
        res = brain.query("what is pi", max_new_tokens=40)
        assert res.dominant_expert == "math"
        res2 = brain.query("what is the capital of france", max_new_tokens=40)
        assert res2.dominant_expert == "geography"

    def test_mixed_char_and_bpe_brain(self, bpe):
        """A brain can mix char-level and BPE experts (different tokenizers)."""
        bpe_expert = Expert.from_qa_pairs(MATH_QA, domain="math", D=3000, tokenizer=bpe)
        char_expert = Expert.from_qa_pairs(GEO_QA, domain="geo", D=3000)
        brain = Brain()
        brain.add_expert(bpe_expert)
        brain.add_expert(char_expert)
        # both should be queryable without crashing
        r1 = brain.query("what is pi")
        r2 = brain.query("what is the capital of france")
        assert isinstance(r1.answer, str)
        assert isinstance(r2.answer, str)