File size: 6,413 Bytes
296a506
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
"""
Frox AI Morph 1.1 — Tokenizer Tests
Run with: pytest tests/test_tokenizer.py -v

Uses the synthetic-corpus fallback path (no network needed) so this
suite runs offline. Builds a small real BPE tokenizer once per test
session rather than mocking it — the whole point is to catch real
tokenizer-training regressions (e.g. a special token silently missing
an ID, or the chat template drifting out of sync with the parser in
morph_engine.py).
"""
from __future__ import annotations

import sys
from pathlib import Path

import pytest

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from tokenizer.morph_tokenizer import (
    ALL_SPECIAL_TOKENS, MORPH_SPECIAL_TOKENS,
    apply_chat_template, build_morph_tokenizer,
    format_dpo_pair, format_thinking, format_tool_call, format_tool_result,
    _make_synthetic_corpus, _train_bpe, _wrap_tokenizer,
)


@pytest.fixture(scope="session")
def tokenizer(tmp_path_factory):
    """Build one small tokenizer for the whole test session (BPE training isn't free)."""
    corpus_dir = tmp_path_factory.mktemp("corpus")
    corpus_path = _make_synthetic_corpus(corpus_dir)
    tok_model = _train_bpe([corpus_path], vocab_size=800)
    return _wrap_tokenizer(tok_model)


class TestSpecialTokens:
    def test_all_special_tokens_present(self, tokenizer):
        vocab = tokenizer.get_vocab()
        missing = [t for t in ALL_SPECIAL_TOKENS if t not in vocab]
        assert not missing, f"Missing special tokens: {missing}"

    def test_morph_specific_tokens_have_unique_ids(self, tokenizer):
        vocab = tokenizer.get_vocab()
        ids = [vocab[t] for t in MORPH_SPECIAL_TOKENS]
        assert len(ids) == len(set(ids)), "Two Morph special tokens collapsed to the same ID"

    def test_core_attributes_set(self, tokenizer):
        assert tokenizer.pad_token == "<|pad|>"
        assert tokenizer.bos_token == "<|begin_of_text|>"
        assert tokenizer.eos_token == "<|end_of_text|>"
        assert tokenizer.pad_token_id is not None
        assert tokenizer.eos_token_id is not None


class TestEncodeDecodeRoundTrip:
    def test_simple_text_round_trips(self, tokenizer):
        text = "Neural networks compute features through multiple layers."
        ids = tokenizer.encode(text, add_special_tokens=False)
        decoded = tokenizer.decode(ids, skip_special_tokens=True)
        # BPE round trip isn't always byte-exact for punctuation/whitespace edge
        # cases, but the alphanumeric content must survive intact.
        assert "Neural" in decoded or "neural" in decoded.lower()
        assert len(ids) > 0

    def test_empty_string_does_not_crash(self, tokenizer):
        ids = tokenizer.encode("", add_special_tokens=False)
        assert isinstance(ids, list)


class TestChatTemplate:
    def test_produces_expected_structure(self, tokenizer):
        messages = [{"role": "user", "content": "Hello!"}]
        prompt = apply_chat_template(messages, tokenizer, add_generation_prompt=True)

        assert "<|begin_of_text|>" in prompt
        assert "<|start_header_id|>user<|end_header_id|>" in prompt
        assert "Hello!" in prompt
        assert prompt.rstrip().endswith(
            "<|start_header_id|>assistant<|end_header_id|>"
        ), "Generation prompt (assistant header) must be the last thing in the prompt"

    def test_system_prompt_injected_first(self, tokenizer):
        messages = [{"role": "user", "content": "Hi"}]
        prompt = apply_chat_template(
            messages, tokenizer, add_generation_prompt=True,
            system_prompt="You are a helpful assistant.",
        )
        system_pos = prompt.find("system")
        user_pos = prompt.find("user")
        assert 0 <= system_pos < user_pos, "System turn must come before the user turn"

    def test_existing_system_message_not_duplicated(self, tokenizer):
        messages = [
            {"role": "system", "content": "Custom system prompt"},
            {"role": "user", "content": "Hi"},
        ]
        prompt = apply_chat_template(
            messages, tokenizer, add_generation_prompt=True,
            system_prompt="This should NOT appear since messages[0] is already system",
        )
        assert prompt.count("<|start_header_id|>system<|end_header_id|>") == 1
        assert "Custom system prompt" in prompt
        assert "This should NOT appear" not in prompt

    def test_multi_turn_preserves_order(self, tokenizer):
        messages = [
            {"role": "user", "content": "First question"},
            {"role": "assistant", "content": "First answer"},
            {"role": "user", "content": "Second question"},
        ]
        prompt = apply_chat_template(messages, tokenizer, add_generation_prompt=True)
        assert prompt.find("First question") < prompt.find("First answer")
        assert prompt.find("First answer") < prompt.find("Second question")


class TestToolFormatting:
    def test_tool_call_format_is_parseable_json(self, tokenizer):
        import re, json
        formatted = format_tool_call("web_search", {"query": "test"})
        match = re.search(r"<\|tool_call\|>(.*?)<\|/tool_call\|>", formatted)
        assert match is not None
        parsed = json.loads(match.group(1))
        assert parsed["name"] == "web_search"
        assert parsed["args"]["query"] == "test"

    def test_tool_result_format_is_parseable_json(self, tokenizer):
        import re, json
        formatted = format_tool_result("web_search", {"results": [1, 2, 3]})
        match = re.search(r"<\|tool_result\|>(.*?)<\|/tool_result\|>", formatted)
        assert match is not None
        parsed = json.loads(match.group(1))
        assert parsed["name"] == "web_search"

    def test_thinking_block_wraps_content(self):
        formatted = format_thinking("Let me reason step by step.")
        assert formatted == "<|think|>Let me reason step by step.<|/think|>"


class TestDPOFormatting:
    def test_prompt_chosen_rejected_share_prefix(self, tokenizer):
        messages = [{"role": "user", "content": "What's 2+2?"}]
        pair = format_dpo_pair(
            messages, chosen="4", rejected="I don't know", tokenizer=tokenizer,
        )
        assert pair["chosen"].startswith(pair["prompt"])
        assert pair["rejected"].startswith(pair["prompt"])
        assert pair["chosen"] != pair["rejected"]
        assert "4" in pair["chosen"]
        assert "I don't know" in pair["rejected"]