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

example/example_bpe.py

@@ -1,77 +1,84 @@
+from simple_llm.tokenizer.bpe import BPE
 from simple_llm.tokenizer.simple_bpe import SimpleBPE
 from simple_llm.tokenizer.optimize_bpe import OptimizeBPE
 import time
 
-def tokenize_manually(text, vocab):
-    """Простая ручная токенизация по словарю"""
-    tokens = []
-    i = 0
-    n = len(text)
-    while i < n:
-        found = False
-        # Ищем самый длинный возможный токен из словаря
-        for l in range(min(4, n-i), 0, -1):  # проверяем токены длиной до 4 символов
-            if text[i:i+l] in vocab:
-                tokens.append(text[i:i+l])
-                i += l
-                found = True
-                break
-        if not found:  # если токен не найден, берем один символ
-            tokens.append(text[i])
-            i += 1
-    return tokens
-
-def run_example(text, vocab_size=30):
-    print("\n=== Тестирование токенизаторов ===")
-    print(f"Исходный текст: '{text}'\n")
+def compare_tokenizers(text, vocab_size=50):
+    """Сравнивает разные реализации BPE"""
+    print(f"\n=== Анализ текста: '{text[:20]}...' ===")
     
-    # Simple BPE
+    # 1. Базовая реализация BPE
+    start = time.time()
+    bpe = BPE(vocab_size=vocab_size)
+    bpe.fit(text)
+    base_time = time.time() - start
+    
+    print("\n[Базовая реализация BPE]")
+    print(f"Время обучения: {base_time:.4f} сек")
+    print(f"Размер словаря: {len(bpe.vocab)}")
+    print("Примеры токенов:", list(bpe.vocab)[:10], "...")
+    
+    # 2. SimpleBPE
     start = time.time()
     simple_bpe = SimpleBPE(vocab_size=vocab_size)
     simple_bpe.fit(text)
     simple_time = time.time() - start
     
-    print("SimpleBPE:")
+    print("\n[SimpleBPE]")
     print(f"Время обучения: {simple_time:.4f} сек")
     print(f"Размер словаря: {len(simple_bpe.vocab)}")
-    print(f"Пример словаря: {simple_bpe.vocab[:5]}...")
     
-    # Демонстрация encode/decode
-    test_phrases = [text, text.split()[0], "неизвестное_слово"]
-    for phrase in test_phrases:
-        encoded = simple_bpe.encode(phrase)
-        decoded = simple_bpe.decode(encoded)
-        print(f"\nФраза: '{phrase}'")
-        print(f"Закодировано: {encoded}")
-        print(f"Декодировано: '{decoded}'")
-        print(f"Совпадение: {phrase == decoded}")
-    
-    # Optimize BPE
+    # 3. OptimizeBPE
     start = time.time()
     opt_bpe = OptimizeBPE(vocab_size=vocab_size)
     opt_bpe.fit(text)
     opt_time = time.time() - start
     
-    print("\nOptimizeBPE:")
+    print("\n[OptimizeBPE]")
     print(f"Время обучения: {opt_time:.4f} сек")
     print(f"Размер словаря: {len(opt_bpe.vocab)}")
-    print(f"Пример словаря: {opt_bpe.vocab[:5]}...")
     
-    # Демонстрация encode/decode
+    # Сравнение производительности
+    if opt_time > 0:
+        print(f"\nОптимизированная версия быстрее SimpleBPE в {simple_time/opt_time:.1f} раз")
+    
+    # Демонстрация работы на примерах
+    test_phrases = [
+        text.split()[0],  # первое слово
+        text[:10],       # часть текста
+        "неизвестное_слово",  # OOV
+        "спецсимволы: 123, !@#"
+    ]
+    
+    print("\n=== Примеры кодирования/декодирования ===")
     for phrase in test_phrases:
+        print(f"\nФраза: '{phrase}'")
+        
+        encoded = bpe.encode(phrase)
+        decoded = bpe.decode(encoded)
+        print(f"BPE: {encoded} -> '{decoded}'")
+        
+        encoded = simple_bpe.encode(phrase)
+        decoded = simple_bpe.decode(encoded)
+        print(f"SimpleBPE: {encoded} -> '{decoded}'")
+        
         encoded = opt_bpe.encode(phrase)
         decoded = opt_bpe.decode(encoded)
-        print(f"\nФраза: '{phrase}'")
-        print(f"Закодировано: {encoded}")
-        print(f"Декодировано: '{decoded}'")
-        print(f"Совпадение: {phrase == decoded}")
+        print(f"OptimizeBPE: {encoded} -> '{decoded}'")
+
+def main():
+    # Тестовые тексты разной сложности
+    texts = [
+        "мама мыла раму, папа пил какао",
+        "коты бегают быстро, собаки лают громко",
+        "искусственный интеллект меняет мир вокруг нас",
+        "BPE (Byte Pair Encoding) - популярный алгоритм токенизации"
+    ]
     
-    if opt_time > 0:
-        print(f"\nОптимизированная версия быстрее в {simple_time/opt_time:.1f} раз")
+    for text in texts:
+        compare_tokenizers(text)
+        
+    print("\n=== Тестирование завершено ===")
 
 if __name__ == "__main__":
-    text1 = "мама мыла раму, папа пил какао"
-    text2 = "коты бегают быстро, собаки лают громко"
-    
-    run_example(text1)
-    run_example(text2)
+    main()