File size: 5,375 Bytes
22d1ad7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | """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)
|