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

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()