Обновление BPE: добавлена документация, тесты и улучшен пример использования

tests/test_bpe.py

@@ -1,54 +1,24 @@
 import pytest
-from simple_llm.tokenizer.simple_bpe import SimpleBPE
-from simple_llm.tokenizer.optimize_bpe import OptimizeBPE
+from simple_llm.tokenizer.bpe import BPE
 
-class TestBPE:
-    @pytest.fixture(params=[SimpleBPE, OptimizeBPE])
-    def bpe_class(self, request):
-        return request.param
+def test_basic_bpe():
+    """Базовый тест работы BPE"""
+    tokenizer = BPE(vocab_size=10)
+    text = "мама мыла раму"
     
-    def test_initialization(self, bpe_class):
-        """Тест инициализации BPE-токенизатора"""
-        bpe = bpe_class(vocab_size=100)
-        assert bpe.vocab_size == 100
-        assert bpe.vocab == []
-        assert bpe.token2id == {}
-        assert bpe.id2token == {}
+    # Обучение
+    tokenizer.fit(text)
     
-    def test_fit_simple_text(self, bpe_class):
-        """Тест обучения на простом тексте"""
-        text = "мама мыла раму"
-        bpe = bpe_class(vocab_size=20)
-        bpe.fit(text)
-        
-        # Проверки словаря
-        assert isinstance(bpe.vocab, list)
-        assert len(bpe.vocab) > 0
-        assert len(bpe.vocab) <= 20
-        assert all(isinstance(token, str) for token in bpe.vocab)
-        
-        # Проверка словарей
-        assert len(bpe.vocab) == len(bpe.token2id)
-        assert len(bpe.vocab) == len(bpe.id2token)
-        
-        # Проверка соответствия токенов и ID
-        for token in bpe.vocab:
-            assert bpe.token2id[token] == bpe.vocab.index(token)
-            assert bpe.id2token[bpe.token2id[token]] == token
-
-    @pytest.mark.parametrize("text,expected_min_size", [
-        ("", 0),
-        ("а", 1),
-        ("ааааа", 3)  # Минимум 3 токена
-    ])
-    def test_edge_cases(self, bpe_class, text, expected_min_size):
-        """Тест граничных случаев"""
-        bpe = bpe_class(vocab_size=10)
-        bpe.fit(text)
-        assert len(bpe.vocab) >= expected_min_size
-
-    def test_duplicate_protection(self, bpe_class):
-        """Тест защиты от дубликатов токенов"""
-        bpe = bpe_class(vocab_size=50)
-        bpe.fit("аааааааааа" * 100)  # Много повторений
-        assert len(bpe.vocab) == len(set(bpe.vocab))
+    # Проверка размера словаря
+    assert len(tokenizer.vocab) == 10
+    
+    # Кодирование/декодирование
+    encoded = tokenizer.encode(text)
+    decoded = tokenizer.decode(encoded)
+    
+    assert decoded == text
+    assert len(encoded) > 0
+    
+    # Проверка неизвестных символов
+    unknown_encoded = tokenizer.encode("мама мыла окно")
+    assert -1 in unknown_encoded  # Специальный токен для неизвестных символов

tests/test_bpe_detailed.py

@@ -0,0 +1,74 @@
+import os
+import tempfile
+import pytest
+from simple_llm.tokenizer.bpe import BPE
+
+class TestBPE:
+    @pytest.fixture
+    def sample_text(self):
+        return "ааабббвввггг аааббб дддд ееее жжжж"
+
+    @pytest.fixture
+    def bpe(self):
+        return BPE(vocab_size=20)
+
+    def test_fit(self, bpe, sample_text):
+        """Тест обучения токенизатора"""
+        bpe.fit(sample_text)
+        assert len(bpe.vocab) == bpe.vocab_size
+        assert len(bpe.token2id) == bpe.vocab_size
+        assert len(bpe.id2token) == bpe.vocab_size
+
+    def test_encode_decode(self, bpe, sample_text):
+        """Тест кодирования и декодирования"""
+        bpe.fit(sample_text)
+        encoded = bpe.encode(sample_text)
+        decoded = bpe.decode(encoded)
+        assert decoded == sample_text
+
+    def test_encode_unknown_chars(self, bpe, sample_text):
+        """Тест с неизвестными символами"""
+        bpe.fit(sample_text)
+        test_text = "ааббцц"  # 'цц' нет в обучающем тексте
+        encoded = bpe.encode(test_text)
+        assert -1 in encoded  # Должен содержать специальный токен для неизвестных символов
+        decoded = bpe.decode(encoded)
+        assert "цц" in decoded
+
+    def test_save_load(self, bpe, sample_text):
+        """Тест сохранения и загрузки"""
+        bpe.fit(sample_text)
+        
+        with tempfile.NamedTemporaryFile(delete=False) as tmp:
+            try:
+                bpe.save(tmp.name)
+                loaded = BPE.load(tmp.name)
+                
+                assert loaded.vocab_size == bpe.vocab_size
+                assert loaded.vocab == bpe.vocab
+                assert loaded.token2id == bpe.token2id
+                assert loaded.id2token == bpe.id2token
+                
+                # Проверяем работоспособность после загрузки
+                encoded = loaded.encode(sample_text)
+                decoded = loaded.decode(encoded)
+                assert decoded == sample_text
+            finally:
+                os.unlink(tmp.name)
+
+    def test_pair_merging(self, bpe, sample_text):
+        """Тест правильности объединения пар"""
+        bpe.fit(sample_text)
+        
+        # Проверяем, что самые частые пары были объединены
+        assert 'аа' in bpe.vocab or 'ааа' in bpe.vocab
+        assert 'бб' in bpe.vocab or 'ббб' in bpe.vocab
+
+    def test_vocab_size(self):
+        """Тест обработки слишком маленького vocab_size"""
+        small_bpe = BPE(vocab_size=5)
+        with pytest.raises(ValueError):
+            small_bpe.fit("абвгд")  # Слишком мало для начальных символов
+
+if __name__ == "__main__":
+    pytest.main()