diff --git a/experiments/hf_integration/generate_with_hf_tools.py b/experiments/hf_integration/generate_with_hf_tools.py index 7adae5c..4698316 100644 --- a/experiments/hf_integration/generate_with_hf_tools.py +++ b/experiments/hf_integration/generate_with_hf_tools.py @@ -14,54 +14,50 @@ sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from hf_proxy import HFAdapter, HFTokenizerAdapter, create_hf_pipeline -from shared.configs import ( - TEST_PROMPTS, GENERATION_CONFIG, PATHS -) -from shared.data import ( - print_experiment_info, ensure_directories, ExperimentLogger -) +from shared.configs import TEST_PROMPTS, GENERATION_CONFIG, PATHS +from shared.data import print_experiment_info, ensure_directories, ExperimentLogger def load_hf_model_and_tokenizer() -> tuple: """ Загружает модель и токенизатор в формате HuggingFace. - + Returns: tuple: (hf_model, hf_tokenizer, model_config) """ # Используем упрощенную версию модели model_path = "checkpoints/hf_simple_trained" tokenizer_path = "checkpoints/hf_simple_tokenizer" - + # Проверяем существование файлов if not os.path.exists(model_path): raise FileNotFoundError( f"Модель не найдена: {model_path}\n" f"Сначала обучите модель: uv run python experiments/hf_integration/simple_hf_training.py" ) - + if not os.path.exists(tokenizer_path): - raise FileNotFoundError( - f"Токенизатор не найден: {tokenizer_path}" - ) - + raise FileNotFoundError(f"Токенизатор не найден: {tokenizer_path}") + # Загружаем адаптированный токенизатор print("🔧 Загрузка адаптированного токенизатора...") hf_tokenizer = HFTokenizerAdapter.from_pretrained(tokenizer_path) print(f"✅ Токенизатор загружен (vocab_size={hf_tokenizer.vocab_size})") - + # Загружаем конфигурацию модели import json + config_path = os.path.join(model_path, "config.json") - with open(config_path, 'r', encoding='utf-8') as f: + with open(config_path, "r", encoding="utf-8") as f: model_config = json.load(f) - + # Загружаем модель через HFAdapter с правильной конфигурацией print("🔧 Загрузка адаптированной модели...") model_bin_path = os.path.join(model_path, "pytorch_model.bin") - + # Создаем конфигурацию из сохраненного config.json from hf_proxy import HFAdapterConfig + hf_config = HFAdapterConfig( vocab_size=model_config["vocab_size"], hidden_size=model_config["hidden_size"], @@ -69,26 +65,28 @@ def load_hf_model_and_tokenizer() -> tuple: num_attention_heads=model_config["num_attention_heads"], max_position_embeddings=model_config["max_position_embeddings"], hidden_dropout_prob=model_config.get("hidden_dropout_prob", 0.1), - attention_probs_dropout_prob=model_config.get("attention_probs_dropout_prob", 0.1), + attention_probs_dropout_prob=model_config.get( + "attention_probs_dropout_prob", 0.1 + ), ) - + hf_model = HFAdapter.from_pretrained(model_bin_path, hf_config=hf_config) hf_model.eval() print("✅ Модель загружена") - + return hf_model, hf_tokenizer, model_config def test_hf_pipeline(hf_model, hf_tokenizer): """ Тестирует создание HuggingFace pipeline. - + Args: hf_model: Адаптированная модель hf_tokenizer: Адаптированный токенизатор """ print("\n🧪 Тестирование HuggingFace pipeline...") - + try: # Создаем pipeline pipe = create_hf_pipeline( @@ -97,23 +95,23 @@ def test_hf_pipeline(hf_model, hf_tokenizer): device="cpu", max_length=50, do_sample=True, - temperature=0.7 + temperature=0.7, ) - + print("✅ HuggingFace pipeline создан") - + # Тестируем pipeline test_prompts = TEST_PROMPTS[:3] - + for prompt in test_prompts: print(f"\n🔤 Промпт: '{prompt}'") - + try: result = pipe(prompt, max_new_tokens=20) print(f"🎯 Результат: {result[0]['generated_text']}") except Exception as e: print(f"❌ Ошибка в pipeline: {e}") - + except Exception as e: print(f"❌ Ошибка создания pipeline: {e}") @@ -121,47 +119,49 @@ def test_hf_pipeline(hf_model, hf_tokenizer): def generate_with_hf_model(hf_model, hf_tokenizer, prompt: str, config: dict) -> str: """ Генерирует текст через адаптированную модель HF. - + Args: hf_model: Адаптированная модель hf_tokenizer: Адаптированный токенизатор prompt: Входной текст config: Конфигурация генерации - + Returns: str: Сгенерированный текст """ print(f"🔤 Промпт: '{prompt}'") - print(f"📊 Параметры: max_tokens={config['max_new_tokens']}, " - f"temp={config['temperature']}, sample={config['do_sample']}") - + print( + f"📊 Параметры: max_tokens={config['max_new_tokens']}, " + f"temp={config['temperature']}, sample={config['do_sample']}" + ) + # Кодируем через адаптированный токенизатор inputs = hf_tokenizer(prompt, return_tensors="pt") - + print(f"🎯 Токены промпта: {inputs['input_ids'].tolist()[0]}") print("🔄 Генерация через HF адаптер...") - + # Генерируем через адаптированную модель with torch.no_grad(): generated_ids = hf_model.generate( - input_ids=inputs['input_ids'], + input_ids=inputs["input_ids"], max_new_tokens=config["max_new_tokens"], do_sample=config["do_sample"], temperature=config["temperature"], top_k=config["top_k"], - top_p=config["top_p"] + top_p=config["top_p"], ) - + # Декодируем через адаптированный токенизатор generated_text = hf_tokenizer.decode(generated_ids[0], skip_special_tokens=True) - + return generated_text def test_different_hf_strategies(hf_model, hf_tokenizer, prompt: str): """ Тестирует разные стратегии генерации через HF интерфейс. - + Args: hf_model: Адаптированная модель hf_tokenizer: Адаптированный токенизатор @@ -169,32 +169,38 @@ def test_different_hf_strategies(hf_model, hf_tokenizer, prompt: str): """ print(f"\n🎭 Сравнение стратегий генерации через HF для промпта: '{prompt}'") print("=" * 70) - + strategies = [ {"name": "🎯 Жадный поиск", "do_sample": False, "temperature": 1.0}, {"name": "🎲 Вероятностная (temp=0.7)", "do_sample": True, "temperature": 0.7}, {"name": "🔥 Случайная (temp=1.2)", "do_sample": True, "temperature": 1.2}, - {"name": "❄️ Детерминированная (temp=0.3)", "do_sample": True, "temperature": 0.3}, + { + "name": "❄️ Детерминированная (temp=0.3)", + "do_sample": True, + "temperature": 0.3, + }, ] - + for strategy in strategies: print(f"\n{strategy['name']}:") try: config = GENERATION_CONFIG.copy() - config.update({ - "do_sample": strategy["do_sample"], - "temperature": strategy["temperature"], - "max_new_tokens": 20 - }) - + config.update( + { + "do_sample": strategy["do_sample"], + "temperature": strategy["temperature"], + "max_new_tokens": 20, + } + ) + generated = generate_with_hf_model(hf_model, hf_tokenizer, prompt, config) - + # Выделяем сгенерированную часть - generated_part = generated[len(prompt):] + generated_part = generated[len(prompt) :] print(f" 📤 Промпт: '{prompt}'") print(f" 🎯 Сгенерировано: '{generated_part}'") print(f" 📄 Полный текст: '{generated}'") - + except Exception as e: print(f" ❌ Ошибка: {e}") @@ -202,30 +208,30 @@ def test_different_hf_strategies(hf_model, hf_tokenizer, prompt: str): def analyze_hf_tokenization(hf_tokenizer, texts: list): """ Анализирует токенизацию через адаптированный токенизатор. - + Args: hf_tokenizer: Адаптированный токенизатор texts: Список текстов для анализа """ print(f"\n🔍 Анализ токенизации через HF адаптер:") print("=" * 60) - + for i, text in enumerate(texts): print(f"\nТекст {i+1}: '{text}'") - + # Токенизация через адаптер inputs = hf_tokenizer(text, return_tensors="pt") - tokens = inputs['input_ids'].tolist()[0] + tokens = inputs["input_ids"].tolist()[0] token_strings = hf_tokenizer.tokenize(text) - + print(f" Токены (ID): {tokens}") print(f" Токены (текст): {token_strings}") print(f" Количество токенов: {len(tokens)}") - + # Декодирование обратно decoded = hf_tokenizer.decode(tokens) print(f" Декодированный: '{decoded}'") - + if text == decoded: print(f" ✅ Декодирование корректно") else: @@ -235,51 +241,55 @@ def analyze_hf_tokenization(hf_tokenizer, texts: list): def interactive_hf_generation(hf_model, hf_tokenizer): """ Режим интерактивной генерации через HF интерфейс. - + Args: hf_model: Адаптированная модель hf_tokenizer: Адаптированный токенизатор """ print(f"\n💬 Интерактивная генерация через HF (для выхода введите 'exit')") print("-" * 60) - + while True: try: user_input = input("\n🔤 Введите промпт: ").strip() - - if user_input.lower() in ['exit', 'quit', 'выход']: + + if user_input.lower() in ["exit", "quit", "выход"]: break - + if not user_input: continue - + # Запрашиваем параметры try: max_tokens = int(input("📏 Макс. токенов [50]: ") or "50") temperature = float(input("🌡️ Температура [0.7]: ") or "0.7") do_sample_input = input("🎲 Сэмплирование (y/n) [y]: ").lower() - do_sample = do_sample_input != 'n' + do_sample = do_sample_input != "n" except: max_tokens = 50 temperature = 0.7 do_sample = True print("⚠️ Использую параметры по умолчанию") - + config = GENERATION_CONFIG.copy() - config.update({ - "max_new_tokens": max_tokens, - "temperature": temperature, - "do_sample": do_sample - }) - - generated = generate_with_hf_model(hf_model, hf_tokenizer, user_input, config) - - generated_part = generated[len(user_input):] + config.update( + { + "max_new_tokens": max_tokens, + "temperature": temperature, + "do_sample": do_sample, + } + ) + + generated = generate_with_hf_model( + hf_model, hf_tokenizer, user_input, config + ) + + generated_part = generated[len(user_input) :] print(f"\n🎯 Результат:") print(f" 📤 Промпт: '{user_input}'") print(f" 🎯 Сгенерировано: '{generated_part}'") print(f" 📄 Полный текст: '{generated}'") - + except KeyboardInterrupt: print("\n👋 Завершение работы...") break @@ -295,76 +305,79 @@ def main(): "model": "GPT через HFAdapter", "tokenizer": "BPE через HFTokenizerAdapter", "инструменты": "HuggingFace pipeline & генерация", - "стратегия": "интеграция с HF экосистемой" + "стратегия": "интеграция с HF экосистемой", } - + print_experiment_info(experiment_name, experiment_config) ensure_directories() logger = ExperimentLogger(experiment_name) - + try: # Загружаем модель и токенизатор в HF формате hf_model, hf_tokenizer, model_config = load_hf_model_and_tokenizer() - + # === Анализ токенизации === analysis_texts = [ "Искусственный интеллект", - "Нейронные сети", - "Машинное обучение" + "Нейронные сети", + "Машинное обучение", ] analyze_hf_tokenization(hf_tokenizer, analysis_texts) - + # === Тестирование HF pipeline === test_hf_pipeline(hf_model, hf_tokenizer) - + # === Генерация с разными промптами === print(f"\n🎯 Генерация текста через HF адаптер") print("=" * 60) - + for i, prompt in enumerate(TEST_PROMPTS): print(f"\n📝 Пример {i+1}/{len(TEST_PROMPTS)}") print("-" * 40) - + try: - generated = generate_with_hf_model(hf_model, hf_tokenizer, prompt, GENERATION_CONFIG) - + generated = generate_with_hf_model( + hf_model, hf_tokenizer, prompt, GENERATION_CONFIG + ) + # Выделяем сгенерированную часть - generated_part = generated[len(prompt):] - + generated_part = generated[len(prompt) :] + print(f"📤 Промпт: '{prompt}'") print(f"🎯 Сгенерировано: '{generated_part}'") print(f"📄 Полный текст: '{generated}'") print(f"📏 Длина: {len(generated)} символов") - + # Логируем успешную генерацию logger.log_metric(f"hf_generation_length_{i}", len(generated)) - + except Exception as e: print(f"❌ Ошибка при генерации: {e}") continue - + # === Сравнение стратегий генерации === test_prompt = "Искусственный" test_different_hf_strategies(hf_model, hf_tokenizer, test_prompt) - + # === Интерактивная генерация === interactive_hf_generation(hf_model, hf_tokenizer) - + # === Сохранение результатов === logger.save_logs("checkpoints/hf_integration_generation_logs.json") - + print(f"\n🎉 Эксперимент с HF интеграцией завершен успешно!") print(f"\n📚 Достигнутая интеграция:") print(f" ✅ Загрузка модели и токенизатора в HF формате") print(f" ✅ Использование HF pipeline") print(f" ✅ Генерация через стандартные HF интерфейсы") print(f" ✅ Совместимость с HF экосистемой") - + except FileNotFoundError as e: print(f"❌ {e}") except Exception as e: print(f"❌ Ошибка в эксперименте: {e}") import traceback + traceback.print_exc() diff --git a/experiments/hf_integration/simple_hf_training.py b/experiments/hf_integration/simple_hf_training.py index 6a4d0b1..0c619fc 100644 --- a/experiments/hf_integration/simple_hf_training.py +++ b/experiments/hf_integration/simple_hf_training.py @@ -19,141 +19,139 @@ from llm.tokenizers import BPETokenizer from hf_proxy import HFAdapter, HFTokenizerAdapter from shared.configs import ( - TRAIN_TEXTS, BASE_GPT_CONFIG, BPE_CONFIG, - TRAINING_CONFIG, PATHS, TEST_PROMPTS + TRAIN_TEXTS, + BASE_GPT_CONFIG, + BPE_CONFIG, + TRAINING_CONFIG, + PATHS, + TEST_PROMPTS, ) def create_dataset(hf_tokenizer, texts, max_length=128): """ Создает простой датасет для обучения. - + Args: hf_tokenizer: Адаптированный токенизатор texts: Список текстов max_length: Максимальная длина последовательности - + Returns: list: Список тензоров input_ids """ dataset = [] - + for text in texts: # Токенизируем текст inputs = hf_tokenizer( - text, - max_length=max_length, + text, + max_length=max_length, truncation=True, padding=False, - return_tensors="pt" + return_tensors="pt", ) - - input_ids = inputs['input_ids'][0] - + + input_ids = inputs["input_ids"][0] + # Создаем метки для языкового моделирования labels = input_ids.clone() - - dataset.append({ - 'input_ids': input_ids, - 'labels': labels - }) - + + dataset.append({"input_ids": input_ids, "labels": labels}) + return dataset def manual_training_loop(hf_model, hf_tokenizer, train_texts, val_texts, config): """ Ручной цикл обучения без использования Trainer. - + Args: hf_model: Адаптированная модель hf_tokenizer: Адаптированный токенизатор train_texts: Тексты для обучения val_texts: Тексты для валидации config: Конфигурация обучения - + Returns: dict: Результаты обучения """ print("🎯 Запуск ручного обучения...") - + # Создаем датасеты train_dataset = create_dataset(hf_tokenizer, train_texts) val_dataset = create_dataset(hf_tokenizer, val_texts) - + print(f"📊 Данные: {len(train_dataset)} train, {len(val_dataset)} validation") - + # Оптимизатор - optimizer = torch.optim.AdamW( - hf_model.parameters(), - lr=config["learning_rate"] - ) - + optimizer = torch.optim.AdamW(hf_model.parameters(), lr=config["learning_rate"]) + # Функция потерь loss_fn = nn.CrossEntropyLoss() - + # Обучение hf_model.train() train_losses = [] val_losses = [] - + for epoch in range(config["num_epochs"]): print(f"\n📅 Эпоха {epoch + 1}/{config['num_epochs']}") - + # Обучение epoch_train_loss = 0 for i, batch in enumerate(train_dataset): optimizer.zero_grad() - - input_ids = batch['input_ids'].unsqueeze(0) # [1, seq_len] - labels = batch['labels'].unsqueeze(0) # [1, seq_len] - + + input_ids = batch["input_ids"].unsqueeze(0) # [1, seq_len] + labels = batch["labels"].unsqueeze(0) # [1, seq_len] + # Forward pass outputs = hf_model(input_ids=input_ids, labels=labels) loss = outputs.loss - + # Backward pass loss.backward() optimizer.step() - + epoch_train_loss += loss.item() - + if i % 5 == 0: print(f" Batch {i}/{len(train_dataset)}: loss = {loss.item():.4f}") - + avg_train_loss = epoch_train_loss / len(train_dataset) train_losses.append(avg_train_loss) print(f" 📊 Средняя train loss: {avg_train_loss:.4f}") - + # Валидация hf_model.eval() epoch_val_loss = 0 with torch.no_grad(): for batch in val_dataset: - input_ids = batch['input_ids'].unsqueeze(0) - labels = batch['labels'].unsqueeze(0) - + input_ids = batch["input_ids"].unsqueeze(0) + labels = batch["labels"].unsqueeze(0) + outputs = hf_model(input_ids=input_ids, labels=labels) epoch_val_loss += outputs.loss.item() - + avg_val_loss = epoch_val_loss / len(val_dataset) val_losses.append(avg_val_loss) print(f" 📊 Средняя val loss: {avg_val_loss:.4f}") - + hf_model.train() - + return { - 'train_losses': train_losses, - 'val_losses': val_losses, - 'final_train_loss': train_losses[-1], - 'final_val_loss': val_losses[-1] + "train_losses": train_losses, + "val_losses": val_losses, + "final_train_loss": train_losses[-1], + "final_val_loss": val_losses[-1], } def test_generation_after_training(hf_model, hf_tokenizer, test_prompts): """ Тестирует генерацию после обучения. - + Args: hf_model: Обученная модель hf_tokenizer: Токенизатор @@ -161,24 +159,24 @@ def test_generation_after_training(hf_model, hf_tokenizer, test_prompts): """ print("\n🧪 Тестирование генерации после обучения...") hf_model.eval() - + for prompt in test_prompts[:3]: print(f"\n🔤 Промпт: '{prompt}'") - + try: inputs = hf_tokenizer(prompt, return_tensors="pt") - + with torch.no_grad(): generated = hf_model.generate( - input_ids=inputs['input_ids'], + input_ids=inputs["input_ids"], max_new_tokens=20, do_sample=True, - temperature=0.8 + temperature=0.8, ) - + generated_text = hf_tokenizer.decode(generated[0], skip_special_tokens=True) print(f"🎯 Результат: '{generated_text}'") - + except Exception as e: print(f"❌ Ошибка генерации: {e}") @@ -188,96 +186,102 @@ def main(): print("=" * 60) print("🚀 УПРОЩЕННОЕ ОБУЧЕНИЕ GPT С HF-PROXY") print("=" * 60) - + try: # === Подготовка данных === print("🔧 Подготовка данных...") - train_texts = TRAIN_TEXTS[:10] # Используем меньше данных для быстрого тестирования + train_texts = TRAIN_TEXTS[ + :10 + ] # Используем меньше данных для быстрого тестирования val_texts = TRAIN_TEXTS[10:12] - + print(f"📊 Данные: {len(train_texts)} train, {len(val_texts)} validation") - + # === Подготовка токенизатора === print("🔧 Подготовка токенизатора...") llm_tokenizer = BPETokenizer() llm_tokenizer.train( texts=train_texts, vocab_size=BPE_CONFIG["vocab_size"], - special_tokens=BPE_CONFIG["special_tokens"] + special_tokens=BPE_CONFIG["special_tokens"], ) - + hf_tokenizer = HFTokenizerAdapter(llm_tokenizer) print(f"✅ Токенизатор создан (vocab_size={hf_tokenizer.vocab_size})") - + # === Подготовка модели === print("🔧 Подготовка модели...") model_config = BASE_GPT_CONFIG.copy() model_config["vocab_size"] = hf_tokenizer.vocab_size - + llm_model = GPT(model_config) hf_model = HFAdapter.from_llm_model(llm_model) print(f"✅ Модель создана") - + # === Тестирование до обучения === print("\n🧪 Тестирование до обучения...") test_generation_after_training(hf_model, hf_tokenizer, TEST_PROMPTS) - + # === Обучение === print(f"\n🎯 Обучение модели...") training_config = { "learning_rate": TRAINING_CONFIG["learning_rate"], "num_epochs": 2, # Меньше эпох для быстрого тестирования - "batch_size": TRAINING_CONFIG["batch_size"] + "batch_size": TRAINING_CONFIG["batch_size"], } - + results = manual_training_loop( hf_model, hf_tokenizer, train_texts, val_texts, training_config ) - + print(f"\n📊 Результаты обучения:") print(f" Final train loss: {results['final_train_loss']:.4f}") print(f" Final val loss: {results['final_val_loss']:.4f}") - + # === Тестирование после обучения === print("\n🧪 Тестирование после обучения...") test_generation_after_training(hf_model, hf_tokenizer, TEST_PROMPTS) - + # === Сохранение модели === print(f"\n💾 Сохранение модели...") - + # Создаем директории os.makedirs("checkpoints/hf_simple_trained", exist_ok=True) os.makedirs("checkpoints/hf_simple_tokenizer", exist_ok=True) - + # Сохраняем токенизатор hf_tokenizer.save_pretrained("checkpoints/hf_simple_tokenizer") print("✅ Токенизатор сохранен") - + # Сохраняем модель HFAdapter.save_pretrained( - hf_model, - "checkpoints/hf_simple_trained", - tokenizer=hf_tokenizer + hf_model, "checkpoints/hf_simple_trained", tokenizer=hf_tokenizer ) print("✅ Модель сохранена") - + # Сохраняем результаты results_path = "checkpoints/simple_training_results.json" - with open(results_path, 'w', encoding='utf-8') as f: - json.dump({ - 'training_config': training_config, - 'model_config': model_config, - 'results': results - }, f, indent=2, ensure_ascii=False) + with open(results_path, "w", encoding="utf-8") as f: + json.dump( + { + "training_config": training_config, + "model_config": model_config, + "results": results, + }, + f, + indent=2, + ensure_ascii=False, + ) print(f"✅ Результаты сохранены в {results_path}") - + print(f"\n🎉 Упрощенное обучение завершено успешно!") print(f"\n💡 Для использования обученной модели:") print(f" uv run python experiments/hf_integration/generate_with_hf_tools.py") - + except Exception as e: print(f"❌ Ошибка в эксперименте: {e}") import traceback + traceback.print_exc() diff --git a/experiments/hf_integration/test_hf_proxy.py b/experiments/hf_integration/test_hf_proxy.py index d1123f1..1bc2435 100644 --- a/experiments/hf_integration/test_hf_proxy.py +++ b/experiments/hf_integration/test_hf_proxy.py @@ -16,158 +16,163 @@ from llm.tokenizers import BPETokenizer from hf_proxy import HFAdapter, HFTokenizerAdapter from shared.configs import ( - TRAIN_TEXTS, BASE_GPT_CONFIG, BPE_CONFIG, - TEST_PROMPTS, GENERATION_CONFIG + TRAIN_TEXTS, + BASE_GPT_CONFIG, + BPE_CONFIG, + TEST_PROMPTS, + GENERATION_CONFIG, ) def test_basic_hf_integration(): """Тестирует базовую интеграцию hf-proxy.""" print("🧪 Тестирование базовой интеграции hf-proxy...") - + # === Подготовка токенизатора === print("1. Подготовка токенизатора...") llm_tokenizer = BPETokenizer() llm_tokenizer.train( texts=TRAIN_TEXTS, vocab_size=BPE_CONFIG["vocab_size"], - special_tokens=BPE_CONFIG["special_tokens"] + special_tokens=BPE_CONFIG["special_tokens"], ) - + hf_tokenizer = HFTokenizerAdapter(llm_tokenizer) print(f" ✅ Токенизатор создан (vocab_size={hf_tokenizer.vocab_size})") - + # === Подготовка модели === print("2. Подготовка модели...") model_config = BASE_GPT_CONFIG.copy() model_config["vocab_size"] = hf_tokenizer.vocab_size - + llm_model = GPT(model_config) hf_model = HFAdapter.from_llm_model(llm_model) print(f" ✅ Модель создана") - + # === Тестирование токенизации === print("3. Тестирование токенизации...") test_texts = ["Искусственный интеллект", "Нейронные сети"] - + for text in test_texts: print(f" 📝 Текст: '{text}'") - + # Оригинальный токенизатор original_tokens = llm_tokenizer.encode(text) print(f" Оригинальный: {len(original_tokens)} токенов") - + # HF адаптер hf_inputs = hf_tokenizer(text, return_tensors="pt") print(f" HF адаптер: {hf_inputs['input_ids'].shape}") - + # Декодирование - decoded = hf_tokenizer.decode(hf_inputs['input_ids'][0]) + decoded = hf_tokenizer.decode(hf_inputs["input_ids"][0]) print(f" Декодированный: '{decoded}'") - + # === Тестирование forward pass === print("4. Тестирование forward pass...") for text in test_texts: hf_inputs = hf_tokenizer(text, return_tensors="pt") - + with torch.no_grad(): outputs = hf_model(**hf_inputs) - + print(f" 📝 '{text}' -> logits: {outputs.logits.shape}") - + # === Тестирование генерации === print("5. Тестирование генерации...") hf_model.eval() - + for prompt in TEST_PROMPTS[:3]: print(f" 🔤 Промпт: '{prompt}'") - + try: inputs = hf_tokenizer(prompt, return_tensors="pt") - + with torch.no_grad(): generated = hf_model.generate( - input_ids=inputs['input_ids'], + input_ids=inputs["input_ids"], max_new_tokens=10, do_sample=True, - temperature=0.8 + temperature=0.8, ) - + generated_text = hf_tokenizer.decode(generated[0], skip_special_tokens=True) print(f" 🎯 Результат: '{generated_text}'") - + except Exception as e: print(f" ❌ Ошибка: {e}") - + # === Тестирование сохранения/загрузки === print("6. Тестирование сохранения/загрузки...") try: # Сохраняем токенизатор hf_tokenizer.save_pretrained("test_save/tokenizer") print(" ✅ Токенизатор сохранен") - + # Сохраняем модель HFAdapter.save_pretrained(hf_model, "test_save/model", tokenizer=hf_tokenizer) print(" ✅ Модель сохранена") - + # Загружаем токенизатор loaded_tokenizer = HFTokenizerAdapter.from_pretrained("test_save/tokenizer") print(f" ✅ Токенизатор загружен (vocab_size={loaded_tokenizer.vocab_size})") - + # Загружаем модель model_path = os.path.join("test_save/model", "pytorch_model.bin") loaded_model = HFAdapter.from_pretrained(model_path) print(" ✅ Модель загружена") - + # Проверяем работоспособность загруженной модели test_input = hf_tokenizer("Тест", return_tensors="pt") with torch.no_grad(): loaded_outputs = loaded_model(**test_input) - print(f" ✅ Загруженная модель работает (logits: {loaded_outputs.logits.shape})") - + print( + f" ✅ Загруженная модель работает (logits: {loaded_outputs.logits.shape})" + ) + except Exception as e: print(f" ❌ Ошибка сохранения/загрузки: {e}") - + print("\n🎉 Базовое тестирование hf-proxy завершено!") def test_hf_tokenizer_methods(): """Тестирует различные методы HF токенизатора.""" print("\n🧪 Тестирование методов HF токенизатора...") - + # Создаем токенизатор llm_tokenizer = BPETokenizer() llm_tokenizer.train( texts=TRAIN_TEXTS[:5], vocab_size=500, - special_tokens=BPE_CONFIG["special_tokens"] + special_tokens=BPE_CONFIG["special_tokens"], ) - + hf_tokenizer = HFTokenizerAdapter(llm_tokenizer) - + test_text = "Искусственный интеллект и машинное обучение" - + # Тестируем разные методы print("1. Метод __call__:") result = hf_tokenizer(test_text, return_tensors="pt") print(f" Результат: {result}") - + print("2. Метод encode:") encoded = hf_tokenizer.encode(test_text) print(f" Закодировано: {encoded}") - + print("3. Метод decode:") decoded = hf_tokenizer.decode(encoded) print(f" Декодировано: '{decoded}'") - + print("4. Метод tokenize:") tokens = hf_tokenizer.tokenize(test_text) print(f" Токены: {tokens}") - + print("5. Метод get_vocab:") vocab = hf_tokenizer.get_vocab() print(f" Размер словаря: {len(vocab)}") - + print("✅ Все методы токенизатора работают!") @@ -176,14 +181,14 @@ def main(): print("=" * 60) print("🧪 ТЕСТИРОВАНИЕ HF-PROXY") print("=" * 60) - + try: # Тестируем базовую интеграцию test_basic_hf_integration() - + # Тестируем методы токенизатора test_hf_tokenizer_methods() - + print("\n" + "=" * 60) print("🎉 ВСЕ ТЕСТЫ ПРОЙДЕНЫ УСПЕШНО!") print("=" * 60) @@ -195,10 +200,11 @@ def main(): print(" ✅ Генерация текста") print(" ✅ Сохранение и загрузка моделей") print(" ✅ Все методы HF токенизатора") - + except Exception as e: print(f"\n❌ Ошибка в тестировании: {e}") import traceback + traceback.print_exc() diff --git a/experiments/hf_integration/train_with_hf_trainer.py b/experiments/hf_integration/train_with_hf_trainer.py index 7d81c00..745a093 100644 --- a/experiments/hf_integration/train_with_hf_trainer.py +++ b/experiments/hf_integration/train_with_hf_trainer.py @@ -17,28 +17,34 @@ from llm.tokenizers import BPETokenizer from hf_proxy import HFAdapter, HFTokenizerAdapter from shared.configs import ( - TRAIN_TEXTS, BASE_GPT_CONFIG, BPE_CONFIG, - TRAINING_CONFIG, PATHS, TEST_PROMPTS + TRAIN_TEXTS, + BASE_GPT_CONFIG, + BPE_CONFIG, + TRAINING_CONFIG, + PATHS, + TEST_PROMPTS, ) from shared.data import ( - load_training_data, ensure_directories, - print_experiment_info, ExperimentLogger + load_training_data, + ensure_directories, + print_experiment_info, + ExperimentLogger, ) def setup_hf_training(): """ Настраивает окружение для обучения через HuggingFace Trainer. - + Returns: tuple: (hf_model, hf_tokenizer, llm_tokenizer, model_config) """ print("🔧 Настройка HuggingFace обучения...") - + # === Подготовка данных === train_texts, val_texts = load_training_data() print(f"📊 Данные: {len(train_texts)} train, {len(val_texts)} validation") - + # === Обучение/загрузка токенизатора === if os.path.exists(PATHS["bpe_tokenizer"]): print("📝 Загрузка BPE токенизатора...") @@ -50,55 +56,55 @@ def setup_hf_training(): llm_tokenizer.train( texts=TRAIN_TEXTS, vocab_size=BPE_CONFIG["vocab_size"], - special_tokens=BPE_CONFIG["special_tokens"] + special_tokens=BPE_CONFIG["special_tokens"], ) llm_tokenizer.save(PATHS["bpe_tokenizer"]) print(f"✅ Токенизатор обучен и сохранен") - + # === Создание адаптера токенизатора === print("🔧 Создание адаптера HuggingFace для токенизатора...") hf_tokenizer = HFTokenizerAdapter(llm_tokenizer) print(f"✅ Адаптер токенизатора создан") - + # === Инициализация модели === model_config = BASE_GPT_CONFIG.copy() model_config["vocab_size"] = llm_tokenizer.get_vocab_size() - + print("🔧 Создание GPT модели...") llm_model = GPT(model_config) - + # === Создание адаптера модели === print("🔧 Создание адаптера HuggingFace для модели...") hf_model = HFAdapter.from_llm_model(llm_model) print(f"✅ Адаптер модели создан") - + return hf_model, hf_tokenizer, llm_tokenizer, model_config, train_texts, val_texts def test_hf_integration(hf_model, hf_tokenizer, llm_tokenizer): """ Тестирует интеграцию с HuggingFace инструментами. - + Args: hf_model: Адаптированная модель hf_tokenizer: Адаптированный токенизатор llm_tokenizer: Оригинальный токенизатор """ print("\n🧪 Тестирование интеграции с HuggingFace...") - + test_texts = ["Искусственный интеллект", "Нейронные сети"] - + for text in test_texts: print(f"\n🔤 Текст: '{text}'") - + # Тестируем адаптированный токенизатор hf_inputs = hf_tokenizer(text, return_tensors="pt") print(f" HF токенизатор: {hf_inputs['input_ids'].shape}") - + # Тестируем оригинальный токенизатор для сравнения original_tokens = llm_tokenizer.encode(text) print(f" Оригинальный токенизатор: {len(original_tokens)} токенов") - + # Тестируем forward pass через адаптированную модель try: with torch.no_grad(): @@ -114,28 +120,35 @@ def main(): experiment_name = "Обучение GPT через HF Trainer (с hf-proxy)" experiment_config = { "model": "GPT через HFAdapter", - "tokenizer": "BPE через HFTokenizerAdapter", + "tokenizer": "BPE через HFTokenizerAdapter", "trainer": "HuggingFace Trainer", "vocab_size": BPE_CONFIG["vocab_size"], - "training_epochs": TRAINING_CONFIG["num_epochs"] + "training_epochs": TRAINING_CONFIG["num_epochs"], } - + print_experiment_info(experiment_name, experiment_config) ensure_directories() logger = ExperimentLogger(experiment_name) - + try: # Настраиваем обучение - hf_model, hf_tokenizer, llm_tokenizer, model_config, train_texts, val_texts = setup_hf_training() - + ( + hf_model, + hf_tokenizer, + llm_tokenizer, + model_config, + train_texts, + val_texts, + ) = setup_hf_training() + # Тестируем интеграцию test_hf_integration(hf_model, hf_tokenizer, llm_tokenizer) - + # === Подготовка датасетов HuggingFace === print(f"\n📊 Подготовка датасетов HuggingFace...") - + from datasets import Dataset - + def tokenize_function(examples): """Функция токенизации для HF datasets.""" # Используем адаптированный токенизатор @@ -147,11 +160,11 @@ def main(): ) tokenized["labels"] = tokenized["input_ids"].copy() return tokenized - + # Создаем датасеты train_dataset = Dataset.from_dict({"text": train_texts}) val_dataset = Dataset.from_dict({"text": val_texts}) - + # Токенизируем train_dataset = train_dataset.map( tokenize_function, @@ -163,26 +176,26 @@ def main(): batched=True, remove_columns=val_dataset.column_names, ) - + print(f" Train датасет: {len(train_dataset)} примеров") print(f" Validation датасет: {len(val_dataset)} примеров") - + # === Настройка HuggingFace Trainer === print(f"\n🔧 Настройка HuggingFace Trainer...") - + from transformers import ( - Trainer, + Trainer, TrainingArguments, - DataCollatorForLanguageModeling + DataCollatorForLanguageModeling, ) - + # Data collator для языкового моделирования data_collator = DataCollatorForLanguageModeling( tokenizer=hf_tokenizer, mlm=False, pad_to_multiple_of=8, ) - + # Аргументы обучения training_args = TrainingArguments( output_dir=PATHS["hf_model"], @@ -204,7 +217,7 @@ def main(): dataloader_pin_memory=False, report_to=None, ) - + # Создаем Trainer trainer = Trainer( model=hf_model, @@ -213,84 +226,87 @@ def main(): eval_dataset=val_dataset, data_collator=data_collator, ) - + print("✅ HuggingFace Trainer настроен") - + # === Запуск обучения === print(f"\n🎯 Запуск обучения через HuggingFace Trainer...") - + train_result = trainer.train() - + # Сохраняем лучшую модель trainer.save_model() hf_tokenizer.save_pretrained(PATHS["hf_model"]) - + print("✅ Обучение завершено успешно!") print(f"📊 Final train loss: {train_result.metrics['train_loss']:.4f}") - + if "eval_loss" in train_result.metrics: print(f"📊 Final eval loss: {train_result.metrics['eval_loss']:.4f}") - + # === Сохранение через hf-proxy === print(f"\n💾 Сохранение через hf-proxy...") - + from hf_proxy import convert_to_hf_format - + # Сохраняем токенизатор в HF формате hf_tokenizer_dir = PATHS["hf_tokenizer"] hf_tokenizer.save_pretrained(hf_tokenizer_dir) - + # Сохраняем модель через hf-proxy hf_proxy_dir = PATHS["hf_proxy_model"] HFAdapter.save_pretrained(hf_model, hf_proxy_dir, tokenizer=hf_tokenizer) - + print(f"✅ Модель сохранена в HF формате:") print(f" - {PATHS['hf_model']}: стандартный HF формат") print(f" - {hf_proxy_dir}: через hf-proxy") print(f" - {hf_tokenizer_dir}: токенизатор в HF формате") - + # === Тестирование генерации === print(f"\n🧪 Тестирование генерации после обучения...") hf_model.eval() - + for prompt in TEST_PROMPTS[:3]: print(f"\n🔤 Промпт: '{prompt}'") - + try: inputs = hf_tokenizer(prompt, return_tensors="pt") - + with torch.no_grad(): generated = hf_model.generate( - input_ids=inputs['input_ids'], + input_ids=inputs["input_ids"], max_new_tokens=20, do_sample=True, - temperature=0.8 + temperature=0.8, ) - - generated_text = hf_tokenizer.decode(generated[0], skip_special_tokens=True) + + generated_text = hf_tokenizer.decode( + generated[0], skip_special_tokens=True + ) print(f"🎯 Результат: '{generated_text}'") - + except Exception as e: print(f"❌ Ошибка генерации: {e}") - + # === Сохранение результатов === results = { "experiment": experiment_name, "model_config": model_config, "training_config": TRAINING_CONFIG, - "final_loss": train_result.metrics.get('train_loss', 'N/A'), - "eval_loss": train_result.metrics.get('eval_loss', 'N/A') + "final_loss": train_result.metrics.get("train_loss", "N/A"), + "eval_loss": train_result.metrics.get("eval_loss", "N/A"), } - + logger.save_logs("checkpoints/hf_integration_training_logs.json") - + print(f"\n🎉 Эксперимент с HF интеграцией завершен успешно!") print(f"\n💡 Для использования обученной модели:") print(f" uv run python experiments/hf_integration/generate_with_hf_tools.py") - + except Exception as e: print(f"❌ Ошибка в эксперименте: {e}") import traceback + traceback.print_exc() diff --git a/experiments/llm_only/generate_gpt2_bpe.py b/experiments/llm_only/generate_gpt2_bpe.py index f634970..efd6368 100644 --- a/experiments/llm_only/generate_gpt2_bpe.py +++ b/experiments/llm_only/generate_gpt2_bpe.py @@ -15,18 +15,14 @@ sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from llm.models.gpt import GPT2 from llm.tokenizers import BPETokenizer -from shared.configs import ( - BASE_GPT_CONFIG, TEST_PROMPTS, GENERATION_CONFIG, PATHS -) -from shared.data import ( - print_experiment_info, ensure_directories, ExperimentLogger -) +from shared.configs import BASE_GPT_CONFIG, TEST_PROMPTS, GENERATION_CONFIG, PATHS +from shared.data import print_experiment_info, ensure_directories, ExperimentLogger def load_model_and_tokenizer() -> tuple: """ Загружает обученную модель и токенизатор. - + Returns: tuple: (модель, токенизатор, конфигурация) """ @@ -36,62 +32,60 @@ def load_model_and_tokenizer() -> tuple: f"Модель не найдена: {PATHS['gpt_bpe_model']}\n" f"Сначала обучите модель: uv run python experiments/llm_only/train_gpt_bpe.py" ) - + if not os.path.exists(PATHS["bpe_tokenizer"]): - raise FileNotFoundError( - f"Токенизатор не найден: {PATHS['bpe_tokenizer']}" - ) - + raise FileNotFoundError(f"Токенизатор не найден: {PATHS['bpe_tokenizer']}") + # Загружаем конфигурацию модели import json - with open(PATHS["gpt_bpe_config"], 'r', encoding='utf-8') as f: + + with open(PATHS["gpt_bpe_config"], "r", encoding="utf-8") as f: model_config = json.load(f) - + # Загружаем токенизатор print("🔧 Загрузка BPE токенизатора...") tokenizer = BPETokenizer.load(PATHS["bpe_tokenizer"]) print(f"✅ Токенизатор загружен (vocab_size={tokenizer.get_vocab_size()})") - + # Загружаем модель print("🔧 Загрузка GPT2 модели...") model = GPT2(model_config) - model.load_state_dict(torch.load(PATHS["gpt_bpe_model"], map_location='cpu')) + model.load_state_dict(torch.load(PATHS["gpt_bpe_model"], map_location="cpu")) model.eval() print("✅ Модель загружена") - + return model, tokenizer, model_config def generate_text( - model: GPT2, - tokenizer: BPETokenizer, - prompt: str, - config: dict + model: GPT2, tokenizer: BPETokenizer, prompt: str, config: dict ) -> str: """ Генерирует текст на основе промпта. - + Args: model: Обученная GPT модель tokenizer: BPE токенизатор prompt: Входной текст config: Конфигурация генерации - + Returns: str: Сгенерированный текст """ print(f"🔤 Промпт: '{prompt}'") - print(f"📊 Параметры: max_tokens={config['max_new_tokens']}, " - f"temp={config['temperature']}, sample={config['do_sample']}") - + print( + f"📊 Параметры: max_tokens={config['max_new_tokens']}, " + f"temp={config['temperature']}, sample={config['do_sample']}" + ) + # Кодируем промпт input_ids = tokenizer.encode(prompt, add_special_tokens=False) input_tensor = torch.tensor([input_ids], dtype=torch.long) - + print(f"🎯 Токены промпта: {input_ids}") print(f"🎯 Токены (текст): {tokenizer.tokenize(prompt)}") print("🔄 Генерация...") - + # Генерируем текст with torch.no_grad(): generated_ids = model.generate( @@ -100,19 +94,19 @@ def generate_text( do_sample=config["do_sample"], temperature=config["temperature"], top_k=config["top_k"], - top_p=config["top_p"] + top_p=config["top_p"], ) - + # Декодируем результат generated_text = tokenizer.decode(generated_ids[0].tolist()) - + return generated_text def test_different_strategies(model: GPT2, tokenizer: BPETokenizer, prompt: str): """ Тестирует разные стратегии генерации на одном промпте. - + Args: model: Обученная модель tokenizer: BPE токенизатор @@ -120,32 +114,38 @@ def test_different_strategies(model: GPT2, tokenizer: BPETokenizer, prompt: str) """ print(f"\n🎭 Сравнение стратегий генерации для промпта: '{prompt}'") print("=" * 60) - + strategies = [ {"name": "🎯 Жадный поиск", "do_sample": False, "temperature": 1.0}, {"name": "🎲 Вероятностная (temp=0.7)", "do_sample": True, "temperature": 0.7}, {"name": "🔥 Случайная (temp=1.2)", "do_sample": True, "temperature": 1.2}, - {"name": "❄️ Детерминированная (temp=0.3)", "do_sample": True, "temperature": 0.3}, + { + "name": "❄️ Детерминированная (temp=0.3)", + "do_sample": True, + "temperature": 0.3, + }, ] - + for strategy in strategies: print(f"\n{strategy['name']}:") try: config = GENERATION_CONFIG.copy() - config.update({ - "do_sample": strategy["do_sample"], - "temperature": strategy["temperature"], - "max_new_tokens": 20 - }) - + config.update( + { + "do_sample": strategy["do_sample"], + "temperature": strategy["temperature"], + "max_new_tokens": 20, + } + ) + generated = generate_text(model, tokenizer, prompt, config) - + # Выделяем сгенерированную часть - generated_part = generated[len(prompt):] + generated_part = generated[len(prompt) :] print(f" 📤 Промпт: '{prompt}'") print(f" 🎯 Сгенерировано: '{generated_part}'") print(f" 📄 Полный текст: '{generated}'") - + except Exception as e: print(f" ❌ Ошибка: {e}") @@ -153,26 +153,26 @@ def test_different_strategies(model: GPT2, tokenizer: BPETokenizer, prompt: str) def analyze_tokenization(tokenizer: BPETokenizer, texts: list): """ Анализирует токенизацию различных текстов. - + Args: tokenizer: BPE токенизатор texts: Список текстов для анализа """ print(f"\n🔍 Анализ токенизации BPE:") print("=" * 50) - + for i, text in enumerate(texts): print(f"\nТекст {i+1}: '{text}'") - + # Токенизация tokens = tokenizer.encode(text, add_special_tokens=False) token_strings = tokenizer.tokenize(text) - + print(f" Токены (ID): {tokens}") print(f" Токены (текст): {token_strings}") print(f" Количество токенов: {len(tokens)}") print(f" Эффективность: {len(text)} символов → {len(tokens)} токенов") - + # Декодирование обратно decoded = tokenizer.decode(tokens) if text == decoded: @@ -184,51 +184,53 @@ def analyze_tokenization(tokenizer: BPETokenizer, texts: list): def interactive_generation(model: GPT2, tokenizer: BPETokenizer): """ Режим интерактивной генерации. - + Args: model: Обученная модель tokenizer: BPE токенизатор """ print(f"\n💬 Интерактивная генерация (для выхода введите 'exit')") print("-" * 50) - + while True: try: user_input = input("\n🔤 Введите промпт: ").strip() - - if user_input.lower() in ['exit', 'quit', 'выход']: + + if user_input.lower() in ["exit", "quit", "выход"]: break - + if not user_input: continue - + # Запрашиваем параметры try: max_tokens = int(input("📏 Макс. токенов [50]: ") or "50") temperature = float(input("🌡️ Температура [0.7]: ") or "0.7") do_sample_input = input("🎲 Сэмплирование (y/n) [y]: ").lower() - do_sample = do_sample_input != 'n' + do_sample = do_sample_input != "n" except: max_tokens = 50 temperature = 0.7 do_sample = True print("⚠️ Использую параметры по умолчанию") - + config = GENERATION_CONFIG.copy() - config.update({ - "max_new_tokens": max_tokens, - "temperature": temperature, - "do_sample": do_sample - }) - + config.update( + { + "max_new_tokens": max_tokens, + "temperature": temperature, + "do_sample": do_sample, + } + ) + generated = generate_text(model, tokenizer, user_input, config) - - generated_part = generated[len(user_input):] + + generated_part = generated[len(user_input) :] print(f"\n🎯 Результат:") print(f" 📤 Промпт: '{user_input}'") print(f" 🎯 Сгенерировано: '{generated_part}'") print(f" 📄 Полный текст: '{generated}'") - + except KeyboardInterrupt: print("\n👋 Завершение работы...") break @@ -244,68 +246,69 @@ def main(): "model": "GPT2 с BPE токенизатором", "стратегия": "автономная генерация", "вход": "промпты", - "выход": "сгенерированный текст" + "выход": "сгенерированный текст", } - + print_experiment_info(experiment_name, experiment_config) ensure_directories() logger = ExperimentLogger(experiment_name) - + try: # Загружаем модель и токенизатор model, tokenizer, model_config = load_model_and_tokenizer() - + # === Анализ токенизации === analysis_texts = [ "Искусственный интеллект", "Нейронные сети", - "Машинное обучение", + "Машинное обучение", ] analyze_tokenization(tokenizer, analysis_texts) - + # === Генерация с разными промптами === print(f"\n🎯 Генерация текста с разными промптами") print("=" * 60) - + for i, prompt in enumerate(TEST_PROMPTS): print(f"\n📝 Пример {i+1}/{len(TEST_PROMPTS)}") print("-" * 40) - + try: generated = generate_text(model, tokenizer, prompt, GENERATION_CONFIG) - + # Выделяем сгенерированную часть - generated_part = generated[len(prompt):] - + generated_part = generated[len(prompt) :] + print(f"📤 Промпт: '{prompt}'") print(f"🎯 Сгенерировано: '{generated_part}'") print(f"📄 Полный текст: '{generated}'") print(f"📏 Длина: {len(generated)} символов") - + # Логируем успешную генерацию logger.log_metric(f"generation_length_{i}", len(generated)) - + except Exception as e: print(f"❌ Ошибка при генерации: {e}") continue - + # === Сравнение стратегий генерации === test_prompt = "Искусственный" test_different_strategies(model, tokenizer, test_prompt) - + # === Интерактивная генерация === interactive_generation(model, tokenizer) - + # === Сохранение результатов === logger.save_logs("checkpoints/llm_only_generation_logs.json") - + print(f"\n🎉 Эксперимент генерации завершен успешно!") - + except FileNotFoundError as e: print(f"❌ {e}") except Exception as e: print(f"❌ Ошибка в эксперименте: {e}") import traceback + traceback.print_exc() diff --git a/experiments/llm_only/generate_gpt_bpe.py b/experiments/llm_only/generate_gpt_bpe.py index daa3ae1..b74d401 100644 --- a/experiments/llm_only/generate_gpt_bpe.py +++ b/experiments/llm_only/generate_gpt_bpe.py @@ -15,18 +15,14 @@ sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from llm.models.gpt import GPT from llm.tokenizers import BPETokenizer -from shared.configs import ( - BASE_GPT_CONFIG, TEST_PROMPTS, GENERATION_CONFIG, PATHS -) -from shared.data import ( - print_experiment_info, ensure_directories, ExperimentLogger -) +from shared.configs import BASE_GPT_CONFIG, TEST_PROMPTS, GENERATION_CONFIG, PATHS +from shared.data import print_experiment_info, ensure_directories, ExperimentLogger def load_model_and_tokenizer() -> tuple: """ Загружает обученную модель и токенизатор. - + Returns: tuple: (модель, токенизатор, конфигурация) """ @@ -36,62 +32,60 @@ def load_model_and_tokenizer() -> tuple: f"Модель не найдена: {PATHS['gpt_bpe_model']}\n" f"Сначала обучите модель: uv run python experiments/llm_only/train_gpt_bpe.py" ) - + if not os.path.exists(PATHS["bpe_tokenizer"]): - raise FileNotFoundError( - f"Токенизатор не найден: {PATHS['bpe_tokenizer']}" - ) - + raise FileNotFoundError(f"Токенизатор не найден: {PATHS['bpe_tokenizer']}") + # Загружаем конфигурацию модели import json - with open(PATHS["gpt_bpe_config"], 'r', encoding='utf-8') as f: + + with open(PATHS["gpt_bpe_config"], "r", encoding="utf-8") as f: model_config = json.load(f) - + # Загружаем токенизатор print("🔧 Загрузка BPE токенизатора...") tokenizer = BPETokenizer.load(PATHS["bpe_tokenizer"]) print(f"✅ Токенизатор загружен (vocab_size={tokenizer.get_vocab_size()})") - + # Загружаем модель print("🔧 Загрузка GPT модели...") model = GPT(model_config) - model.load_state_dict(torch.load(PATHS["gpt_bpe_model"], map_location='cpu')) + model.load_state_dict(torch.load(PATHS["gpt_bpe_model"], map_location="cpu")) model.eval() print("✅ Модель загружена") - + return model, tokenizer, model_config def generate_text( - model: GPT, - tokenizer: BPETokenizer, - prompt: str, - config: dict + model: GPT, tokenizer: BPETokenizer, prompt: str, config: dict ) -> str: """ Генерирует текст на основе промпта. - + Args: model: Обученная GPT модель tokenizer: BPE токенизатор prompt: Входной текст config: Конфигурация генерации - + Returns: str: Сгенерированный текст """ print(f"🔤 Промпт: '{prompt}'") - print(f"📊 Параметры: max_tokens={config['max_new_tokens']}, " - f"temp={config['temperature']}, sample={config['do_sample']}") - + print( + f"📊 Параметры: max_tokens={config['max_new_tokens']}, " + f"temp={config['temperature']}, sample={config['do_sample']}" + ) + # Кодируем промпт input_ids = tokenizer.encode(prompt, add_special_tokens=False) input_tensor = torch.tensor([input_ids], dtype=torch.long) - + print(f"🎯 Токены промпта: {input_ids}") print(f"🎯 Токены (текст): {tokenizer.tokenize(prompt)}") print("🔄 Генерация...") - + # Генерируем текст with torch.no_grad(): generated_ids = model.generate( @@ -100,19 +94,19 @@ def generate_text( do_sample=config["do_sample"], temperature=config["temperature"], top_k=config["top_k"], - top_p=config["top_p"] + top_p=config["top_p"], ) - + # Декодируем результат generated_text = tokenizer.decode(generated_ids[0].tolist()) - + return generated_text def test_different_strategies(model: GPT, tokenizer: BPETokenizer, prompt: str): """ Тестирует разные стратегии генерации на одном промпте. - + Args: model: Обученная модель tokenizer: BPE токенизатор @@ -120,32 +114,38 @@ def test_different_strategies(model: GPT, tokenizer: BPETokenizer, prompt: str): """ print(f"\n🎭 Сравнение стратегий генерации для промпта: '{prompt}'") print("=" * 60) - + strategies = [ {"name": "🎯 Жадный поиск", "do_sample": False, "temperature": 1.0}, {"name": "🎲 Вероятностная (temp=0.7)", "do_sample": True, "temperature": 0.7}, {"name": "🔥 Случайная (temp=1.2)", "do_sample": True, "temperature": 1.2}, - {"name": "❄️ Детерминированная (temp=0.3)", "do_sample": True, "temperature": 0.3}, + { + "name": "❄️ Детерминированная (temp=0.3)", + "do_sample": True, + "temperature": 0.3, + }, ] - + for strategy in strategies: print(f"\n{strategy['name']}:") try: config = GENERATION_CONFIG.copy() - config.update({ - "do_sample": strategy["do_sample"], - "temperature": strategy["temperature"], - "max_new_tokens": 20 - }) - + config.update( + { + "do_sample": strategy["do_sample"], + "temperature": strategy["temperature"], + "max_new_tokens": 20, + } + ) + generated = generate_text(model, tokenizer, prompt, config) - + # Выделяем сгенерированную часть - generated_part = generated[len(prompt):] + generated_part = generated[len(prompt) :] print(f" 📤 Промпт: '{prompt}'") print(f" 🎯 Сгенерировано: '{generated_part}'") print(f" 📄 Полный текст: '{generated}'") - + except Exception as e: print(f" ❌ Ошибка: {e}") @@ -153,26 +153,26 @@ def test_different_strategies(model: GPT, tokenizer: BPETokenizer, prompt: str): def analyze_tokenization(tokenizer: BPETokenizer, texts: list): """ Анализирует токенизацию различных текстов. - + Args: tokenizer: BPE токенизатор texts: Список текстов для анализа """ print(f"\n🔍 Анализ токенизации BPE:") print("=" * 50) - + for i, text in enumerate(texts): print(f"\nТекст {i+1}: '{text}'") - + # Токенизация tokens = tokenizer.encode(text, add_special_tokens=False) token_strings = tokenizer.tokenize(text) - + print(f" Токены (ID): {tokens}") print(f" Токены (текст): {token_strings}") print(f" Количество токенов: {len(tokens)}") print(f" Эффективность: {len(text)} символов → {len(tokens)} токенов") - + # Декодирование обратно decoded = tokenizer.decode(tokens) if text == decoded: @@ -184,51 +184,53 @@ def analyze_tokenization(tokenizer: BPETokenizer, texts: list): def interactive_generation(model: GPT, tokenizer: BPETokenizer): """ Режим интерактивной генерации. - + Args: model: Обученная модель tokenizer: BPE токенизатор """ print(f"\n💬 Интерактивная генерация (для выхода введите 'exit')") print("-" * 50) - + while True: try: user_input = input("\n🔤 Введите промпт: ").strip() - - if user_input.lower() in ['exit', 'quit', 'выход']: + + if user_input.lower() in ["exit", "quit", "выход"]: break - + if not user_input: continue - + # Запрашиваем параметры try: max_tokens = int(input("📏 Макс. токенов [50]: ") or "50") temperature = float(input("🌡️ Температура [0.7]: ") or "0.7") do_sample_input = input("🎲 Сэмплирование (y/n) [y]: ").lower() - do_sample = do_sample_input != 'n' + do_sample = do_sample_input != "n" except: max_tokens = 50 temperature = 0.7 do_sample = True print("⚠️ Использую параметры по умолчанию") - + config = GENERATION_CONFIG.copy() - config.update({ - "max_new_tokens": max_tokens, - "temperature": temperature, - "do_sample": do_sample - }) - + config.update( + { + "max_new_tokens": max_tokens, + "temperature": temperature, + "do_sample": do_sample, + } + ) + generated = generate_text(model, tokenizer, user_input, config) - - generated_part = generated[len(user_input):] + + generated_part = generated[len(user_input) :] print(f"\n🎯 Результат:") print(f" 📤 Промпт: '{user_input}'") print(f" 🎯 Сгенерировано: '{generated_part}'") print(f" 📄 Полный текст: '{generated}'") - + except KeyboardInterrupt: print("\n👋 Завершение работы...") break @@ -244,68 +246,69 @@ def main(): "model": "GPT с BPE токенизатором", "стратегия": "автономная генерация", "вход": "промпты", - "выход": "сгенерированный текст" + "выход": "сгенерированный текст", } - + print_experiment_info(experiment_name, experiment_config) ensure_directories() logger = ExperimentLogger(experiment_name) - + try: # Загружаем модель и токенизатор model, tokenizer, model_config = load_model_and_tokenizer() - + # === Анализ токенизации === analysis_texts = [ "Искусственный интеллект", "Нейронные сети", - "Машинное обучение", + "Машинное обучение", ] analyze_tokenization(tokenizer, analysis_texts) - + # === Генерация с разными промптами === print(f"\n🎯 Генерация текста с разными промптами") print("=" * 60) - + for i, prompt in enumerate(TEST_PROMPTS): print(f"\n📝 Пример {i+1}/{len(TEST_PROMPTS)}") print("-" * 40) - + try: generated = generate_text(model, tokenizer, prompt, GENERATION_CONFIG) - + # Выделяем сгенерированную часть - generated_part = generated[len(prompt):] - + generated_part = generated[len(prompt) :] + print(f"📤 Промпт: '{prompt}'") print(f"🎯 Сгенерировано: '{generated_part}'") print(f"📄 Полный текст: '{generated}'") print(f"📏 Длина: {len(generated)} символов") - + # Логируем успешную генерацию logger.log_metric(f"generation_length_{i}", len(generated)) - + except Exception as e: print(f"❌ Ошибка при генерации: {e}") continue - + # === Сравнение стратегий генерации === test_prompt = "Искусственный" test_different_strategies(model, tokenizer, test_prompt) - + # === Интерактивная генерация === interactive_generation(model, tokenizer) - + # === Сохранение результатов === logger.save_logs("checkpoints/llm_only_generation_logs.json") - + print(f"\n🎉 Эксперимент генерации завершен успешно!") - + except FileNotFoundError as e: print(f"❌ {e}") except Exception as e: print(f"❌ Ошибка в эксперименте: {e}") import traceback + traceback.print_exc() diff --git a/experiments/llm_only/train_gpt2_bpe.py b/experiments/llm_only/train_gpt2_bpe.py index 7fb03fe..b6aa703 100644 --- a/experiments/llm_only/train_gpt2_bpe.py +++ b/experiments/llm_only/train_gpt2_bpe.py @@ -18,71 +18,77 @@ from llm.training.dataset import TextDataset from llm.training.trainer import Trainer from shared.configs import ( - TRAIN_TEXTS, BASE_GPT_CONFIG, BPE_CONFIG, - TRAINING_CONFIG, PATHS, TEST_PROMPTS + TRAIN_TEXTS, + BASE_GPT_CONFIG, + BPE_CONFIG, + TRAINING_CONFIG, + PATHS, + TEST_PROMPTS, ) from shared.data import ( - load_training_data, ensure_directories, - print_experiment_info, ExperimentLogger + load_training_data, + ensure_directories, + print_experiment_info, + ExperimentLogger, ) def train_bpe_tokenizer(texts: list, config: dict) -> BPETokenizer: """ Обучает BPE токенизатор на текстах. - + Args: texts: Список текстов для обучения config: Конфигурация токенизатора - + Returns: BPETokenizer: Обученный токенизатор """ print("🔧 Обучение BPE токенизатора...") - + tokenizer = BPETokenizer() tokenizer.train( texts=texts, vocab_size=config["vocab_size"], - special_tokens=config["special_tokens"] + special_tokens=config["special_tokens"], ) - + # Сохраняем токенизатор os.makedirs(os.path.dirname(PATHS["bpe_tokenizer"]), exist_ok=True) tokenizer.save(PATHS["bpe_tokenizer"]) - + print(f"✅ BPE токенизатор обучен и сохранен: {PATHS['bpe_tokenizer']}") print(f"📊 Размер словаря: {tokenizer.get_vocab_size()}") - + return tokenizer def test_tokenizer(tokenizer: BPETokenizer, texts: list): """ Тестирует токенизатор на примерах. - + Args: tokenizer: Обученный токенизатор texts: Список тестовых текстов """ print("\n🧪 Тестирование токенизатора:") - + for i, text in enumerate(texts[:3]): print(f"\nПример {i+1}:") print(f" Исходный текст: '{text}'") - + # Кодирование tokens = tokenizer.encode(text) token_strings = tokenizer.tokenize(text) - + print(f" Токены (ID): {tokens}") print(f" Токены (текст): {token_strings}") print(f" Количество токенов: {len(tokens)}") - + # Декодирование decoded = tokenizer.decode(tokens) print(f" Декодированный: '{decoded}'") - + if text == decoded: print(" ✅ Кодирование/декодирование корректно") else: @@ -95,22 +101,22 @@ def main(): experiment_name = "Обучение GPT2 с BPE токенизатором (только llm)" experiment_config = { "model": "GPT2", - "tokenizer": "BPE", + "tokenizer": "BPE", "vocab_size": BPE_CONFIG["vocab_size"], "training_epochs": TRAINING_CONFIG["num_epochs"], "batch_size": TRAINING_CONFIG["batch_size"], - "learning_rate": TRAINING_CONFIG["learning_rate"] + "learning_rate": TRAINING_CONFIG["learning_rate"], } - + print_experiment_info(experiment_name, experiment_config) ensure_directories() logger = ExperimentLogger(experiment_name) - + try: # === Подготовка данных === train_texts, val_texts = load_training_data() print(f"📊 Данные: {len(train_texts)} train, {len(val_texts)} validation") - + # === Обучение токенизатора === if os.path.exists(PATHS["bpe_tokenizer"]): print("📝 Загрузка предварительно обученного токенизатора...") @@ -118,112 +124,112 @@ def main(): print(f"✅ Токенизатор загружен (vocab_size={tokenizer.get_vocab_size()})") else: tokenizer = train_bpe_tokenizer(TRAIN_TEXTS, BPE_CONFIG) - + # Тестируем токенизатор test_tokenizer(tokenizer, TEST_PROMPTS[:3]) - + # === Инициализация модели === model_config = BASE_GPT_CONFIG.copy() model_config["vocab_size"] = tokenizer.get_vocab_size() - + print(f"\n🔧 Инициализация GPT2 модели...") print(f" Размер словаря: {model_config['vocab_size']}") print(f" Размер эмбеддингов: {model_config['embed_dim']}") print(f" Количество слоев: {model_config['num_layers']}") print(f" Количество голов внимания: {model_config['num_heads']}") - + model = GPT2(model_config) - + # === Подготовка датасета === print(f"\n📊 Подготовка датасета...") train_dataset = TextDataset( - train_texts, - tokenizer, - block_size=model_config["max_position_embeddings"] + train_texts, tokenizer, block_size=model_config["max_position_embeddings"] ) print(f" Размер train датасета: {len(train_dataset)} примеров") - + # === Обучение модели === print(f"\n🎯 Начало обучения GPT2 модели...") - + trainer = Trainer( model=model, train_dataset=train_dataset, lr=TRAINING_CONFIG["learning_rate"], batch_size=TRAINING_CONFIG["batch_size"], num_epochs=TRAINING_CONFIG["num_epochs"], - warmup_steps=TRAINING_CONFIG["warmup_steps"] + warmup_steps=TRAINING_CONFIG["warmup_steps"], ) - + # Запускаем обучение trainer.train() - + # === Сохранение модели === print(f"\n💾 Сохранение модели...") os.makedirs(os.path.dirname(PATHS["gpt_bpe_model"]), exist_ok=True) - + # Сохраняем модель torch.save(model.state_dict(), PATHS["gpt_bpe_model"]) - + # Сохраняем конфигурацию import json - with open(PATHS["gpt_bpe_config"], 'w', encoding='utf-8') as f: + + with open(PATHS["gpt_bpe_config"], "w", encoding="utf-8") as f: json.dump(model_config, f, indent=2, ensure_ascii=False) - + print(f"✅ Модель сохранена:") print(f" - {PATHS['gpt_bpe_model']}: веса модели") print(f" - {PATHS['gpt_bpe_config']}: конфигурация модели") print(f" - {PATHS['bpe_tokenizer']}: токенизатор") - + # === Тестирование генерации === print(f"\n🧪 Тестирование генерации текста...") model.eval() - + for prompt in TEST_PROMPTS[:3]: print(f"\n🔤 Промпт: '{prompt}'") - + try: # Кодируем промпт input_ids = tokenizer.encode(prompt, add_special_tokens=False) input_tensor = torch.tensor([input_ids], dtype=torch.long) - + # Генерируем текст with torch.no_grad(): generated_ids = model.generate( x=input_tensor, max_new_tokens=20, do_sample=True, - temperature=0.8 + temperature=0.8, ) - + # Декодируем результат generated_text = tokenizer.decode(generated_ids[0].tolist()) - generated_part = generated_text[len(prompt):] - + generated_part = generated_text[len(prompt) :] + print(f"🎯 Сгенерировано: '{generated_part}'") print(f"📄 Полный текст: '{generated_text}'") - + except Exception as e: print(f"❌ Ошибка генерации: {e}") - + # === Сохранение результатов === results = { "experiment": experiment_name, "model_config": model_config, "training_config": TRAINING_CONFIG, "tokenizer_vocab_size": tokenizer.get_vocab_size(), - "final_loss": "см. логи обучения" # В реальном эксперименте можно сохранить final loss + "final_loss": "см. логи обучения", # В реальном эксперименте можно сохранить final loss } - + logger.save_logs("checkpoints/llm_only_training_logs.json") - + print(f"\n🎉 Эксперимент завершен успешно!") print(f"\n💡 Для использования обученной модели:") print(f" uv run python experiments/llm_only/generate_gpt_bpe.py") - + except Exception as e: print(f"❌ Ошибка в эксперименте: {e}") import traceback + traceback.print_exc() diff --git a/experiments/llm_only/train_gpt_bpe.py b/experiments/llm_only/train_gpt_bpe.py index 11a69e1..42355f8 100644 --- a/experiments/llm_only/train_gpt_bpe.py +++ b/experiments/llm_only/train_gpt_bpe.py @@ -18,71 +18,77 @@ from llm.training.dataset import TextDataset from llm.training.trainer import Trainer from shared.configs import ( - TRAIN_TEXTS, BASE_GPT_CONFIG, BPE_CONFIG, - TRAINING_CONFIG, PATHS, TEST_PROMPTS + TRAIN_TEXTS, + BASE_GPT_CONFIG, + BPE_CONFIG, + TRAINING_CONFIG, + PATHS, + TEST_PROMPTS, ) from shared.data import ( - load_training_data, ensure_directories, - print_experiment_info, ExperimentLogger + load_training_data, + ensure_directories, + print_experiment_info, + ExperimentLogger, ) def train_bpe_tokenizer(texts: list, config: dict) -> BPETokenizer: """ Обучает BPE токенизатор на текстах. - + Args: texts: Список текстов для обучения config: Конфигурация токенизатора - + Returns: BPETokenizer: Обученный токенизатор """ print("🔧 Обучение BPE токенизатора...") - + tokenizer = BPETokenizer() tokenizer.train( texts=texts, vocab_size=config["vocab_size"], - special_tokens=config["special_tokens"] + special_tokens=config["special_tokens"], ) - + # Сохраняем токенизатор os.makedirs(os.path.dirname(PATHS["bpe_tokenizer"]), exist_ok=True) tokenizer.save(PATHS["bpe_tokenizer"]) - + print(f"✅ BPE токенизатор обучен и сохранен: {PATHS['bpe_tokenizer']}") print(f"📊 Размер словаря: {tokenizer.get_vocab_size()}") - + return tokenizer def test_tokenizer(tokenizer: BPETokenizer, texts: list): """ Тестирует токенизатор на примерах. - + Args: tokenizer: Обученный токенизатор texts: Список тестовых текстов """ print("\n🧪 Тестирование токенизатора:") - + for i, text in enumerate(texts[:3]): print(f"\nПример {i+1}:") print(f" Исходный текст: '{text}'") - + # Кодирование tokens = tokenizer.encode(text) token_strings = tokenizer.tokenize(text) - + print(f" Токены (ID): {tokens}") print(f" Токены (текст): {token_strings}") print(f" Количество токенов: {len(tokens)}") - + # Декодирование decoded = tokenizer.decode(tokens) print(f" Декодированный: '{decoded}'") - + if text == decoded: print(" ✅ Кодирование/декодирование корректно") else: @@ -95,22 +101,22 @@ def main(): experiment_name = "Обучение GPT с BPE токенизатором (только llm)" experiment_config = { "model": "GPT", - "tokenizer": "BPE", + "tokenizer": "BPE", "vocab_size": BPE_CONFIG["vocab_size"], "training_epochs": TRAINING_CONFIG["num_epochs"], "batch_size": TRAINING_CONFIG["batch_size"], - "learning_rate": TRAINING_CONFIG["learning_rate"] + "learning_rate": TRAINING_CONFIG["learning_rate"], } - + print_experiment_info(experiment_name, experiment_config) ensure_directories() logger = ExperimentLogger(experiment_name) - + try: # === Подготовка данных === train_texts, val_texts = load_training_data() print(f"📊 Данные: {len(train_texts)} train, {len(val_texts)} validation") - + # === Обучение токенизатора === if os.path.exists(PATHS["bpe_tokenizer"]): print("📝 Загрузка предварительно обученного токенизатора...") @@ -118,112 +124,112 @@ def main(): print(f"✅ Токенизатор загружен (vocab_size={tokenizer.get_vocab_size()})") else: tokenizer = train_bpe_tokenizer(TRAIN_TEXTS, BPE_CONFIG) - + # Тестируем токенизатор test_tokenizer(tokenizer, TEST_PROMPTS[:3]) - + # === Инициализация модели === model_config = BASE_GPT_CONFIG.copy() model_config["vocab_size"] = tokenizer.get_vocab_size() - + print(f"\n🔧 Инициализация GPT модели...") print(f" Размер словаря: {model_config['vocab_size']}") print(f" Размер эмбеддингов: {model_config['embed_dim']}") print(f" Количество слоев: {model_config['num_layers']}") print(f" Количество голов внимания: {model_config['num_heads']}") - + model = GPT(model_config) - + # === Подготовка датасета === print(f"\n📊 Подготовка датасета...") train_dataset = TextDataset( - train_texts, - tokenizer, - block_size=model_config["max_position_embeddings"] + train_texts, tokenizer, block_size=model_config["max_position_embeddings"] ) print(f" Размер train датасета: {len(train_dataset)} примеров") - + # === Обучение модели === print(f"\n🎯 Начало обучения GPT модели...") - + trainer = Trainer( model=model, train_dataset=train_dataset, lr=TRAINING_CONFIG["learning_rate"], batch_size=TRAINING_CONFIG["batch_size"], num_epochs=TRAINING_CONFIG["num_epochs"], - warmup_steps=TRAINING_CONFIG["warmup_steps"] + warmup_steps=TRAINING_CONFIG["warmup_steps"], ) - + # Запускаем обучение trainer.train() - + # === Сохранение модели === print(f"\n💾 Сохранение модели...") os.makedirs(os.path.dirname(PATHS["gpt_bpe_model"]), exist_ok=True) - + # Сохраняем модель torch.save(model.state_dict(), PATHS["gpt_bpe_model"]) - + # Сохраняем конфигурацию import json - with open(PATHS["gpt_bpe_config"], 'w', encoding='utf-8') as f: + + with open(PATHS["gpt_bpe_config"], "w", encoding="utf-8") as f: json.dump(model_config, f, indent=2, ensure_ascii=False) - + print(f"✅ Модель сохранена:") print(f" - {PATHS['gpt_bpe_model']}: веса модели") print(f" - {PATHS['gpt_bpe_config']}: конфигурация модели") print(f" - {PATHS['bpe_tokenizer']}: токенизатор") - + # === Тестирование генерации === print(f"\n🧪 Тестирование генерации текста...") model.eval() - + for prompt in TEST_PROMPTS[:3]: print(f"\n🔤 Промпт: '{prompt}'") - + try: # Кодируем промпт input_ids = tokenizer.encode(prompt, add_special_tokens=False) input_tensor = torch.tensor([input_ids], dtype=torch.long) - + # Генерируем текст with torch.no_grad(): generated_ids = model.generate( x=input_tensor, max_new_tokens=20, do_sample=True, - temperature=0.8 + temperature=0.8, ) - + # Декодируем результат generated_text = tokenizer.decode(generated_ids[0].tolist()) - generated_part = generated_text[len(prompt):] - + generated_part = generated_text[len(prompt) :] + print(f"🎯 Сгенерировано: '{generated_part}'") print(f"📄 Полный текст: '{generated_text}'") - + except Exception as e: print(f"❌ Ошибка генерации: {e}") - + # === Сохранение результатов === results = { "experiment": experiment_name, "model_config": model_config, "training_config": TRAINING_CONFIG, "tokenizer_vocab_size": tokenizer.get_vocab_size(), - "final_loss": "см. логи обучения" # В реальном эксперименте можно сохранить final loss + "final_loss": "см. логи обучения", # В реальном эксперименте можно сохранить final loss } - + logger.save_logs("checkpoints/llm_only_training_logs.json") - + print(f"\n🎉 Эксперимент завершен успешно!") print(f"\n💡 Для использования обученной модели:") print(f" uv run python experiments/llm_only/generate_gpt_bpe.py") - + except Exception as e: print(f"❌ Ошибка в эксперименте: {e}") import traceback + traceback.print_exc() diff --git a/experiments/llm_only/train_llama_bpe.py b/experiments/llm_only/train_llama_bpe.py index 76e8690..4c609ac 100644 --- a/experiments/llm_only/train_llama_bpe.py +++ b/experiments/llm_only/train_llama_bpe.py @@ -18,71 +18,77 @@ from llm.training.dataset import TextDataset from llm.training.trainer import Trainer from shared.configs import ( - TRAIN_TEXTS, BASE_GPT_CONFIG, BPE_CONFIG, - TRAINING_CONFIG, PATHS, TEST_PROMPTS + TRAIN_TEXTS, + BASE_GPT_CONFIG, + BPE_CONFIG, + TRAINING_CONFIG, + PATHS, + TEST_PROMPTS, ) from shared.data import ( - load_training_data, ensure_directories, - print_experiment_info, ExperimentLogger + load_training_data, + ensure_directories, + print_experiment_info, + ExperimentLogger, ) def train_bpe_tokenizer(texts: list, config: dict) -> BPETokenizer: """ Обучает BPE токенизатор на текстах. - + Args: texts: Список текстов для обучения config: Конфигурация токенизатора - + Returns: BPETokenizer: Обученный токенизатор """ print("🔧 Обучение BPE токенизатора...") - + tokenizer = BPETokenizer() tokenizer.train( texts=texts, vocab_size=config["vocab_size"], - special_tokens=config["special_tokens"] + special_tokens=config["special_tokens"], ) - + # Сохраняем токенизатор os.makedirs(os.path.dirname(PATHS["bpe_tokenizer"]), exist_ok=True) tokenizer.save(PATHS["bpe_tokenizer"]) - + print(f"✅ BPE токенизатор обучен и сохранен: {PATHS['bpe_tokenizer']}") print(f"📊 Размер словаря: {tokenizer.get_vocab_size()}") - + return tokenizer def test_tokenizer(tokenizer: BPETokenizer, texts: list): """ Тестирует токенизатор на примерах. - + Args: tokenizer: Обученный токенизатор texts: Список тестовых текстов """ print("\n🧪 Тестирование токенизатора:") - + for i, text in enumerate(texts[:3]): print(f"\nПример {i+1}:") print(f" Исходный текст: '{text}'") - + # Кодирование tokens = tokenizer.encode(text) token_strings = tokenizer.tokenize(text) - + print(f" Токены (ID): {tokens}") print(f" Токены (текст): {token_strings}") print(f" Количество токенов: {len(tokens)}") - + # Декодирование decoded = tokenizer.decode(tokens) print(f" Декодированный: '{decoded}'") - + if text == decoded: print(" ✅ Кодирование/декодирование корректно") else: @@ -95,22 +101,22 @@ def main(): experiment_name = "Обучение Llama с BPE токенизатором (только llm)" experiment_config = { "model": "Llama", - "tokenizer": "BPE", + "tokenizer": "BPE", "vocab_size": BPE_CONFIG["vocab_size"], "training_epochs": TRAINING_CONFIG["num_epochs"], "batch_size": TRAINING_CONFIG["batch_size"], - "learning_rate": TRAINING_CONFIG["learning_rate"] + "learning_rate": TRAINING_CONFIG["learning_rate"], } - + print_experiment_info(experiment_name, experiment_config) ensure_directories() logger = ExperimentLogger(experiment_name) - + try: # === Подготовка данных === train_texts, val_texts = load_training_data() print(f"📊 Данные: {len(train_texts)} train, {len(val_texts)} validation") - + # === Обучение токенизатора === if os.path.exists(PATHS["bpe_tokenizer"]): print("📝 Загрузка предварительно обученного токенизатора...") @@ -118,112 +124,112 @@ def main(): print(f"✅ Токенизатор загружен (vocab_size={tokenizer.get_vocab_size()})") else: tokenizer = train_bpe_tokenizer(TRAIN_TEXTS, BPE_CONFIG) - + # Тестируем токенизатор test_tokenizer(tokenizer, TEST_PROMPTS[:3]) - + # === Инициализация модели === model_config = BASE_GPT_CONFIG.copy() model_config["vocab_size"] = tokenizer.get_vocab_size() - + print(f"\n🔧 Инициализация Llama модели...") print(f" Размер словаря: {model_config['vocab_size']}") print(f" Размер эмбеддингов: {model_config['embed_dim']}") print(f" Количество слоев: {model_config['num_layers']}") print(f" Количество голов внимания: {model_config['num_heads']}") - + model = Llama(model_config) - + # === Подготовка датасета === print(f"\n📊 Подготовка датасета...") train_dataset = TextDataset( - train_texts, - tokenizer, - block_size=model_config["max_position_embeddings"] + train_texts, tokenizer, block_size=model_config["max_position_embeddings"] ) print(f" Размер train датасета: {len(train_dataset)} примеров") - + # === Обучение модели === print(f"\n🎯 Начало обучения Llama модели...") - + trainer = Trainer( model=model, train_dataset=train_dataset, lr=TRAINING_CONFIG["learning_rate"], batch_size=TRAINING_CONFIG["batch_size"], num_epochs=TRAINING_CONFIG["num_epochs"], - warmup_steps=TRAINING_CONFIG["warmup_steps"] + warmup_steps=TRAINING_CONFIG["warmup_steps"], ) - + # Запускаем обучение trainer.train() - + # === Сохранение модели === print(f"\n💾 Сохранение модели...") os.makedirs(os.path.dirname(PATHS["gpt_bpe_model"]), exist_ok=True) - + # Сохраняем модель torch.save(model.state_dict(), PATHS["gpt_bpe_model"]) - + # Сохраняем конфигурацию import json - with open(PATHS["gpt_bpe_config"], 'w', encoding='utf-8') as f: + + with open(PATHS["gpt_bpe_config"], "w", encoding="utf-8") as f: json.dump(model_config, f, indent=2, ensure_ascii=False) - + print(f"✅ Модель сохранена:") print(f" - {PATHS['gpt_bpe_model']}: веса модели") print(f" - {PATHS['gpt_bpe_config']}: конфигурация модели") print(f" - {PATHS['bpe_tokenizer']}: токенизатор") - + # === Тестирование генерации === print(f"\n🧪 Тестирование генерации текста...") model.eval() - + for prompt in TEST_PROMPTS[:3]: print(f"\n🔤 Промпт: '{prompt}'") - + try: # Кодируем промпт input_ids = tokenizer.encode(prompt, add_special_tokens=False) input_tensor = torch.tensor([input_ids], dtype=torch.long) - + # Генерируем текст with torch.no_grad(): generated_ids = model.generate( x=input_tensor, max_new_tokens=20, do_sample=True, - temperature=0.8 + temperature=0.8, ) - + # Декодируем результат generated_text = tokenizer.decode(generated_ids[0].tolist()) - generated_part = generated_text[len(prompt):] - + generated_part = generated_text[len(prompt) :] + print(f"🎯 Сгенерировано: '{generated_part}'") print(f"📄 Полный текст: '{generated_text}'") - + except Exception as e: print(f"❌ Ошибка генерации: {e}") - + # === Сохранение результатов === results = { "experiment": experiment_name, "model_config": model_config, "training_config": TRAINING_CONFIG, "tokenizer_vocab_size": tokenizer.get_vocab_size(), - "final_loss": "см. логи обучения" # В реальном эксперименте можно сохранить final loss + "final_loss": "см. логи обучения", # В реальном эксперименте можно сохранить final loss } - + logger.save_logs("checkpoints/llm_only_training_logs.json") - + print(f"\n🎉 Эксперимент завершен успешно!") print(f"\n💡 Для использования обученной модели:") print(f" uv run python experiments/llm_only/generate_gpt_bpe.py") - + except Exception as e: print(f"❌ Ошибка в эксперименте: {e}") import traceback + traceback.print_exc() diff --git a/experiments/shared/configs.py b/experiments/shared/configs.py index 1ff6726..85348bc 100644 --- a/experiments/shared/configs.py +++ b/experiments/shared/configs.py @@ -30,7 +30,7 @@ BASE_GPT_CONFIG = { "num_heads": 4, "num_layers": 4, "max_position_embeddings": 128, - "dropout": 0.1 + "dropout": 0.1, } # Конфигурация для маленькой модели (быстрое тестирование) @@ -40,7 +40,7 @@ SMALL_GPT_CONFIG = { "num_heads": 2, "num_layers": 2, "max_position_embeddings": 64, - "dropout": 0.1 + "dropout": 0.1, } # Конфигурация для большой модели (качественное обучение) @@ -50,13 +50,13 @@ LARGE_GPT_CONFIG = { "num_heads": 8, "num_layers": 6, "max_position_embeddings": 256, - "dropout": 0.1 + "dropout": 0.1, } # === Конфигурации токенизатора === BPE_CONFIG = { "vocab_size": 1000, - "special_tokens": ["", "", "", ""] + "special_tokens": ["", "", "", ""], } # === Конфигурации обучения === @@ -65,7 +65,7 @@ TRAINING_CONFIG = { "batch_size": 2, "num_epochs": 3, "warmup_steps": 50, - "gradient_clip": 1.0 + "gradient_clip": 1.0, } # === Конфигурации генерации === @@ -74,7 +74,7 @@ GENERATION_CONFIG = { "temperature": 0.7, "do_sample": True, "top_k": None, - "top_p": None + "top_p": None, } # === Пути для сохранения === @@ -84,7 +84,7 @@ PATHS = { "gpt_bpe_config": "checkpoints/gpt-bpe/config.json", "hf_tokenizer": "checkpoints/hf-bpe-tokenizer", "hf_model": "checkpoints/hf-trained", - "hf_proxy_model": "checkpoints/hf-trained-proxy" + "hf_proxy_model": "checkpoints/hf-trained-proxy", } # === Тестовые промпты === diff --git a/experiments/shared/data.py b/experiments/shared/data.py index 4c18220..0019882 100644 --- a/experiments/shared/data.py +++ b/experiments/shared/data.py @@ -10,17 +10,17 @@ from .configs import TRAIN_TEXTS, PATHS def load_training_data(split_ratio: float = 0.8) -> Tuple[List[str], List[str]]: """ Загружает данные для обучения и разделяет на train/validation. - + Args: split_ratio: Доля данных для обучения - + Returns: Tuple: (train_texts, val_texts) """ train_size = int(len(TRAIN_TEXTS) * split_ratio) train_data = TRAIN_TEXTS[:train_size] val_data = TRAIN_TEXTS[train_size:] - + return train_data, val_data @@ -28,13 +28,13 @@ def ensure_directories(): """Создает необходимые директории если они не существуют.""" directories = [ "checkpoints", - "checkpoints/gpt-bpe", + "checkpoints/gpt-bpe", "checkpoints/hf-bpe-tokenizer", "checkpoints/hf-trained", "checkpoints/hf-trained-proxy", - "logs" + "logs", ] - + for directory in directories: os.makedirs(directory, exist_ok=True) @@ -42,33 +42,34 @@ def ensure_directories(): def get_model_paths(experiment_type: str = "llm_only") -> dict: """ Возвращает пути для конкретного типа эксперимента. - + Args: experiment_type: Тип эксперимента ('llm_only' или 'hf_integration') - + Returns: dict: Словарь с путями """ base_paths = PATHS.copy() - + if experiment_type == "hf_integration": - base_paths.update({ - "model": base_paths["hf_model"], - "tokenizer": base_paths["hf_tokenizer"] - }) + base_paths.update( + {"model": base_paths["hf_model"], "tokenizer": base_paths["hf_tokenizer"]} + ) else: # llm_only - base_paths.update({ - "model": base_paths["gpt_bpe_model"], - "tokenizer": base_paths["bpe_tokenizer"] - }) - + base_paths.update( + { + "model": base_paths["gpt_bpe_model"], + "tokenizer": base_paths["bpe_tokenizer"], + } + ) + return base_paths def print_experiment_info(experiment_name: str, config: dict): """ Выводит информацию о запускаемом эксперименте. - + Args: experiment_name: Название эксперимента config: Конфигурация эксперимента @@ -85,35 +86,35 @@ def print_experiment_info(experiment_name: str, config: dict): def save_experiment_results(results: dict, filepath: str): """ Сохраняет результаты эксперимента в файл. - + Args: results: Словарь с результатами filepath: Путь для сохранения """ import json - - with open(filepath, 'w', encoding='utf-8') as f: + + with open(filepath, "w", encoding="utf-8") as f: json.dump(results, f, ensure_ascii=False, indent=2) - + print(f"✅ Результаты эксперимента сохранены: {filepath}") def load_experiment_results(filepath: str) -> dict: """ Загружает результаты эксперимента из файла. - + Args: filepath: Путь к файлу с результатами - + Returns: dict: Загруженные результаты """ import json - + if not os.path.exists(filepath): return {} - - with open(filepath, 'r', encoding='utf-8') as f: + + with open(filepath, "r", encoding="utf-8") as f: return json.load(f) @@ -121,42 +122,39 @@ class ExperimentLogger: """ Логгер для экспериментов. """ - + def __init__(self, experiment_name: str): self.experiment_name = experiment_name self.metrics = {} - + def log_metric(self, name: str, value: float): """Логирует метрику.""" if name not in self.metrics: self.metrics[name] = [] self.metrics[name].append(value) print(f"📈 {name}: {value:.4f}") - + def log_step(self, step: int, loss: float, **kwargs): """Логирует шаг обучения.""" print(f"📊 Step {step}: loss={loss:.4f}", end="") for key, value in kwargs.items(): print(f", {key}={value:.4f}", end="") print() - + def log_epoch(self, epoch: int, train_loss: float, val_loss: float = None): """Логирует завершение эпохи.""" print(f"🎯 Epoch {epoch}: train_loss={train_loss:.4f}", end="") if val_loss is not None: print(f", val_loss={val_loss:.4f}", end="") print() - + def save_logs(self, filepath: str): """Сохраняет логи эксперимента.""" import json - - logs = { - "experiment_name": self.experiment_name, - "metrics": self.metrics - } - - with open(filepath, 'w', encoding='utf-8') as f: + + logs = {"experiment_name": self.experiment_name, "metrics": self.metrics} + + with open(filepath, "w", encoding="utf-8") as f: json.dump(logs, f, ensure_ascii=False, indent=2) - + print(f"✅ Логи эксперимента сохранены: {filepath}") diff --git a/hf-proxy/src/hf_proxy/__init__.py b/hf-proxy/src/hf_proxy/__init__.py index 9dcc511..628eb9d 100644 --- a/hf-proxy/src/hf_proxy/__init__.py +++ b/hf-proxy/src/hf_proxy/__init__.py @@ -27,16 +27,13 @@ __all__ = [ # Основные классы адаптера "HFAdapter", "HFGPTAdapter", - # Конфигурации - "HFAdapterConfig", + "HFAdapterConfig", "HFPretrainedConfig", - # Адаптеры токенизаторов "HFTokenizerAdapter", - "create_hf_tokenizer", + "create_hf_tokenizer", "convert_to_hf_format", - # Утилиты "HFUtils", "TokenizerWrapper", diff --git a/hf-proxy/src/hf_proxy/hf_adapter.py b/hf-proxy/src/hf_proxy/hf_adapter.py index 3b12124..0b0689a 100644 --- a/hf-proxy/src/hf_proxy/hf_adapter.py +++ b/hf-proxy/src/hf_proxy/hf_adapter.py @@ -6,12 +6,12 @@ import torch import torch.nn as nn from typing import Optional, Tuple, Union, List from transformers import ( - PreTrainedModel, + PreTrainedModel, GPT2LMHeadModel, GPT2Config, GenerationConfig, LogitsProcessorList, - StoppingCriteriaList + StoppingCriteriaList, ) from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions @@ -24,38 +24,39 @@ class HFGPTAdapter(PreTrainedModel): Адаптер для модели GPT из библиотеки llm. Позволяет использовать кастомные GPT модели с HuggingFace Transformers. """ + config_class = HFPretrainedConfig - + def __init__(self, config: HFPretrainedConfig, llm_model: Optional[GPT] = None): """ Инициализация адаптера. - + Args: config: Конфигурация HuggingFace llm_model: Опционально, предварительно созданная модель llm """ super().__init__(config) - + # Преобразуем HF конфигурацию в формат llm llm_config = self._hf_to_llm_config(config) - + # Создаем или используем переданную модель if llm_model is None: self.llm_model = GPT(llm_config) else: self.llm_model = llm_model - + # Устанавливаем веса если они есть в конфигурации - if hasattr(config, 'state_dict') and config.state_dict is not None: + if hasattr(config, "state_dict") and config.state_dict is not None: self.llm_model.load_state_dict(config.state_dict) - + def _hf_to_llm_config(self, hf_config: HFPretrainedConfig) -> dict: """ Преобразует конфигурацию HF в формат llm. - + Args: hf_config: Конфигурация HuggingFace - + Returns: dict: Конфигурация для llm модели """ @@ -67,7 +68,7 @@ class HFGPTAdapter(PreTrainedModel): "max_position_embeddings": hf_config.max_position_embeddings, "dropout": hf_config.hidden_dropout_prob, } - + def forward( self, input_ids: Optional[torch.Tensor] = None, @@ -78,11 +79,11 @@ class HFGPTAdapter(PreTrainedModel): output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, - **kwargs + **kwargs, ) -> Union[Tuple, CausalLMOutputWithCrossAttentions]: """ Прямой проход модели. - + Args: input_ids: Входные токены [batch_size, seq_len] attention_mask: Маска внимания [batch_size, seq_len] @@ -92,38 +93,39 @@ class HFGPTAdapter(PreTrainedModel): output_attentions: Возвращать веса внимания output_hidden_states: Возвращать скрытые состояния return_dict: Возвращать словарь вместо кортежа - + Returns: CausalLMOutputWithCrossAttentions или кортеж """ - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + # Основной forward pass outputs = self.llm_model(input_ids) if isinstance(outputs, tuple): logits = outputs[0] else: logits = outputs - + loss = None if labels is not None: # Сдвигаем логиты и метки для языкового моделирования shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() - + # Вычисляем cross-entropy loss loss_fct = nn.CrossEntropyLoss() loss = loss_fct( - shift_logits.view(-1, shift_logits.size(-1)), - shift_labels.view(-1) + shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1) ) - + if not return_dict: output = (logits,) if loss is not None: output = (loss,) + output return output - + return CausalLMOutputWithCrossAttentions( loss=loss, logits=logits, @@ -132,30 +134,27 @@ class HFGPTAdapter(PreTrainedModel): attentions=None, cross_attentions=None, ) - + def prepare_inputs_for_generation( - self, - input_ids: torch.Tensor, - past_key_values: Optional[Tuple] = None, - **kwargs + self, input_ids: torch.Tensor, past_key_values: Optional[Tuple] = None, **kwargs ) -> dict: """ Подготавливает входные данные для генерации. - + Args: input_ids: Входные токены past_key_values: Кешированные ключи и значения - + Returns: dict: Подготовленные входные данные """ # Наша простая реализация пока не поддерживает past_key_values return {"input_ids": input_ids} - + def can_generate(self) -> bool: """Проверяет, может ли модель генерировать текст.""" return True - + def generate( self, input_ids: Optional[torch.Tensor] = None, @@ -163,32 +162,32 @@ class HFGPTAdapter(PreTrainedModel): generation_config: Optional[GenerationConfig] = None, logits_processor: Optional[LogitsProcessorList] = None, stopping_criteria: Optional[StoppingCriteriaList] = None, - **kwargs + **kwargs, ) -> torch.Tensor: """ Генерация текста с поддержкой HuggingFace интерфейса. - + Args: input_ids: Входные токены attention_mask: Маска внимания generation_config: Конфигурация генерации logits_processor: Процессоры логитов stopping_criteria: Критерии остановки - + Returns: torch.Tensor: Сгенерированные токены """ # Извлекаем обязательные параметры из kwargs или используем значения по умолчанию - max_new_tokens = kwargs.pop('max_new_tokens', 50) - do_sample = kwargs.pop('do_sample', True) - + max_new_tokens = kwargs.pop("max_new_tokens", 50) + do_sample = kwargs.pop("do_sample", True) + # Используем встроенную генерацию llm модели return self.llm_model.generate( x=input_ids, max_new_tokens=max_new_tokens, do_sample=do_sample, attention_mask=attention_mask, - **kwargs + **kwargs, ) @@ -196,64 +195,66 @@ class HFAdapter: """ Основной класс адаптера для преобразования моделей llm в формат HuggingFace. """ - + @staticmethod def from_llm_model( - llm_model: GPT, - hf_config: Optional[HFAdapterConfig] = None + llm_model: GPT, hf_config: Optional[HFAdapterConfig] = None ) -> HFGPTAdapter: """ Создает адаптер из существующей llm модели. - + Args: llm_model: Обученная модель из библиотеки llm hf_config: Конфигурация для HuggingFace - + Returns: HFGPTAdapter: Адаптированная модель """ if hf_config is None: # Создаем конфигурацию из модели llm hf_config = HFAdapterConfig.from_llm_config(llm_model.config) - + # Преобразуем в PretrainedConfig pretrained_config = HFPretrainedConfig(**hf_config.to_dict()) - + return HFGPTAdapter(pretrained_config, llm_model) - + @staticmethod def from_pretrained( - model_path: str, - hf_config: Optional[HFAdapterConfig] = None + model_path: str, hf_config: Optional[HFAdapterConfig] = None ) -> HFGPTAdapter: """ Загружает модель из чекпоинта и создает адаптер. - + Args: model_path: Путь к сохраненной модели hf_config: Конфигурация для HuggingFace - + Returns: HFGPTAdapter: Адаптированная модель """ # Загружаем состояние модели - state_dict = torch.load(model_path, map_location='cpu') - + state_dict = torch.load(model_path, map_location="cpu") + # Определяем конфигурацию из состояния модели или используем переданную if hf_config is None: # Пытаемся определить конфигурацию из состояния модели # Это упрощенный подход - в реальности нужно сохранять конфигурацию отдельно - vocab_size = state_dict.get('_token_embeddings._embedding.weight', torch.zeros(50257, 768)).shape[0] - embed_dim = state_dict.get('_token_embeddings._embedding.weight', torch.zeros(50257, 768)).shape[1] - + vocab_size = state_dict.get( + "_token_embeddings._embedding.weight", torch.zeros(50257, 768) + ).shape[0] + embed_dim = state_dict.get( + "_token_embeddings._embedding.weight", torch.zeros(50257, 768) + ).shape[1] + hf_config = HFAdapterConfig( vocab_size=vocab_size, hidden_size=embed_dim, # Остальные параметры можно установить по умолчанию ) - + pretrained_config = HFPretrainedConfig(**hf_config.to_dict()) - + # Создаем модель llm и загружаем веса llm_config = { "vocab_size": hf_config.vocab_size, @@ -263,21 +264,17 @@ class HFAdapter: "max_position_embeddings": hf_config.max_position_embeddings, "dropout": hf_config.hidden_dropout_prob, } - + llm_model = GPT(llm_config) llm_model.load_state_dict(state_dict) - + return HFGPTAdapter(pretrained_config, llm_model) - + @staticmethod - def save_pretrained( - model: HFGPTAdapter, - save_directory: str, - **kwargs - ): + def save_pretrained(model: HFGPTAdapter, save_directory: str, **kwargs): """ Сохраняет адаптированную модель в формате HuggingFace. - + Args: model: Адаптированная модель save_directory: Директория для сохранения @@ -285,19 +282,19 @@ class HFAdapter: """ import os import json - + # Создаем директорию если не существует os.makedirs(save_directory, exist_ok=True) - + # Сохраняем конфигурацию config_path = os.path.join(save_directory, "config.json") - with open(config_path, 'w', encoding='utf-8') as f: + with open(config_path, "w", encoding="utf-8") as f: json.dump(model.config.to_dict(), f, indent=2, ensure_ascii=False) - + # Сохраняем веса модели model_path = os.path.join(save_directory, "pytorch_model.bin") torch.save(model.llm_model.state_dict(), model_path) - + # Сохраняем токенизатор если передан - if hasattr(kwargs, 'tokenizer') and kwargs['tokenizer'] is not None: - kwargs['tokenizer'].save_pretrained(save_directory) + if hasattr(kwargs, "tokenizer") and kwargs["tokenizer"] is not None: + kwargs["tokenizer"].save_pretrained(save_directory) diff --git a/hf-proxy/src/hf_proxy/hf_config.py b/hf-proxy/src/hf_proxy/hf_config.py index 2cd206d..3088126 100644 --- a/hf-proxy/src/hf_proxy/hf_config.py +++ b/hf-proxy/src/hf_proxy/hf_config.py @@ -6,11 +6,12 @@ from dataclasses import dataclass, field from typing import Dict, Any, Optional from transformers import PretrainedConfig + @dataclass class HFAdapterConfig: """ Конфигурация для адаптера HuggingFace. - + Параметры: model_type: Тип модели (gpt, llama, etc.) vocab_size: Размер словаря @@ -28,6 +29,7 @@ class HFAdapterConfig: eos_token_id: ID токена конца строки bos_token_id: ID токена начала строки """ + model_type: str = "gpt" vocab_size: int = 50257 hidden_size: int = 768 @@ -43,49 +45,50 @@ class HFAdapterConfig: pad_token_id: int = 50256 eos_token_id: int = 50256 bos_token_id: int = 50256 - + # Дополнительные параметры для совместимости architectures: list = field(default_factory=lambda: ["GPT2LMHeadModel"]) torch_dtype: str = "float32" transformers_version: str = "4.44.0" - + def to_dict(self) -> Dict[str, Any]: """Преобразует конфигурацию в словарь.""" return { - k: v for k, v in self.__dict__.items() - if not k.startswith('_') and not callable(v) + k: v + for k, v in self.__dict__.items() + if not k.startswith("_") and not callable(v) } - + @classmethod def from_llm_config(cls, llm_config: Dict[str, Any]) -> "HFAdapterConfig": """ Создает конфигурацию HF из конфигурации llm. - + Args: llm_config: Конфигурация модели из библиотеки llm - + Returns: HFAdapterConfig: Конфигурация для HuggingFace """ # Маппинг параметров из llm в HF формат mapping = { "embed_dim": "hidden_size", - "num_layers": "num_hidden_layers", + "num_layers": "num_hidden_layers", "num_heads": "num_attention_heads", "max_position_embeddings": "max_position_embeddings", "dropout": "hidden_dropout_prob", - "vocab_size": "vocab_size" + "vocab_size": "vocab_size", } - + hf_config_dict = {} for llm_key, hf_key in mapping.items(): if llm_key in llm_config: hf_config_dict[hf_key] = llm_config[llm_key] - + # Устанавливаем промежуточный размер (обычно 4x hidden_size) if "hidden_size" in hf_config_dict: hf_config_dict["intermediate_size"] = hf_config_dict["hidden_size"] * 4 - + return cls(**hf_config_dict) @@ -94,8 +97,9 @@ class HFPretrainedConfig(PretrainedConfig): Конфигурация для предобученных моделей HuggingFace. Наследуется от PretrainedConfig для полной совместимости. """ + model_type = "gpt" - + def __init__( self, vocab_size=50257, @@ -112,15 +116,15 @@ class HFPretrainedConfig(PretrainedConfig): pad_token_id=50256, eos_token_id=50256, bos_token_id=50256, - **kwargs + **kwargs, ): super().__init__( pad_token_id=pad_token_id, eos_token_id=eos_token_id, bos_token_id=bos_token_id, - **kwargs + **kwargs, ) - + self.vocab_size = vocab_size self.hidden_size = hidden_size self.num_hidden_layers = num_hidden_layers diff --git a/hf-proxy/src/hf_proxy/hf_tokenizer.py b/hf-proxy/src/hf_proxy/hf_tokenizer.py index 34d66aa..106d87f 100644 --- a/hf-proxy/src/hf_proxy/hf_tokenizer.py +++ b/hf-proxy/src/hf_proxy/hf_tokenizer.py @@ -12,84 +12,82 @@ class HFTokenizerAdapter: Упрощенный адаптер для кастомных токенизаторов llm. Предоставляет совместимый с HuggingFace интерфейс. """ - + def __init__(self, llm_tokenizer: BaseTokenizer): """ Инициализация адаптера. - + Args: llm_tokenizer: Кастомный токенизатор из llm """ self.llm_tokenizer = llm_tokenizer - + # Получаем словарь и размер self._vocab = llm_tokenizer.get_vocab() self.vocab_size = llm_tokenizer.get_vocab_size() - + # Устанавливаем специальные токены - self.pad_token = getattr(llm_tokenizer, 'pad_token', '') - self.unk_token = getattr(llm_tokenizer, 'unk_token', '') - self.bos_token = getattr(llm_tokenizer, 'bos_token', '') - self.eos_token = getattr(llm_tokenizer, 'eos_token', '') - + self.pad_token = getattr(llm_tokenizer, "pad_token", "") + self.unk_token = getattr(llm_tokenizer, "unk_token", "") + self.bos_token = getattr(llm_tokenizer, "bos_token", "") + self.eos_token = getattr(llm_tokenizer, "eos_token", "") + # Сохраняем ID специальных токенов - self.pad_token_id = getattr(llm_tokenizer, 'pad_token_id', 0) - self.unk_token_id = getattr(llm_tokenizer, 'unk_token_id', 1) - self.bos_token_id = getattr(llm_tokenizer, 'bos_token_id', 2) - self.eos_token_id = getattr(llm_tokenizer, 'eos_token_id', 3) - + self.pad_token_id = getattr(llm_tokenizer, "pad_token_id", 0) + self.unk_token_id = getattr(llm_tokenizer, "unk_token_id", 1) + self.bos_token_id = getattr(llm_tokenizer, "bos_token_id", 2) + self.eos_token_id = getattr(llm_tokenizer, "eos_token_id", 3) + def __call__(self, text: str, **kwargs): """ Вызов токенизатора с параметрами как у HuggingFace. - + Args: text: Входной текст **kwargs: Параметры токенизации - + Returns: dict: Словарь с токенами """ - return_tensors = kwargs.get('return_tensors', None) - padding = kwargs.get('padding', False) - truncation = kwargs.get('truncation', False) - max_length = kwargs.get('max_length', None) - add_special_tokens = kwargs.get('add_special_tokens', True) - + return_tensors = kwargs.get("return_tensors", None) + padding = kwargs.get("padding", False) + truncation = kwargs.get("truncation", False) + max_length = kwargs.get("max_length", None) + add_special_tokens = kwargs.get("add_special_tokens", True) + # Кодируем текст - #input_ids = self.llm_tokenizer.encode( - # text, + # input_ids = self.llm_tokenizer.encode( + # text, # add_special_tokens=add_special_tokens - #) + # ) if isinstance(text, str): input_ids = self.llm_tokenizer.encode( - text, - add_special_tokens=add_special_tokens + text, add_special_tokens=add_special_tokens ) input_ids = [input_ids] # <-- оборачиваем в batch else: # Список строк, батч-режим! input_ids = [ - self.llm_tokenizer.encode( - t, - add_special_tokens=add_special_tokens - ) for t in text + self.llm_tokenizer.encode(t, add_special_tokens=add_special_tokens) + for t in text ] - + # Применяем truncation if truncation and max_length is not None and len(input_ids) > max_length: input_ids = input_ids[:max_length] - + # Применяем padding if padding and max_length is not None and len(input_ids) < max_length: input_ids = input_ids + [self.pad_token_id] * (max_length - len(input_ids)) - + # Конвертируем в тензоры если нужно if return_tensors == "pt": import torch + input_ids = torch.tensor([input_ids]) - + return {"input_ids": input_ids} - + def encode( self, text: str, @@ -99,11 +97,11 @@ class HFTokenizerAdapter: truncation: bool = False, max_length: Optional[int] = None, return_tensors: Optional[str] = None, - **kwargs + **kwargs, ) -> Union[List[int], List[List[int]]]: """ Кодирует текст в последовательность токенов. - + Args: text: Входной текст text_pair: Второй текст (для парных задач) @@ -112,84 +110,91 @@ class HFTokenizerAdapter: truncation: Обрезать последовательность max_length: Максимальная длина return_tensors: Возвращать тензоры - + Returns: Список токенов или список списков токенов """ # Кодируем основной текст token_ids = self.llm_tokenizer.encode( - text, - add_special_tokens=add_special_tokens + text, add_special_tokens=add_special_tokens ) - + # Обрабатываем text_pair если есть if text_pair is not None: - pair_ids = self.llm_tokenizer.encode( - text_pair, - add_special_tokens=False - ) + pair_ids = self.llm_tokenizer.encode(text_pair, add_special_tokens=False) token_ids.extend(pair_ids) - + # Применяем truncation if truncation and max_length is not None and len(token_ids) > max_length: token_ids = token_ids[:max_length] - + # Применяем padding if padding and max_length is not None and len(token_ids) < max_length: token_ids = token_ids + [self.pad_token_id] * (max_length - len(token_ids)) - + # Конвертируем в тензоры если нужно if return_tensors == "pt": import torch + return torch.tensor([token_ids]) elif return_tensors == "np": import numpy as np + return np.array([token_ids]) - + return token_ids - + def decode( self, token_ids: Union[int, List[int], List[List[int]]], skip_special_tokens: bool = True, - **kwargs + **kwargs, ) -> str: """ Декодирует последовательность токенов в текст. - + Args: token_ids: ID токенов skip_special_tokens: Пропускать специальные токены - + Returns: str: Декодированный текст """ # Обрабатываем разные форматы входных данных if isinstance(token_ids, int): token_ids = [token_ids] - elif isinstance(token_ids, list) and len(token_ids) > 0 and isinstance(token_ids[0], list): + elif ( + isinstance(token_ids, list) + and len(token_ids) > 0 + and isinstance(token_ids[0], list) + ): # Список списков - берем первый элемент token_ids = token_ids[0] - + # Фильтруем специальные токены если нужно if skip_special_tokens: - special_ids = {self.pad_token_id, self.unk_token_id, self.bos_token_id, self.eos_token_id} + special_ids = { + self.pad_token_id, + self.unk_token_id, + self.bos_token_id, + self.eos_token_id, + } token_ids = [tid for tid in token_ids if tid not in special_ids] - + return self.llm_tokenizer.decode(token_ids) - + def tokenize(self, text: str, **kwargs) -> List[str]: """ Токенизирует текст в список строковых токенов. - + Args: text: Входной текст - + Returns: List[str]: Список токенов """ return self.llm_tokenizer.tokenize(text) - + def pad( self, encoded_inputs, @@ -202,7 +207,7 @@ class HFTokenizerAdapter: ): """ Pad a list of encoded inputs. - + Args: encoded_inputs: List of encoded inputs padding: Padding strategy @@ -211,7 +216,7 @@ class HFTokenizerAdapter: return_attention_mask: Return attention mask return_tensors: Return tensors verbose: Verbose mode - + Returns: Padded inputs """ @@ -224,47 +229,62 @@ class HFTokenizerAdapter: # Обрабатываем разные типы данных if isinstance(input_ids, int): seq_len = 1 - elif hasattr(input_ids, 'shape'): - seq_len = input_ids.shape[-1] if len(input_ids.shape) > 1 else len(input_ids) + elif hasattr(input_ids, "shape"): + seq_len = ( + input_ids.shape[-1] + if len(input_ids.shape) > 1 + else len(input_ids) + ) else: seq_len = len(input_ids) max_len = max(max_len, seq_len) - + if max_length is not None: max_len = min(max_len, max_length) - + # Применяем padding for item in encoded_inputs: input_ids = item["input_ids"] - + # Получаем текущую длину if isinstance(input_ids, int): current_len = 1 - elif hasattr(input_ids, 'shape'): - current_len = input_ids.shape[-1] if len(input_ids.shape) > 1 else len(input_ids) + elif hasattr(input_ids, "shape"): + current_len = ( + input_ids.shape[-1] + if len(input_ids.shape) > 1 + else len(input_ids) + ) else: current_len = len(input_ids) - + if current_len < max_len: # Дополняем pad_token_id padding_length = max_len - current_len - + # Обрабатываем разные типы данных if isinstance(input_ids, int): - item["input_ids"] = [input_ids] + [self.pad_token_id] * padding_length - elif hasattr(input_ids, 'shape'): + item["input_ids"] = [input_ids] + [ + self.pad_token_id + ] * padding_length + elif hasattr(input_ids, "shape"): import torch - padding_tensor = torch.full((padding_length,), self.pad_token_id, dtype=input_ids.dtype) + + padding_tensor = torch.full( + (padding_length,), self.pad_token_id, dtype=input_ids.dtype + ) item["input_ids"] = torch.cat([input_ids, padding_tensor]) else: - item["input_ids"] = input_ids + [self.pad_token_id] * padding_length - + item["input_ids"] = ( + input_ids + [self.pad_token_id] * padding_length + ) + # Добавляем attention_mask если требуется if "attention_mask" in item: mask = item["attention_mask"] if isinstance(mask, int): item["attention_mask"] = [mask] + [0] * padding_length - elif hasattr(mask, 'shape'): + elif hasattr(mask, "shape"): padding_mask = torch.zeros(padding_length, dtype=mask.dtype) item["attention_mask"] = torch.cat([mask, padding_mask]) else: @@ -272,44 +292,49 @@ class HFTokenizerAdapter: elif return_attention_mask: if isinstance(input_ids, int): item["attention_mask"] = [1] + [0] * padding_length - elif hasattr(input_ids, 'shape'): + elif hasattr(input_ids, "shape"): attention_mask = torch.ones(current_len, dtype=torch.long) padding_mask = torch.zeros(padding_length, dtype=torch.long) - item["attention_mask"] = torch.cat([attention_mask, padding_mask]) + item["attention_mask"] = torch.cat( + [attention_mask, padding_mask] + ) else: - item["attention_mask"] = [1] * current_len + [0] * padding_length - + item["attention_mask"] = [1] * current_len + [ + 0 + ] * padding_length + # Конвертируем в тензоры если требуется if return_tensors == "pt": import torch + for key in list(encoded_inputs[0].keys()): if isinstance(encoded_inputs[0][key], list): for i in range(len(encoded_inputs)): encoded_inputs[i][key] = torch.tensor(encoded_inputs[i][key]) - + return encoded_inputs - + def get_vocab(self) -> Dict[str, int]: """Возвращает словарь токенизатора.""" return self._vocab - + def __len__(self) -> int: """Возвращает размер словаря.""" return self.vocab_size - + def save_pretrained(self, save_directory: str, **kwargs): """ Сохраняет токенизатор в формате HuggingFace. - + Args: save_directory: Директория для сохранения **kwargs: Дополнительные параметры """ import os - + # Создаем директорию если не существует os.makedirs(save_directory, exist_ok=True) - + # Сохраняем конфигурацию токенизатора tokenizer_config = { "tokenizer_class": self.__class__.__name__, @@ -324,77 +349,81 @@ class HFTokenizerAdapter: "bos_token_id": self.bos_token_id, "eos_token_id": self.eos_token_id, } - + config_path = os.path.join(save_directory, "tokenizer_config.json") - with open(config_path, 'w', encoding='utf-8') as f: + with open(config_path, "w", encoding="utf-8") as f: json.dump(tokenizer_config, f, ensure_ascii=False, indent=2) - + # Сохраняем словарь vocab_path = os.path.join(save_directory, "vocab.json") - with open(vocab_path, 'w', encoding='utf-8') as f: + with open(vocab_path, "w", encoding="utf-8") as f: json.dump(self._vocab, f, ensure_ascii=False, indent=2) - + print(f"✅ Токенизатор сохранен в {save_directory}") - + @classmethod def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs): """ Загружает адаптированный токенизатор. - + Args: pretrained_model_name_or_path: Путь к сохраненному токенизатору **kwargs: Дополнительные параметры - + Returns: HFTokenizerAdapter: Загруженный адаптер """ import os - + # Проверяем, является ли путь директорией с файлами токенизатора if os.path.isdir(pretrained_model_name_or_path): # Загружаем из директории - config_path = os.path.join(pretrained_model_name_or_path, "tokenizer_config.json") + config_path = os.path.join( + pretrained_model_name_or_path, "tokenizer_config.json" + ) vocab_path = os.path.join(pretrained_model_name_or_path, "vocab.json") - + if not os.path.exists(config_path) or not os.path.exists(vocab_path): raise FileNotFoundError( f"Файлы токенизатора не найдены в {pretrained_model_name_or_path}" ) - + # Загружаем конфигурацию - with open(config_path, 'r', encoding='utf-8') as f: + with open(config_path, "r", encoding="utf-8") as f: config = json.load(f) - + # Определяем тип токенизатора llm llm_tokenizer_type = config.get("llm_tokenizer_type", "BPETokenizer") - + if llm_tokenizer_type == "BPETokenizer": # Создаем BPETokenizer и загружаем словарь llm_tokenizer = BPETokenizer() - + # Загружаем словарь - with open(vocab_path, 'r', encoding='utf-8') as f: + with open(vocab_path, "r", encoding="utf-8") as f: vocab = json.load(f) - + llm_tokenizer.vocab = vocab llm_tokenizer.inverse_vocab = {v: k for k, v in vocab.items()} llm_tokenizer.vocab_size = len(vocab) - + # Устанавливаем специальные токены llm_tokenizer.pad_token = config.get("pad_token", "") llm_tokenizer.unk_token = config.get("unk_token", "") llm_tokenizer.bos_token = config.get("bos_token", "") llm_tokenizer.eos_token = config.get("eos_token", "") - + llm_tokenizer.pad_token_id = config.get("pad_token_id", 0) llm_tokenizer.unk_token_id = config.get("unk_token_id", 1) llm_tokenizer.bos_token_id = config.get("bos_token_id", 2) llm_tokenizer.eos_token_id = config.get("eos_token_id", 3) - + return cls(llm_tokenizer, **kwargs) else: - raise ValueError(f"Неподдерживаемый тип токенизатора: {llm_tokenizer_type}") - + raise ValueError( + f"Неподдерживаемый тип токенизатора: {llm_tokenizer_type}" + ) + else: # Пытаемся загрузить как файл llm токенизатора try: @@ -409,10 +438,10 @@ class HFTokenizerAdapter: def create_hf_tokenizer(llm_tokenizer: BaseTokenizer) -> HFTokenizerAdapter: """ Создает адаптер HuggingFace для кастомного токенизатора. - + Args: llm_tokenizer: Токенизатор из библиотеки llm - + Returns: HFTokenizerAdapter: Адаптированный токенизатор """ @@ -422,7 +451,7 @@ def create_hf_tokenizer(llm_tokenizer: BaseTokenizer) -> HFTokenizerAdapter: def convert_to_hf_format(llm_tokenizer: BaseTokenizer, save_directory: str): """ Конвертирует кастомный токенизатор в формат HuggingFace. - + Args: llm_tokenizer: Токенизатор из llm save_directory: Директория для сохранения diff --git a/hf-proxy/src/hf_proxy/hf_utils.py b/hf-proxy/src/hf_proxy/hf_utils.py index f956d36..85bb958 100644 --- a/hf-proxy/src/hf_proxy/hf_utils.py +++ b/hf-proxy/src/hf_proxy/hf_utils.py @@ -14,55 +14,57 @@ class HFUtils: """ Утилиты для работы с HuggingFace адаптером. """ - + @staticmethod def create_hf_config_from_llm(llm_config: Dict[str, Any]) -> HFPretrainedConfig: """ Создает конфигурацию HuggingFace из конфигурации llm. - + Args: llm_config: Конфигурация модели из библиотеки llm - + Returns: HFPretrainedConfig: Конфигурация для HuggingFace """ adapter_config = HFAdapterConfig.from_llm_config(llm_config) return HFPretrainedConfig(**adapter_config.to_dict()) - + @staticmethod def convert_to_hf_format( - llm_model, - tokenizer = None, - model_name: str = "custom-gpt" + llm_model, tokenizer=None, model_name: str = "custom-gpt" ) -> tuple: """ Конвертирует llm модель в формат HuggingFace. - + Args: llm_model: Модель из библиотеки llm tokenizer: Токенизатор (HF или кастомный) model_name: Имя модели для сохранения - + Returns: tuple: (адаптированная модель, токенизатор) """ # Создаем адаптер hf_model = HFAdapter.from_llm_model(llm_model) - + # Если токенизатор не передан, создаем стандартный if tokenizer is None: from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained("gpt2") # Устанавливаем специальные токены if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token - elif hasattr(tokenizer, '__class__') and 'BPETokenizer' in str(tokenizer.__class__): + elif hasattr(tokenizer, "__class__") and "BPETokenizer" in str( + tokenizer.__class__ + ): # Если передан наш кастомный токенизатор, создаем адаптер from .hf_tokenizer import create_hf_tokenizer + tokenizer = create_hf_tokenizer(tokenizer) - + return hf_model, tokenizer - + @staticmethod def push_to_hub( model: HFGPTAdapter, @@ -70,11 +72,11 @@ class HFUtils: repo_name: str, organization: Optional[str] = None, private: bool = False, - **kwargs + **kwargs, ): """ Загружает модель в HuggingFace Hub. - + Args: model: Адаптированная модель tokenizer: Токенизатор @@ -85,23 +87,23 @@ class HFUtils: """ try: from huggingface_hub import HfApi, ModelCard, create_repo - + # Создаем репозиторий if organization: repo_id = f"{organization}/{repo_name}" else: repo_id = repo_name - + create_repo(repo_id, private=private, exist_ok=True) - + # Сохраняем модель локально import tempfile import os - + with tempfile.TemporaryDirectory() as tmp_dir: # Сохраняем модель HFAdapter.save_pretrained(model, tmp_dir, tokenizer=tokenizer) - + # Создаем Model Card card = ModelCard.from_template( model_name=repo_name, @@ -110,46 +112,43 @@ class HFUtils: tags=["llm", "gpt", "custom"], ) card.save(os.path.join(tmp_dir, "README.md")) - + # Загружаем в Hub api = HfApi() api.upload_folder( folder_path=tmp_dir, repo_id=repo_id, - commit_message="Initial commit with custom GPT model" + commit_message="Initial commit with custom GPT model", ) - + print(f"✅ Модель успешно загружена в HuggingFace Hub: {repo_id}") - + except ImportError: raise ImportError( "Для загрузки в HuggingFace Hub установите huggingface_hub: " "pip install huggingface_hub" ) - + @staticmethod - def load_from_hub( - repo_id: str, - **kwargs - ) -> tuple: + def load_from_hub(repo_id: str, **kwargs) -> tuple: """ Загружает модель из HuggingFace Hub. - + Args: repo_id: ID репозитория **kwargs: Дополнительные параметры - + Returns: tuple: (модель, токенизатор) """ from transformers import AutoTokenizer - + # Загружаем токенизатор tokenizer = AutoTokenizer.from_pretrained(repo_id, **kwargs) - + # Загружаем конфигурацию config = AutoConfig.from_pretrained(repo_id, **kwargs) - + # Создаем модель llm на основе конфигурации llm_config = { "vocab_size": config.vocab_size, @@ -159,63 +158,56 @@ class HFUtils: "max_position_embeddings": config.max_position_embeddings, "dropout": config.hidden_dropout_prob, } - + # Загружаем модель через адаптер model = HFAdapter.from_pretrained( - f"{repo_id}/pytorch_model.bin", - HFAdapterConfig.from_llm_config(llm_config) + f"{repo_id}/pytorch_model.bin", HFAdapterConfig.from_llm_config(llm_config) ) - + return model, tokenizer - + @staticmethod def compare_with_hf_model( - llm_model, - hf_model_name: str = "gpt2", - test_input: str = "Hello world" + llm_model, hf_model_name: str = "gpt2", test_input: str = "Hello world" ) -> Dict[str, Any]: """ Сравнивает llm модель с эталонной моделью из HuggingFace. - + Args: llm_model: Модель из библиотеки llm hf_model_name: Имя модели HuggingFace для сравнения test_input: Тестовый вход - + Returns: Dict: Результаты сравнения """ from transformers import AutoModelForCausalLM, AutoTokenizer - + # Загружаем эталонную модель hf_tokenizer = AutoTokenizer.from_pretrained(hf_model_name) hf_model = AutoModelForCausalLM.from_pretrained(hf_model_name) - + # Подготавливаем входные данные inputs = hf_tokenizer(test_input, return_tensors="pt") - + # Получаем логиты от обеих моделей with torch.no_grad(): hf_logits = hf_model(**inputs).logits - llm_logits = llm_model(inputs['input_ids']) - + llm_logits = llm_model(inputs["input_ids"]) + # Сравниваем результаты hf_probs = torch.softmax(hf_logits[0, -1], dim=-1) llm_probs = torch.softmax(llm_logits[0, -1], dim=-1) - + # Вычисляем метрики kl_divergence = torch.nn.functional.kl_div( - torch.log(llm_probs + 1e-8), - hf_probs, - reduction='batchmean' + torch.log(llm_probs + 1e-8), hf_probs, reduction="batchmean" ) - + cosine_similarity = torch.nn.functional.cosine_similarity( - hf_logits.flatten(), - llm_logits.flatten(), - dim=0 + hf_logits.flatten(), llm_logits.flatten(), dim=0 ) - + return { "kl_divergence": kl_divergence.item(), "cosine_similarity": cosine_similarity.item(), @@ -228,58 +220,52 @@ class TokenizerWrapper: """ Обертка для токенизатора с дополнительными утилитами. """ - + def __init__(self, tokenizer): self.tokenizer = tokenizer - + def encode_batch(self, texts: List[str], **kwargs) -> Dict[str, torch.Tensor]: """ Кодирует батч текстов. - + Args: texts: Список текстов **kwargs: Дополнительные параметры токенизации - + Returns: Dict: Токенизированные данные """ return self.tokenizer( - texts, - padding=True, - truncation=True, - return_tensors="pt", - **kwargs + texts, padding=True, truncation=True, return_tensors="pt", **kwargs ) - + def decode_batch(self, token_ids: torch.Tensor, **kwargs) -> List[str]: """ Декодирует батч токенов. - + Args: token_ids: Тензор с токенами **kwargs: Дополнительные параметры декодирования - + Returns: List[str]: Декодированные тексты """ if token_ids.dim() == 1: token_ids = token_ids.unsqueeze(0) - + texts = [] for i in range(token_ids.size(0)): text = self.tokenizer.decode( - token_ids[i], - skip_special_tokens=True, - **kwargs + token_ids[i], skip_special_tokens=True, **kwargs ) texts.append(text) - + return texts - + def get_vocab_size(self) -> int: """Возвращает размер словаря.""" return len(self.tokenizer) - + def get_special_tokens(self) -> Dict[str, int]: """Возвращает специальные токены.""" return { @@ -290,36 +276,27 @@ class TokenizerWrapper: } -def create_hf_pipeline( - llm_model, - tokenizer=None, - device: str = "auto", - **kwargs -): +def create_hf_pipeline(llm_model, tokenizer=None, device: str = "auto", **kwargs): """ Создает HuggingFace pipeline из llm модели. - + Args: llm_model: Модель из библиотеки llm tokenizer: Токенизатор device: Устройство для вычислений **kwargs: Дополнительные параметры pipeline - + Returns: transformers.Pipeline: Готовый pipeline """ from transformers import pipeline - + # Конвертируем модель в HF формат hf_model, tokenizer = HFUtils.convert_to_hf_format(llm_model, tokenizer) - + # Создаем pipeline pipe = pipeline( - "text-generation", - model=hf_model, - tokenizer=tokenizer, - device=device, - **kwargs + "text-generation", model=hf_model, tokenizer=tokenizer, device=device, **kwargs ) - + return pipe diff --git a/llm/src/llm/core/base_model.py b/llm/src/llm/core/base_model.py index 0560b53..422809f 100644 --- a/llm/src/llm/core/base_model.py +++ b/llm/src/llm/core/base_model.py @@ -19,23 +19,25 @@ from abc import ABC, abstractmethod from typing import Optional, Tuple import torch + class BaseModel(nn.Module, ABC): """ Абстрактный класс — стандарт для всех архитектур LLM. - + Научная идея: Реализация унифицированного входа/выхода для поддержки построения и обучения любых современных языковых моделей. - + Args: config (dict): Параметры архитектуры (размерность эмбеддингов, число слоев, heads и т.д.) - + Attributes: config (dict): Конфиг модели """ + def __init__(self, config: dict): """ Инициализация модели. - + Args: config (dict): Настройки архитектуры модели (размеры слоев, типы блоков и т.д.) """ @@ -43,10 +45,12 @@ class BaseModel(nn.Module, ABC): self.config = config @abstractmethod - def forward(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor: + def forward( + self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None + ) -> torch.Tensor: """ Прямой проход — получение логитов для входных токенов. - + Args: input_ids (Tensor[int]): Индексы токенов [batch, seq_len] attention_mask (Optional[Tensor[bool]]): Маска разрешенных позиций (если требуется) [batch, seq_len] @@ -59,7 +63,7 @@ class BaseModel(nn.Module, ABC): def generate(self, input_ids: torch.Tensor, max_length: int = 50) -> torch.Tensor: """ Генерация текста (авторегрессивно, greedy или sampling). - + Args: input_ids (Tensor[int]): Начальные токены [batch, start_len] max_length (int): Максимальная длина последовательности diff --git a/llm/src/llm/core/cached_decoder.py b/llm/src/llm/core/cached_decoder.py index dcc5a4e..7cac0ff 100644 --- a/llm/src/llm/core/cached_decoder.py +++ b/llm/src/llm/core/cached_decoder.py @@ -6,6 +6,7 @@ from .feed_forward import FeedForward from .multi_head_attention import MultiHeadAttention from .rope import RoPE + class CachedDecoder(nn.Module): """ Универсальный декодерный блок для современных LLM (GPT, LLaMA, др.), поддерживает кэширование key-value для эффективной генерации. @@ -28,7 +29,7 @@ class CachedDecoder(nn.Module): norm_layer (тип nn.Module): Normalization слой (LayerNorm или RMSNorm) dropout (float): Dropout rope (RoPE|None): Экземпляр RoPE (для LLaMA) - + Пример (GPT2 style): >>> decoder = CachedDecoder( ... feed_forward_layer=FeedForward(...), @@ -36,6 +37,7 @@ class CachedDecoder(nn.Module): ... num_heads=4, emb_size=256, head_size=64, max_seq_len=128) >>> out, cache = decoder(x, use_cache=True) """ + def __init__( self, feed_forward_layer: nn.Module, @@ -49,7 +51,7 @@ class CachedDecoder(nn.Module): ): """ Инициализация декодера с кэшированием. - + Поведение аналогично блоку TransformerDecoderLayer, но с гибкой возможностью подмены любых подкомпонент (активация, norm, позиции). @@ -85,7 +87,7 @@ class CachedDecoder(nn.Module): ): """ Прямой проход с поддержкой кэша. - + Args: x (Tensor[float]): [batch, seq_len, emb_size] — скрытые состояния mask (Optional[Tensor]): маска внимания (или causal mask), shape [seq_len, seq_len] @@ -111,4 +113,4 @@ class CachedDecoder(nn.Module): if use_cache: return (result, kv_caches) else: - return (result, None) \ No newline at end of file + return (result, None) diff --git a/llm/src/llm/core/decoder.py b/llm/src/llm/core/decoder.py index 40cb9dd..396342f 100644 --- a/llm/src/llm/core/decoder.py +++ b/llm/src/llm/core/decoder.py @@ -3,6 +3,7 @@ import torch from .feed_forward import FeedForward from .multi_head_attention import MultiHeadAttention + class Decoder(nn.Module): """ Базовый автогерессивный блок-декодер трансформера (без кэша KV). @@ -24,12 +25,14 @@ class Decoder(nn.Module): >>> out = decoder(x) >>> print(out.shape) # torch.Size([1, 10, 512]) """ - def __init__(self, + + def __init__( + self, num_heads: int, emb_size: int, head_size: int, max_seq_len: int, - dropout: float = 0.1 + dropout: float = 0.1, ): """ Инициализация декодера. @@ -43,11 +46,11 @@ class Decoder(nn.Module): """ super().__init__() self._heads = MultiHeadAttention( - num_heads=num_heads, - emb_size=emb_size, - head_size=head_size, - max_seq_len=max_seq_len, - dropout=dropout + num_heads=num_heads, + emb_size=emb_size, + head_size=head_size, + max_seq_len=max_seq_len, + dropout=dropout, ) self._ff = FeedForward(emb_size=emb_size, dropout=dropout) self._norm1 = nn.LayerNorm(emb_size) @@ -73,7 +76,7 @@ class Decoder(nn.Module): # Self-Attention блок attention, _ = self._heads(x, mask, use_cache=False, cache=None) out = self._norm1(attention + x) - + # FeedForward блок ffn_out = self._ff(out) - return self._norm2(ffn_out + out) \ No newline at end of file + return self._norm2(ffn_out + out) diff --git a/llm/src/llm/core/feed_forward.py b/llm/src/llm/core/feed_forward.py index 3168dea..f13018f 100644 --- a/llm/src/llm/core/feed_forward.py +++ b/llm/src/llm/core/feed_forward.py @@ -16,7 +16,7 @@ class FeedForward(nn.Module): - После внимания каждому токену применяется одинаковая двухслойная нейросеть. - Дает глубокую нелинейность; позволяет модели не только сопоставлять, но и моделировать сложные связи между токенами. - Изначально предложен в «Attention is All You Need» (Vaswani et al., 2017). - + Формула: FFN(x) = Dropout(W2·act(W1·x)) где act — ReLU, GELU и др., обычно expansion x4. @@ -33,22 +33,23 @@ class FeedForward(nn.Module): - Добавляет нелинейность в архитектуру трансформера - Обеспечивает взаимодействие между различными размерностями эмбеддингов - Работает независимо для каждого токена в последовательности - + Args: emb_size (int): размерность входных эмбеддингов dropout (float): вероятность(dropout) activation (str): нелинейная функция (relu, gelu, gelu_exact) - + Пример: >>> ff = FeedForward(emb_size=512, dropout=0.1) >>> x = torch.randn(32, 10, 512) >>> output = ff(x) >>> print(output.shape) # torch.Size([32, 10, 512]) """ + def __init__(self, emb_size: int, dropout: float = 0.1, activation: str = "relu"): """ Инициализация слоя Feed Forward Network. - + Args: emb_size: Размерность входных эмбеддингов dropout: Вероятность dropout для регуляризации (по умолчанию: 0.1) @@ -73,23 +74,23 @@ class FeedForward(nn.Module): def forward(self, x: torch.Tensor): """ Прямой проход через слой Feed Forward Network. - + Args: x: Входной тензор размерности [batch_size, seq_len, emb_size] - + Returns: Тензор той же размерности, что и входной """ # Сохраняем dtype входных данных input_dtype = x.dtype - + # Приводим веса к нужному типу если необходимо if input_dtype != self._layer1.weight.dtype: self._layer1 = self._layer1.to(dtype=input_dtype) self._layer2 = self._layer2.to(dtype=input_dtype) - + # Пропустим тензор x по очереди через все созданные слои x = self._layer1(x) x = self._activation(x) x = self._layer2(x) - return self._dropout(x) \ No newline at end of file + return self._dropout(x) diff --git a/llm/src/llm/core/gelu.py b/llm/src/llm/core/gelu.py index 95772f6..b04aa55 100644 --- a/llm/src/llm/core/gelu.py +++ b/llm/src/llm/core/gelu.py @@ -1,6 +1,7 @@ import torch from torch import nn + class GELU(nn.Module): """ Гауссовская Эрф-активация (GELU, Gaussian Error Linear Unit). @@ -17,11 +18,14 @@ class GELU(nn.Module): >>> y = gelu(torch.tensor([-1.0, 0.0, 1.0])) >>> print(y) """ + def __init__(self): super().__init__() self.sqrt_2_over_pi = torch.sqrt(torch.tensor(2.0) / math.pi) - + def forward(self, x: torch.Tensor) -> torch.Tensor: - return 0.5 * x * (1 + torch.tanh( - self.sqrt_2_over_pi * (x + 0.044715 * torch.pow(x, 3)) - )) \ No newline at end of file + return ( + 0.5 + * x + * (1 + torch.tanh(self.sqrt_2_over_pi * (x + 0.044715 * torch.pow(x, 3)))) + ) diff --git a/llm/src/llm/core/head_attention.py b/llm/src/llm/core/head_attention.py index 4b32032..410f499 100644 --- a/llm/src/llm/core/head_attention.py +++ b/llm/src/llm/core/head_attention.py @@ -4,6 +4,7 @@ import torch.nn.functional as F from math import sqrt from .rope import RoPE + class HeadAttention(nn.Module): """ Одноголовый механизм внимания (scaled dot-product attention) — фундаментальный строительный блок всех современных Transformer. @@ -11,11 +12,11 @@ class HeadAttention(nn.Module): Научная суть: - Attention учит модель самостоятельно "выбирать" важные связи между словами, независимо от их положения. - Механизм causal mask гарантирует невозможность "заглядывания в будущее" при генерации (авторегрессия). - + Формула: Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) · V (Q — запросы, K — ключи, V — значения; d_k — размерность ключа) - + Поддерживает Rotary Position Encoding (RoPE) для относительного позиционного кодирования. Args: @@ -28,14 +29,17 @@ class HeadAttention(nn.Module): - Использует нижнетреугольную маску для предотвращения "заглядывания в будущее" - Автоматически адаптируется к разным версиям PyTorch - Поддерживает batch-обработку входных данных - + Пример использования: >>> attention = HeadAttention(emb_size=64, head_size=32, max_seq_len=128) >>> x = torch.randn(1, 10, 64) >>> output, _ = attention(x) >>> print(output.shape) # torch.Size([1, 10, 32]) """ - def __init__(self, emb_size: int, head_size: int, max_seq_len: int, rope: RoPE = None): + + def __init__( + self, emb_size: int, head_size: int, max_seq_len: int, rope: RoPE = None + ): super().__init__() self._emb_size = emb_size self._head_size = head_size @@ -49,21 +53,25 @@ class HeadAttention(nn.Module): # Создание causal маски mask = torch.tril(torch.ones(max_seq_len, max_seq_len)) - self.register_buffer('_tril_mask', mask.bool() if hasattr(torch, 'bool') else mask.byte()) + self.register_buffer( + "_tril_mask", mask.bool() if hasattr(torch, "bool") else mask.byte() + ) - def forward(self, x: torch.Tensor, use_cache: bool = True, cache: tuple = None) -> tuple: + def forward( + self, x: torch.Tensor, use_cache: bool = True, cache: tuple = None + ) -> tuple: """ Прямой проход через слой внимания. - + Аргументы: x (torch.Tensor): Входной тензор формы [batch_size, seq_len, emb_size] - + Возвращает: torch.Tensor: Выходной тензор формы [batch_size, seq_len, head_size] - + Исключения: ValueError: Если длина последовательности превышает max_seq_len - + Пример внутренних преобразований: Для входа x.shape = [2, 5, 64]: 1. Q/K/V преобразования -> [2, 5, 32] @@ -73,7 +81,9 @@ class HeadAttention(nn.Module): """ seq_len = x.shape[1] if seq_len > self._max_seq_len: - raise ValueError(f"Длина последовательности {seq_len} превышает максимум {self._max_seq_len}") + raise ValueError( + f"Длина последовательности {seq_len} превышает максимум {self._max_seq_len}" + ) k = self._k(x) # [B, T, hs] q = self._q(x) # [B, T, hs] @@ -88,16 +98,18 @@ class HeadAttention(nn.Module): k_cache, v_cache = cache k = torch.cat([k_cache, k], dim=1) # [B, cache_len + T, hs] v = torch.cat([v_cache, v], dim=1) # [B, cache_len + T, hs] - + scores = q @ k.transpose(-2, -1) / sqrt(self._head_size) - + if cache is None: - scores = scores.masked_fill(~self._tril_mask[:seq_len, :seq_len], float('-inf')) - + scores = scores.masked_fill( + ~self._tril_mask[:seq_len, :seq_len], float("-inf") + ) + weights = F.softmax(scores, dim=-1) x_out = weights @ v # [B, T, hs] if use_cache is True: return (x_out, (k, v)) else: - return (x_out, None) \ No newline at end of file + return (x_out, None) diff --git a/llm/src/llm/core/multi_head_attention.py b/llm/src/llm/core/multi_head_attention.py index c788359..21fc03f 100644 --- a/llm/src/llm/core/multi_head_attention.py +++ b/llm/src/llm/core/multi_head_attention.py @@ -3,6 +3,7 @@ import torch from .head_attention import HeadAttention from .rope import RoPE + class MultiHeadAttention(nn.Module): """ Мультиголовый (многоголовый) механизм внимания — ключевой компонент любого Transformer. @@ -12,7 +13,7 @@ class MultiHeadAttention(nn.Module): чтобы видеть разные связи в последовательности (разный контекст, локально/глобально). - Каждый attention блок работает независимо, выход конкатенируется. - Механизм предложен в статье "Attention is All You Need" (Vaswani et al., 2017). - + Формула внимания для одной головы: Attention(Q, K, V) = softmax(QK^T/sqrt(d_k))·V Мультиголовый: @@ -32,7 +33,16 @@ class MultiHeadAttention(nn.Module): >>> out, cache = mha(x) >>> print(out.shape) """ - def __init__(self, num_heads: int, emb_size: int, head_size: int, max_seq_len: int, rope: RoPE = None, dropout: float = 0.1): + + def __init__( + self, + num_heads: int, + emb_size: int, + head_size: int, + max_seq_len: int, + rope: RoPE = None, + dropout: float = 0.1, + ): """ Инициализация многоголового внимания. @@ -49,18 +59,27 @@ class MultiHeadAttention(nn.Module): - max_seq_len зависит от задачи (512 для BERT, 2048 для GPT-3) """ super().__init__() - self._heads = nn.ModuleList([ - HeadAttention( - emb_size=emb_size, - head_size=head_size, - max_seq_len=max_seq_len, - rope=rope, - ) for _ in range(num_heads) - ]) + self._heads = nn.ModuleList( + [ + HeadAttention( + emb_size=emb_size, + head_size=head_size, + max_seq_len=max_seq_len, + rope=rope, + ) + for _ in range(num_heads) + ] + ) self._layer = nn.Linear(head_size * num_heads, emb_size) self._dropout = nn.Dropout(dropout) - def forward(self, x: torch.Tensor, mask: torch.Tensor = None, use_cache: bool = True, cache: list = None): + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor = None, + use_cache: bool = True, + cache: list = None, + ): """ Прямой проход (forward): Для каждого токена оценивает "важность" остальных токенов сразу через несколько attention-блоков. @@ -76,7 +95,7 @@ class MultiHeadAttention(nn.Module): 4. Линейная проекция: - Выход: [batch_size, seq_len, emb_size] 5. Применение dropout - + Args: x (Tensor[float]): [batch, seq_len, emb_size] — вход mask (Optional[Tensor[bool]]): маска позиции [seq_len, seq_len] @@ -94,11 +113,11 @@ class MultiHeadAttention(nn.Module): Пример преобразований для emb_size=512, num_heads=8: Вход: [4, 100, 512] -> Каждая голова: [4, 100, 64] - -> После внимания: 8 x [4, 100, 64] + -> После внимания: 8 x [4, 100, 64] -> Конкатенация: [4, 100, 512] -> Проекция: [4, 100, 512] -> Dropout: [4, 100, 512] - + Пример: >>> out, caches = mha(x) >>> out.shape # [batch, seq_len, emb_size] @@ -109,20 +128,20 @@ class MultiHeadAttention(nn.Module): head_cache = cache[i] if cache is not None else None result = head(x, use_cache=use_cache, cache=head_cache) attention_results.append(result) - + outputs, caches = zip(*attention_results) attention_outputs = list(outputs) kv_caches = list(caches) - + # 2. Объединяем результаты всех голов concatenated_attention = torch.cat(attention_outputs, dim=-1) - + # 3. Проецируем в пространство эмбеддингов projected_output = self._layer(concatenated_attention) - + # 4. Применяем dropout для регуляризации final_output = self._dropout(projected_output) - + if use_cache is True: return (final_output, kv_caches) else: diff --git a/llm/src/llm/core/positional_embeddings.py b/llm/src/llm/core/positional_embeddings.py index fb0e7fb..de728dc 100644 --- a/llm/src/llm/core/positional_embeddings.py +++ b/llm/src/llm/core/positional_embeddings.py @@ -1,13 +1,14 @@ import torch from torch import nn, Tensor + class PositionalEmbeddings(nn.Module): """ Обучаемые позиционные эмбеддинги (learnable positional embeddings). - Позиционные эмбеддинги используются в нейросетях для передачи информации + Позиционные эмбеддинги используются в нейросетях для передачи информации о позиции элементов в последовательности (например, в Transformer). - + Научная суть: - Трансформеры не используют рекуррентность, а значит сами по себе не различают порядок слов. - Позиционные эмбеддинги добавляются к токеновым, чтобы сеть понимала, в каком месте последовательности находится каждый токен. @@ -16,7 +17,7 @@ class PositionalEmbeddings(nn.Module): Args: max_seq_len (int): максимальная длина последовательности emb_size (int): размер вектора позиции - + Пример использования: >>> pos_encoder = PositionalEmbeddings(max_seq_len=100, emb_size=256) >>> # Получить эмбеддинги для последовательности из 10 элементов @@ -36,23 +37,22 @@ class PositionalEmbeddings(nn.Module): self.max_seq_len = max_seq_len self.emb_size = emb_size self.embedding = nn.Embedding( - num_embeddings=max_seq_len, - embedding_dim=emb_size + num_embeddings=max_seq_len, embedding_dim=emb_size ) def forward(self, seq_len: int, start_pos: int = 0) -> Tensor: """ Возвращает позиционные эмбеддинги для заданной длины последовательности. - + Args: seq_len (int): Длина последовательности (1 <= seq_len <= max_seq_len) - + Returns: Tensor: Тензор позиционных эмбеддингов формы [seq_len, emb_size] - + Raises: IndexError: Если seq_len выходит за допустимые границы - + Пример: >>> pos_encoder = PositionalEmbeddings(100, 64) >>> emb = pos_encoder(10) # Тензор 10x64 @@ -62,5 +62,9 @@ class PositionalEmbeddings(nn.Module): if start_pos == 0: positions = torch.arange(seq_len, device=self.embedding.weight.device) else: - positions = torch.arange(start=start_pos, end=start_pos + seq_len, device=self.embedding.weight.device) + positions = torch.arange( + start=start_pos, + end=start_pos + seq_len, + device=self.embedding.weight.device, + ) return self.embedding(positions) diff --git a/llm/src/llm/core/rms_norm.py b/llm/src/llm/core/rms_norm.py index 47def26..6b2a5fe 100644 --- a/llm/src/llm/core/rms_norm.py +++ b/llm/src/llm/core/rms_norm.py @@ -33,23 +33,23 @@ class RMSNorm(nn.Module): - Упрощенный вариант LayerNorm без вычисления среднего, только деление на rms. - Лучшая численная стабильность на больших моделях, меньше вычислений. - Применяется в LLaMA, PaLM и др. - + Формула: RMSNorm(x) = (x / sqrt(mean(x²) + eps)) * w (w — обучаемый вектор) Args: dim (int): размер последнего измерения (обычно emb_size) eps (float): для численной устойчивости - + Пример: >>> norm = RMSNorm(emb_size) >>> out = norm(x) """ - + def __init__(self, dim: int, eps: float = 1e-6): """ Инициализация RMSNorm слоя. - + Args: dim: Размерность нормализуемого измерения eps: Малое значение для численной стабильности (по умолчанию 1e-6) @@ -57,27 +57,27 @@ class RMSNorm(nn.Module): super().__init__() self._eps = eps self._w = nn.Parameter(torch.ones(dim)) - + def forward(self, x: torch.Tensor) -> torch.Tensor: """ Прямой проход через RMSNorm слой. - + Args: x: Входной тензор формы [..., dim] - + Returns: Нормализованный тензор той же формы, что и входной - + Формула: output = w * (x / sqrt(mean(x²) + eps)) """ # Вычисление RMS (Root Mean Square) по последнему измерению rms = (x.pow(2).mean(-1, keepdim=True) + self._eps) ** 0.5 - + # Нормализация и масштабирование norm_x = x / rms return self._w * norm_x - + def extra_repr(self) -> str: """Строковое представление для отладки.""" - return f'dim={self._w.shape[0]}, eps={self._eps}' \ No newline at end of file + return f"dim={self._w.shape[0]}, eps={self._eps}" diff --git a/llm/src/llm/core/rope.py b/llm/src/llm/core/rope.py index d07b348..ebf188a 100644 --- a/llm/src/llm/core/rope.py +++ b/llm/src/llm/core/rope.py @@ -26,72 +26,72 @@ from typing import Optional class RoPE(nn.Module): """ Rotary Positional Embeddings (RoPE) для механизма внимания. - + Кодирует позиционную информацию через вращение векторов запросов и ключей в многомерном пространстве с использованием синусов и косинусов. - + Args: head_size: Размерность головы внимания (должен быть четным) max_seq_len: Максимальная длина последовательности base: Базовое значение для вычисления частот (по умолчанию 10000) - + Attributes: cos_matrix: Буферизованная матрица косинусов формы [max_seq_len, head_size//2] sin_matrix: Буферизованная матрица синусов формы [max_seq_len, head_size//2] """ - + def __init__(self, head_size: int, max_seq_len: int, base: int = 10_000): """ Инициализация RoPE эмбеддингов. - + Args: head_size: Размерность головы внимания (должен быть четным) max_seq_len: Максимальная поддерживаемая длина последовательности base: Базовое значение для вычисления частот (типично 10000) - + Raises: AssertionError: Если head_size не четный """ super().__init__() assert head_size % 2 == 0, "head_size должен быть четным" - + # Вычисление частот: θ_i = base^(-2i/d) для i ∈ [0, d/2-1] freqs = 1.0 / (base ** (2 * torch.arange(head_size // 2).float() / head_size)) - + # Позиции от 0 до max_seq_len-1 positions = torch.arange(max_seq_len).float() - + # Внешнее произведение: m * θ_i для всех позиций и частот freq_matrix = positions.unsqueeze(1) * freqs.unsqueeze(0) # Предвычисление матриц косинусов и синусов - self.register_buffer('cos_matrix', torch.cos(freq_matrix)) - self.register_buffer('sin_matrix', torch.sin(freq_matrix)) + self.register_buffer("cos_matrix", torch.cos(freq_matrix)) + self.register_buffer("sin_matrix", torch.sin(freq_matrix)) def forward(self, x: torch.Tensor) -> torch.Tensor: """ Применение ротационного позиционного кодирования к входному тензору. - + Args: x: Входной тензор формы [batch_size, seq_len, head_size] - + Returns: Тензор с примененным RoPE формы [batch_size, seq_len, head_size] - + Алгоритм: 1. Разделение векторов на четные и нечетные компоненты 2. Применение вращения через синусы и косинусы 3. Объединение компонент обратно """ seq_len = x.size(1) - + # Берем нужную часть матриц и приводим к типу x cos = self.cos_matrix[:seq_len].to(x.dtype) # [seq_len, head_size//2] sin = self.sin_matrix[:seq_len].to(x.dtype) # [seq_len, head_size//2] - + # Разделяем на четные и нечетные компоненты x_even = x[:, :, 0::2] # [batch_size, seq_len, head_size//2] - x_odd = x[:, :, 1::2] # [batch_size, seq_len, head_size//2] + x_odd = x[:, :, 1::2] # [batch_size, seq_len, head_size//2] # Применяем поворот: q' = q * cos(mθ) + rotate(q) * sin(mθ) x_rotated_even = x_even * cos - x_odd * sin @@ -101,4 +101,4 @@ class RoPE(nn.Module): x_rotated = torch.stack([x_rotated_even, x_rotated_odd], dim=-1) x_rotated = x_rotated.flatten(-2) # [batch_size, seq_len, head_size] - return x_rotated \ No newline at end of file + return x_rotated diff --git a/llm/src/llm/core/silu.py b/llm/src/llm/core/silu.py index c23d542..ade010a 100644 --- a/llm/src/llm/core/silu.py +++ b/llm/src/llm/core/silu.py @@ -1,19 +1,21 @@ import torch from torch import nn + class SiLU(nn.Module): """ SiLU (Swish) — современная активационная функция для нейросетей. - + Научная суть: - Формула: $SiLU(x) = x * \sigm(x)$, где $\sigm(x)$ — сигмоида. - Более гладкая альтернатива ReLU, улучшает поток градиентов в глубоких сетях. - Используется во многих «state-of-the-art» архитектурах (SwiGLU, PaLM, LLaMA). - - Также известна как Swish (Ramachandran et al, 2017). + - Также известна как Swish (Ramachandran et al, 2017). Пример: >>> act = SiLU() >>> x = torch.tensor([-1.0, 0.0, 1.0]) >>> print(act(x)) """ + def forward(self, x: torch.Tensor): - return torch.sigmoid(x) * x \ No newline at end of file + return torch.sigmoid(x) * x diff --git a/llm/src/llm/core/swi_glu.py b/llm/src/llm/core/swi_glu.py index ace972d..d0820b0 100644 --- a/llm/src/llm/core/swi_glu.py +++ b/llm/src/llm/core/swi_glu.py @@ -27,13 +27,13 @@ class SwiGLU(nn.Module): SwiGLU (Swish-Gated Linear Unit) — современная нелинейность для архитектур LLM (LLaMA, PaLM). Реализация SwiGLU активационной функции. - + Состоит из трех линейных слоев и активации SiLU: 1. Gate слой + SiLU активация 2. Up слой (линейное преобразование) 3. Element-wise multiplication gate и up 4. Down слой (линейная проекция) - + Научная суть: - Сохраняет преимущества GLU (раздельные гейтом и телом) + мощность Swish/SiLU активации. - Дает надежную гладкую активацию, хорошо работает на больших масштабах. @@ -50,11 +50,11 @@ class SwiGLU(nn.Module): >>> ff = SwiGLU(emb_size=512, dropout=0.1) >>> y = ff(torch.randn(2,10,512)) """ - + def __init__(self, emb_size: int, dropout: float = 0.1): """ Инициализация SwiGLU слоя. - + Args: emb_size: Размерность входных/выходных эмбеддингов dropout: Вероятность dropout (по умолчанию 0.1) @@ -69,13 +69,13 @@ class SwiGLU(nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: """ Прямой проход через SwiGLU слой. - + Args: x: Входной тензор формы [batch_size, seq_len, emb_size] - + Returns: Выходной тензор формы [batch_size, seq_len, emb_size] - + Алгоритм: 1. gate = SiLU(linear_gate(x)) 2. up = linear_up(x) @@ -83,19 +83,19 @@ class SwiGLU(nn.Module): 4. apply dropout """ # Gate ветвь: линейное преобразование + активация - gate_out = self._gate(x) # [batch, seq, 4*emb] - activation_out = self._activation(gate_out) # [batch, seq, 4*emb] - + gate_out = self._gate(x) # [batch, seq, 4*emb] + activation_out = self._activation(gate_out) # [batch, seq, 4*emb] + # Up ветвь: линейное преобразование - up_out = self._up(x) # [batch, seq, 4*emb] - + up_out = self._up(x) # [batch, seq, 4*emb] + # Element-wise multiplication (gating mechanism) - out = up_out * activation_out # поэлементное умножение! - + out = up_out * activation_out # поэлементное умножение! + # Final projection and dropout - out = self._down(out) # [batch, seq, emb] + out = self._down(out) # [batch, seq, emb] return self._dropout(out) - + def extra_repr(self) -> str: """Строковое представление для отладки.""" - return f'emb_size={self._gate.in_features}, dropout={self._dropout.p}' \ No newline at end of file + return f"emb_size={self._gate.in_features}, dropout={self._dropout.p}" diff --git a/llm/src/llm/core/token_embeddings.py b/llm/src/llm/core/token_embeddings.py index 5674e6f..56e414d 100644 --- a/llm/src/llm/core/token_embeddings.py +++ b/llm/src/llm/core/token_embeddings.py @@ -2,13 +2,14 @@ import torch from torch import nn from torch import Tensor + class TokenEmbeddings(nn.Module): """ Токеновые эмбеддинги — обучаемые векторные представления для каждого токена словаря. Преобразует целочисленные индексы токенов в обучаемые векторные представления фиксированного размера. Обычно используется как первый слой в нейронных сетях для задач NLP. - + Научная суть: - Первый шаг для любого NLP-модуля: вместо индекса токена подаём его dense-вектор. - Эти вектора изучаются в процессе обучения и отражают скрытые взаимосвязи между токенами. @@ -22,18 +23,18 @@ class TokenEmbeddings(nn.Module): Примечание: - Индексы должны быть в диапазоне [0, vocab_size-1] - Эмбеддинги инициализируются случайно и обучаются в процессе тренировки модели - + Пример: >>> emb = TokenEmbeddings(vocab_size=10000, emb_size=256) >>> tokens = torch.tensor([[1, 2, 3]]) >>> vecs = emb(tokens) >>> vecs.shape # torch.Size([1, 3, 256]) """ + def __init__(self, vocab_size: int, emb_size: int): super().__init__() self._embedding = nn.Embedding( - num_embeddings=vocab_size, - embedding_dim=emb_size + num_embeddings=vocab_size, embedding_dim=emb_size ) def forward(self, x: Tensor) -> Tensor: @@ -55,10 +56,7 @@ if __name__ == "__main__": embedding = TokenEmbeddings(vocab_size=100, emb_size=128) # Создаем тензор с индексами в пределах vocab_size (0-99) - tensor = torch.tensor([ - [11, 45, 76, 34], - [34, 67, 45, 54] - ]) + tensor = torch.tensor([[11, 45, 76, 34], [34, 67, 45, 54]]) # Проверяем индексы if (tensor >= 100).any(): @@ -66,4 +64,4 @@ if __name__ == "__main__": output = embedding(tensor) print("Embeddings shape:", output.shape) - print(f"{output.shape} | {output.mean().item():.11f}") # Формат как в ТЗ \ No newline at end of file + print(f"{output.shape} | {output.mean().item():.11f}") # Формат как в ТЗ diff --git a/llm/src/llm/models/gpt/gpt.py b/llm/src/llm/models/gpt/gpt.py index 2121da8..0ff743f 100644 --- a/llm/src/llm/models/gpt/gpt.py +++ b/llm/src/llm/models/gpt/gpt.py @@ -34,10 +34,10 @@ from llm.core.positional_embeddings import PositionalEmbeddings class GPT(BaseModel): """ Original GPT (Generative Pre-trained Transformer) модель. - + Первая версия трансформерной архитектуры от OpenAI, предназначенная для генеративного предобучения на текстовых данных. - + Args: config: Словарь конфигурации с параметрами: - vocab_size: Размер словаря токенов @@ -46,7 +46,7 @@ class GPT(BaseModel): - num_layers: Количество декодерных слоев - max_position_embeddings: Максимальная длина последовательности - dropout: Вероятность dropout - + Attributes: _token_embeddings: Слой векторных представлений токенов _position_embeddings: Слой позиционных эмбеддингов @@ -54,30 +54,34 @@ class GPT(BaseModel): _norm: Финальный слой нормализации _linear: Выходной линейный слой """ + def __init__(self, config): super().__init__(config) # Инициализация слоев self._max_seq_len = config["max_position_embeddings"] self._token_embeddings = TokenEmbeddings( - vocab_size=config["vocab_size"], - emb_size=config["embed_dim"] + vocab_size=config["vocab_size"], emb_size=config["embed_dim"] ) self._position_embeddings = PositionalEmbeddings( - max_seq_len=config["max_position_embeddings"], - emb_size=config["embed_dim"] + max_seq_len=config["max_position_embeddings"], emb_size=config["embed_dim"] ) self._dropout = nn.Dropout(config["dropout"]) # head_size = emb_size // num_heads - self._decoders = nn.ModuleList([Decoder( - num_heads=config["num_heads"], - emb_size=config["embed_dim"], - head_size=config["embed_dim"] // config["num_heads"], - max_seq_len=config["max_position_embeddings"], - dropout=config["dropout"] - ) for _ in range(config["num_layers"])]) + self._decoders = nn.ModuleList( + [ + Decoder( + num_heads=config["num_heads"], + emb_size=config["embed_dim"], + head_size=config["embed_dim"] // config["num_heads"], + max_seq_len=config["max_position_embeddings"], + dropout=config["dropout"], + ) + for _ in range(config["num_layers"]) + ] + ) self._linear = nn.Linear(config["embed_dim"], config["vocab_size"]) - + @property def max_seq_len(self): """Возвращает максимальную длину последовательности.""" @@ -85,57 +89,60 @@ class GPT(BaseModel): def forward(self, x: torch.Tensor, attention_mask=None) -> torch.Tensor: """Прямой проход через GPT - + Args: x: Входной тензор [batch_size, seq_len] - + Returns: Тензор логитов [batch_size, seq_len, vocab_size] """ # Проверка длины последовательности if x.size(1) > self._max_seq_len: - raise ValueError(f"Длина последовательности {x.size(1)} превышает максимальную {self._max_seq_len}") - + raise ValueError( + f"Длина последовательности {x.size(1)} превышает максимальную {self._max_seq_len}" + ) + # Эмбеддинги токенов и позиций tok_out = self._token_embeddings(x) # [batch, seq_len, emb_size] pos_out = self._position_embeddings(x.size(1)) # [seq_len, emb_size] - + # Комбинирование - out = self._dropout(tok_out + pos_out.unsqueeze(0)) # [batch, seq_len, emb_size] - + out = self._dropout( + tok_out + pos_out.unsqueeze(0) + ) # [batch, seq_len, emb_size] + # Стек декодеров for decoder in self._decoders: out = decoder(out) - + return self._linear(out) # [batch, seq_len, vocab_size] + # def forward(self, input_ids, attention_mask=None): + # B, T = input_ids.size() + # pos = torch.arange(0, T, device=input_ids.device).unsqueeze(0) + # + # x = self.token_emb(input_ids) + self.pos_emb(pos) + # + # for block in self.blocks: + # x = block(x, attention_mask) + # + # x = self.ln_f(x) + # logits = self.head(x) + # return logits -# def forward(self, input_ids, attention_mask=None): -# B, T = input_ids.size() -# pos = torch.arange(0, T, device=input_ids.device).unsqueeze(0) -# -# x = self.token_emb(input_ids) + self.pos_emb(pos) -# -# for block in self.blocks: -# x = block(x, attention_mask) -# -# x = self.ln_f(x) -# logits = self.head(x) -# return logits - - - def generate(self, - x: torch.Tensor, - max_new_tokens: int, + def generate( + self, + x: torch.Tensor, + max_new_tokens: int, do_sample: bool, temperature: float = 1.0, top_k: int = None, top_p: float = None, attention_mask: torch.Tensor = None, # Добавляем для совместимости с HF - **kwargs # Игнорируем остальные параметры + **kwargs, # Игнорируем остальные параметры ) -> torch.Tensor: """Авторегрессивная генерация текста. - + Параметры: x: Входной тензор с индексами токенов формы [batch_size, seq_len], где batch_size - размер батча, seq_len - длина последовательности. @@ -157,9 +164,9 @@ class GPT(BaseModel): - Гарантируется, что хотя бы один токен остаётся (даже если его вероятность > top_p) - None: отключено (по умолчанию) - Должен быть в диапазоне (0, 1] - + Возвращает: - torch.Tensor: Тензор с расширенной последовательностью токенов формы + torch.Tensor: Тензор с расширенной последовательностью токенов формы [batch_size, seq_len + max_new_tokens] Исключения: @@ -172,7 +179,7 @@ class GPT(BaseModel): Примеры: >>> # Жадная генерация >>> output = model.generate(input_ids, max_new_tokens=10, do_sample=False) - >>> + >>> >>> # Вероятностная генерация с top-k >>> output = model.generate(input_ids, max_new_tokens=10, do_sample=True, top_k=50) >>> @@ -180,11 +187,11 @@ class GPT(BaseModel): >>> output = model.generate(input_ids, max_new_tokens=10, do_sample=True, top_p=0.9) >>> >>> # Комбинация температуры и top-k - >>> output = model.generate(input_ids, max_new_tokens=10, do_sample=True, + >>> output = model.generate(input_ids, max_new_tokens=10, do_sample=True, ... temperature=0.7, top_k=50) Примечания: - 1. Для детерминированных результатов в режиме сэмплирования + 1. Для детерминированных результатов в режиме сэмплирования зафиксируйте random seed (torch.manual_seed). 2. Температура влияет только на режим сэмплирования (do_sample=True). 3. Одновременное использование top_k и top_p запрещено. @@ -204,7 +211,7 @@ class GPT(BaseModel): Должна быть > 0 (по умолчанию: 1.0) Returns: - torch.Tensor: Тензор с расширенной последовательностью токенов формы + torch.Tensor: Тензор с расширенной последовательностью токенов формы [batch_size, seq_len + max_new_tokens] Raises: @@ -222,13 +229,13 @@ class GPT(BaseModel): >>> output = model.generate(input_ids, max_new_tokens=10, do_sample=True, temperature=1.5) Note: - Для детерминированных результатов в режиме сэмплирования + Для детерминированных результатов в режиме сэмплирования зафиксируйте random seed (torch.manual_seed). Температура влияет только на режим сэмплирования (do_sample=True). """ for _ in range(max_new_tokens): # 1. Обрезаем вход, если последовательность слишком длинная - x_cond = x[:, -self._max_seq_len:] + x_cond = x[:, -self._max_seq_len :] # 2. Передаем последовательность в метод forward класса GPT и полуаем логиты. logits = self.forward(x_cond) @@ -250,9 +257,14 @@ class GPT(BaseModel): vocab_size = logits_scaled.size(-1) # создаём маску: True, если токен НЕ в topk_indices - mask = torch.ones_like(logits_scaled, dtype=torch.bool if hasattr(torch, 'bool') else torch.uint8) - mask.scatter_(1, topk_indices, False if hasattr(torch, 'bool') else 0) # False там, где top-k индексы - masked_logits[mask] = float('-inf') + mask = torch.ones_like( + logits_scaled, + dtype=torch.bool if hasattr(torch, "bool") else torch.uint8, + ) + mask.scatter_( + 1, topk_indices, False if hasattr(torch, "bool") else 0 + ) # False там, где top-k индексы + masked_logits[mask] = float("-inf") logits_scaled = masked_logits @@ -260,36 +272,42 @@ class GPT(BaseModel): # 1. Применим softmax, чтобы получить вероятности: probs = F.softmax(logits_scaled, dim=-1) # [B, vocab_size] # 2. Отсортируем токены по убыванию вероятностей: - sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=-1) + sorted_probs, sorted_indices = torch.sort( + probs, descending=True, dim=-1 + ) # 3. Посчитаем кумулятивную сумму вероятностей: cum_probs = torch.cumsum(sorted_probs, dim=-1) # [B, vocab_size] # 4. Определим маску: оставить токены, пока сумма < top_p - sorted_mask = (cum_probs <= top_p) # [B, vocab_size] + sorted_mask = cum_probs <= top_p # [B, vocab_size] # Гарантируем, что хотя бы первый токен останется sorted_mask[:, 0] = True # 5. Преобразуем маску обратно в оригинальный порядок: # Создаём полную маску из False - mask = torch.zeros_like(probs, dtype=torch.bool if hasattr(torch, 'bool') else torch.uint8) + mask = torch.zeros_like( + probs, dtype=torch.bool if hasattr(torch, "bool") else torch.uint8 + ) # Устанавливаем True в местах нужных токенов mask.scatter_(dim=1, index=sorted_indices, src=sorted_mask) # 6. Зануляем логиты токенов вне топ-p: - logits_scaled[~mask] = float('-inf') + logits_scaled[~mask] = float("-inf") # 4. Применяем Softmax probs = F.softmax(logits_scaled, dim=-1) # [batch_size, vocab_size] - if do_sample == True: # 5. Если do_sample равен True, то отбираем токен случайно с помощью torch.multinomial next_token = torch.multinomial(probs, num_samples=1) # [batch_size, 1] else: # 5. Если do_sample равен False, то выбираем токен с максимальной вероятностью - next_token = torch.argmax(probs, dim=-1, keepdim=True) # [batch_size, 1] - + next_token = torch.argmax( + probs, dim=-1, keepdim=True + ) # [batch_size, 1] + # 6. Добавляем его к последовательности x = torch.cat([x, next_token], dim=1) # [batch_size, seq_len+1] return x + # def generate(self, input_ids, max_length=50): # for _ in range(max_length): # logits = self.forward(input_ids) diff --git a/llm/src/llm/models/gpt/gpt2.py b/llm/src/llm/models/gpt/gpt2.py index 54c60a0..0b8d532 100644 --- a/llm/src/llm/models/gpt/gpt2.py +++ b/llm/src/llm/models/gpt/gpt2.py @@ -27,6 +27,7 @@ from llm.core.positional_embeddings import PositionalEmbeddings from llm.core.cached_decoder import CachedDecoder from llm.core.feed_forward import FeedForward + class GPT2(BaseModel): """ GPT2 — автогерессивная языковая модель, архитектура Transformer, предложенная OpenAI. @@ -35,7 +36,7 @@ class GPT2(BaseModel): - Масштабируемый автогерессивный трансформер для предсказания токенов слева направо. - Главное отличие от классической GPT: порядок layer normalization ПЕРЕД attention и FFN. - Используется GELU, efficient KV-cache, несет наследие классической GPT, но делает архитектуру глубже/шире. - + Args: config (dict): параметры архитектуры (vocab_size, embed_dim, num_heads, num_layers, max_position_embeddings, dropout) @@ -44,37 +45,43 @@ class GPT2(BaseModel): >>> logits = model(input_ids) >>> out = model.generate(input_ids, max_length=20) """ + def __init__(self, config): super().__init__(config) # Инициализация слоев self._max_seq_len = config["max_position_embeddings"] self._token_embeddings = TokenEmbeddings( - vocab_size=config["vocab_size"], - emb_size=config["embed_dim"] + vocab_size=config["vocab_size"], emb_size=config["embed_dim"] ) self._position_embeddings = PositionalEmbeddings( - max_seq_len=config["max_position_embeddings"], - emb_size=config["embed_dim"] + max_seq_len=config["max_position_embeddings"], emb_size=config["embed_dim"] ) self._dropout = nn.Dropout(config["dropout"]) # head_size = emb_size // num_heads - self._decoders = nn.ModuleList([CachedDecoder( - num_heads=config["num_heads"], - emb_size=config["embed_dim"], - head_size=config["embed_dim"] // config["num_heads"], - feed_forward_layer=FeedForward( - emb_size=config["embed_dim"], - dropout=config["dropout"], - activation="gelu" - ), - max_seq_len=config["max_position_embeddings"], - dropout=config["dropout"] - ) for _ in range(config["num_layers"])]) + self._decoders = nn.ModuleList( + [ + CachedDecoder( + num_heads=config["num_heads"], + emb_size=config["embed_dim"], + head_size=config["embed_dim"] // config["num_heads"], + feed_forward_layer=FeedForward( + emb_size=config["embed_dim"], + dropout=config["dropout"], + activation="gelu", + ), + max_seq_len=config["max_position_embeddings"], + dropout=config["dropout"], + ) + for _ in range(config["num_layers"]) + ] + ) self._norm = nn.LayerNorm(config["embed_dim"]) self._linear = nn.Linear(config["embed_dim"], config["vocab_size"]) - def forward(self, x: torch.Tensor, use_cache: bool = True, cache: list = None) -> tuple: + def forward( + self, x: torch.Tensor, use_cache: bool = True, cache: list = None + ) -> tuple: """ Прямой проход GPT2: - Все слои работают как autoregressive transformer (masked self-attention). @@ -91,9 +98,10 @@ class GPT2(BaseModel): """ # Проверка длины последовательности (только при отсутствии кэша) if cache is None and x.size(1) > self._max_seq_len: - raise ValueError(f"Длина последовательности {x.size(1)} превышает максимальную {self.max_seq_len}") - - + raise ValueError( + f"Длина последовательности {x.size(1)} превышает максимальную {self.max_seq_len}" + ) + # Вычисление start_pos из кэша (если кэш передан) if cache is not None: # При кэше обрабатываем только один токен (последний) @@ -111,11 +119,15 @@ class GPT2(BaseModel): # Эмбеддинги токенов и позиций tok_out = self._token_embeddings(x) # [batch, seq_len, emb_size] - pos_out = self._position_embeddings(seq_len, start_pos=start_pos) # [seq_len, emb_size] - + pos_out = self._position_embeddings( + seq_len, start_pos=start_pos + ) # [seq_len, emb_size] + # Комбинирование - out = self._dropout(tok_out + pos_out.unsqueeze(0)) # [batch, seq_len, emb_size] - + out = self._dropout( + tok_out + pos_out.unsqueeze(0) + ) # [batch, seq_len, emb_size] + # Стек декодеров с передачей кэша new_cache = [] for i, decoder in enumerate(self._decoders): @@ -131,21 +143,22 @@ class GPT2(BaseModel): out = self._norm(out) logits = self._linear(out) - + # Возвращаем результат с учетом use_cache if use_cache: return (logits, new_cache) else: return (logits, None) - def generate(self, - x: torch.Tensor, - max_new_tokens: int, + def generate( + self, + x: torch.Tensor, + max_new_tokens: int, do_sample: bool, temperature: float = 1.0, top_k: int = None, top_p: float = None, - use_cache: bool = True + use_cache: bool = True, ) -> torch.Tensor: """ Генерация текста с использованием autoregressive трансформера (GPT2). @@ -174,10 +187,10 @@ class GPT2(BaseModel): else: # Первая итерация или кэш отключен - передаем всю последовательность x_input = x - + # Прямой проход с кэшем logits, new_cache = self.forward(x_input, use_cache=use_cache, cache=cache) - + # Обновляем кэш для следующей итерации if use_cache: cache = new_cache @@ -200,7 +213,7 @@ class GPT2(BaseModel): # создаём маску: 1, если токен НЕ в topk_indices mask = torch.ones_like(logits_scaled, dtype=torch.uint8) mask.scatter_(1, topk_indices, 0) # 0 там, где top-k индексы - masked_logits[mask.byte()] = float('-inf') + masked_logits[mask.byte()] = float("-inf") logits_scaled = masked_logits @@ -208,7 +221,9 @@ class GPT2(BaseModel): # 1. Применим softmax, чтобы получить вероятности: probs = F.softmax(logits_scaled, dim=-1) # [B, vocab_size] # 2. Отсортируем токены по убыванию вероятностей: - sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=-1) + sorted_probs, sorted_indices = torch.sort( + probs, descending=True, dim=-1 + ) # 3. Посчитаем кумулятивную сумму вероятностей: cum_probs = torch.cumsum(sorted_probs, dim=-1) # [B, vocab_size] # 4. Определим маску: оставить токены, пока сумма < top_p @@ -221,23 +236,24 @@ class GPT2(BaseModel): # Устанавливаем 1 в местах нужных токенов mask.scatter_(dim=1, index=sorted_indices, src=sorted_mask) # 6. Зануляем логиты токенов вне топ-p: - logits_scaled[~mask] = float('-inf') + logits_scaled[~mask] = float("-inf") # 4. Применяем Softmax probs = F.softmax(logits_scaled, dim=-1) # [batch_size, vocab_size] - if do_sample == True: # 5. Если do_sample равен True, то отбираем токен случайно с помощью torch.multinomial next_token = torch.multinomial(probs, num_samples=1) # [batch_size, 1] else: # 5. Если do_sample равен False, то выбираем токен с максимальной вероятностью - next_token = torch.argmax(probs, dim=-1, keepdim=True) # [batch_size, 1] - + next_token = torch.argmax( + probs, dim=-1, keepdim=True + ) # [batch_size, 1] + # 6. Добавляем его к последовательности x = torch.cat([x, next_token], dim=1) # [batch_size, seq_len+1] return x @property def max_seq_len(self) -> int: - return self._max_seq_len \ No newline at end of file + return self._max_seq_len diff --git a/llm/src/llm/models/llama/llama.py b/llm/src/llm/models/llama/llama.py index 7e63221..78486fb 100644 --- a/llm/src/llm/models/llama/llama.py +++ b/llm/src/llm/models/llama/llama.py @@ -10,7 +10,6 @@ from llm.core.rope import RoPE from llm.core.cached_decoder import CachedDecoder - class Llama(BaseModel): """ LLaMA (Large Language Model Meta AI) — высокоэффективная масштабируемая языковая модель, разработанная Meta AI Research. @@ -29,38 +28,45 @@ class Llama(BaseModel): >>> logits, cache = model(input_ids, use_cache=True) >>> out = model.generate(input_ids, max_new_tokens=20) """ - def __init__(self,config): + + def __init__(self, config): super().__init__(config) # Инициализация слоев self._max_seq_len = config["max_position_embeddings"] self._token_embeddings = TokenEmbeddings( - vocab_size=config["vocab_size"], - emb_size=config["embed_dim"] + vocab_size=config["vocab_size"], emb_size=config["embed_dim"] ) self._position_embeddings = RoPE( head_size=config["embed_dim"] // config["num_heads"], - max_seq_len=config["max_position_embeddings"] + max_seq_len=config["max_position_embeddings"], ) self._dropout = nn.Dropout(config["dropout"]) - self._decoders = nn.ModuleList([CachedDecoder( - norm_layer=RMSNorm, - num_heads=config["num_heads"], - emb_size=config["embed_dim"], - head_size=config["embed_dim"] // config["num_heads"], - feed_forward_layer=SwiGLU( - emb_size=config["embed_dim"], - dropout=config["dropout"], - ), - max_seq_len=config["max_position_embeddings"], - rope=self._position_embeddings, - dropout=config["dropout"], - ) for _ in range(config["num_layers"])]) + self._decoders = nn.ModuleList( + [ + CachedDecoder( + norm_layer=RMSNorm, + num_heads=config["num_heads"], + emb_size=config["embed_dim"], + head_size=config["embed_dim"] // config["num_heads"], + feed_forward_layer=SwiGLU( + emb_size=config["embed_dim"], + dropout=config["dropout"], + ), + max_seq_len=config["max_position_embeddings"], + rope=self._position_embeddings, + dropout=config["dropout"], + ) + for _ in range(config["num_layers"]) + ] + ) self._norm = RMSNorm(config["embed_dim"]) self._linear = nn.Linear(config["embed_dim"], config["vocab_size"]) - def forward(self, x: torch.Tensor, use_cache: bool = True, cache: list = None) -> tuple: + def forward( + self, x: torch.Tensor, use_cache: bool = True, cache: list = None + ) -> tuple: """ Прямой проход через LLaMA (inference/train): авторегрессионное предсказание токенов. @@ -76,11 +82,12 @@ class Llama(BaseModel): """ # Проверка длины последовательности (только при отсутствии кэша) if cache is None and x.size(1) > self._max_seq_len: - raise ValueError(f"Длина последовательности {x.size(1)} превышает максимальную {self.max_seq_len}") - - + raise ValueError( + f"Длина последовательности {x.size(1)} превышает максимальную {self.max_seq_len}" + ) + # Вычисление start_pos из кэша (если кэш передан) - #if cache is not None: + # if cache is not None: # # При кэше обрабатываем только один токен (последний) # seq_len = 1 # # Вычисляем start_pos из самого нижнего уровня кэша @@ -89,18 +96,18 @@ class Llama(BaseModel): # start_pos = key_cache.size(1) # cache_len # else: # start_pos = 0 - #else: + # else: # # Без кэша работаем как раньше # start_pos = 0 # seq_len = x.size(1) # Эмбеддинги токенов и позиций tok_out = self._token_embeddings(x) # [batch, seq_len, emb_size] - #pos_out = self._position_embeddings(x) # [batch, seq_len, emb_size] - + # pos_out = self._position_embeddings(x) # [batch, seq_len, emb_size] + # Комбинирование out = self._dropout(tok_out) # [batch, seq_len, emb_size] - + # Стек декодеров с передачей кэша new_cache = [] for i, decoder in enumerate(self._decoders): @@ -116,42 +123,43 @@ class Llama(BaseModel): out = self._norm(out) logits = self._linear(out) - + # Возвращаем результат с учетом use_cache if use_cache: return (logits, new_cache) else: return (logits, None) - def generate(self, - x: torch.Tensor, - max_new_tokens: int, + def generate( + self, + x: torch.Tensor, + max_new_tokens: int, do_sample: bool, temperature: float = 1.0, top_k: int = None, top_p: float = None, - use_cache: bool = True + use_cache: bool = True, ) -> torch.Tensor: """ - Генерация текста c помощью LLaMA (autoregressive Transformer). - Поддерживается: - - greedy и вероятностное сэмплирование (top-k, top-p, temperature) - - кэш attention для ускорения генерации длинных последовательностей + Генерация текста c помощью LLaMA (autoregressive Transformer). + Поддерживается: + - greedy и вероятностное сэмплирование (top-k, top-p, temperature) + - кэш attention для ускорения генерации длинных последовательностей - Args: - x (Tensor[int]): начальная последовательность [batch, seq_len] - max_new_tokens (int): сколько новых токенов сгенерировать - do_sample (bool): использовать стохастику (True) или жадный выбор (False) - temperature (float): масштаб для softmax (важно для sampling) - top_k (int|None): ограничение на количество кандидатов (top-k sampling) - top_p (float|None): nucleus sampling - use_cache (bool): ускоряет autoregressive при длинной генерации - Returns: - output (Tensor[int]): [batch, seq_len + max_new_tokens] - Пример: - >>> prompt = tokenizer.encode('Meta AI', return_tensors="pt") - >>> generated = model.generate(prompt, max_new_tokens=30, do_sample=True) - >>> print(tokenizer.decode(generated[0])) + Args: + x (Tensor[int]): начальная последовательность [batch, seq_len] + max_new_tokens (int): сколько новых токенов сгенерировать + do_sample (bool): использовать стохастику (True) или жадный выбор (False) + temperature (float): масштаб для softmax (важно для sampling) + top_k (int|None): ограничение на количество кандидатов (top-k sampling) + top_p (float|None): nucleus sampling + use_cache (bool): ускоряет autoregressive при длинной генерации + Returns: + output (Tensor[int]): [batch, seq_len + max_new_tokens] + Пример: + >>> prompt = tokenizer.encode('Meta AI', return_tensors="pt") + >>> generated = model.generate(prompt, max_new_tokens=30, do_sample=True) + >>> print(tokenizer.decode(generated[0])) """ cache = None @@ -162,10 +170,10 @@ class Llama(BaseModel): else: # Первая итерация или кэш отключен - передаем всю последовательность x_input = x - + # Прямой проход с кэшем logits, new_cache = self.forward(x_input, use_cache=use_cache, cache=cache) - + # Обновляем кэш для следующей итерации if use_cache: cache = new_cache @@ -188,7 +196,7 @@ class Llama(BaseModel): # создаём маску: 1, если токен НЕ в topk_indices mask = torch.ones_like(logits_scaled, dtype=torch.uint8) mask.scatter_(1, topk_indices, 0) # 0 там, где top-k индексы - masked_logits[mask.byte()] = float('-inf') + masked_logits[mask.byte()] = float("-inf") logits_scaled = masked_logits @@ -196,7 +204,9 @@ class Llama(BaseModel): # 1. Применим softmax, чтобы получить вероятности: probs = F.softmax(logits_scaled, dim=-1) # [B, vocab_size] # 2. Отсортируем токены по убыванию вероятностей: - sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=-1) + sorted_probs, sorted_indices = torch.sort( + probs, descending=True, dim=-1 + ) # 3. Посчитаем кумулятивную сумму вероятностей: cum_probs = torch.cumsum(sorted_probs, dim=-1) # [B, vocab_size] # 4. Определим маску: оставить токены, пока сумма < top_p @@ -209,25 +219,24 @@ class Llama(BaseModel): # Устанавливаем 1 в местах нужных токенов mask.scatter_(dim=1, index=sorted_indices, src=sorted_mask) # 6. Зануляем логиты токенов вне топ-p: - logits_scaled[~mask] = float('-inf') + logits_scaled[~mask] = float("-inf") # 4. Применяем Softmax probs = F.softmax(logits_scaled, dim=-1) # [batch_size, vocab_size] - if do_sample == True: # 5. Если do_sample равен True, то отбираем токен случайно с помощью torch.multinomial next_token = torch.multinomial(probs, num_samples=1) # [batch_size, 1] else: # 5. Если do_sample равен False, то выбираем токен с максимальной вероятностью - next_token = torch.argmax(probs, dim=-1, keepdim=True) # [batch_size, 1] - + next_token = torch.argmax( + probs, dim=-1, keepdim=True + ) # [batch_size, 1] + # 6. Добавляем его к последовательности x = torch.cat([x, next_token], dim=1) # [batch_size, seq_len+1] return x - - @property def max_seq_len(self) -> int: - return self._max_seq_len \ No newline at end of file + return self._max_seq_len diff --git a/llm/src/llm/tokenizers/base_tokenizer.py b/llm/src/llm/tokenizers/base_tokenizer.py index 1b4c1da..b1b8367 100644 --- a/llm/src/llm/tokenizers/base_tokenizer.py +++ b/llm/src/llm/tokenizers/base_tokenizer.py @@ -10,92 +10,94 @@ import json class BaseTokenizer(ABC): """ Абстрактный базовый класс для всех токенизаторов. - + Определяет общий интерфейс для токенизации текста. """ - + def __init__(self): self.vocab: Dict[str, int] = {} self.inverse_vocab: Dict[int, str] = {} self.vocab_size: int = 0 - + # Специальные токены self.pad_token = "" self.unk_token = "" self.bos_token = "" self.eos_token = "" - + self.pad_token_id: Optional[int] = None self.unk_token_id: Optional[int] = None self.bos_token_id: Optional[int] = None self.eos_token_id: Optional[int] = None - + @abstractmethod def train(self, texts: List[str], vocab_size: int = 1000, **kwargs): """ Обучение токенизатора на текстах. - + Args: texts: Список текстов для обучения vocab_size: Желаемый размер словаря **kwargs: Дополнительные параметры обучения """ pass - + @abstractmethod def encode(self, text: str, **kwargs) -> List[int]: """ Кодирование текста в последовательность токенов. - + Args: text: Входной текст **kwargs: Дополнительные параметры кодирования - + Returns: List[int]: Список идентификаторов токенов """ pass - + @abstractmethod def decode(self, tokens: List[int], **kwargs) -> str: """ Декодирование последовательности токенов в текст. - + Args: tokens: Список идентификаторов токенов **kwargs: Дополнительные параметры декодирования - + Returns: str: Декодированный текст """ pass - + def tokenize(self, text: str, **kwargs) -> List[str]: """ Токенизация текста в список строковых токенов. - + Args: text: Входной текст **kwargs: Дополнительные параметры - + Returns: List[str]: Список токенов """ token_ids = self.encode(text, **kwargs) - return [self.inverse_vocab.get(token_id, self.unk_token) for token_id in token_ids] - + return [ + self.inverse_vocab.get(token_id, self.unk_token) for token_id in token_ids + ] + def get_vocab(self) -> Dict[str, int]: """Возвращает словарь токенизатора.""" return self.vocab.copy() - + def get_vocab_size(self) -> int: """Возвращает размер словаря.""" return self.vocab_size - + def add_special_tokens(self, special_tokens: List[str]): """ Добавляет специальные токены в словарь. - + Args: special_tokens: Список специальных токенов """ @@ -105,70 +107,70 @@ class BaseTokenizer(ABC): self.vocab[token] = token_id self.inverse_vocab[token_id] = token self.vocab_size += 1 - + # Обновляем ID специальных токенов self.pad_token_id = self.vocab.get(self.pad_token) self.unk_token_id = self.vocab.get(self.unk_token) self.bos_token_id = self.vocab.get(self.bos_token) self.eos_token_id = self.vocab.get(self.eos_token) - + def save(self, filepath: str): """ Сохраняет токенизатор в файл. - + Args: filepath: Путь для сохранения """ config = { - 'vocab': self.vocab, - 'vocab_size': self.vocab_size, - 'pad_token': self.pad_token, - 'unk_token': self.unk_token, - 'bos_token': self.bos_token, - 'eos_token': self.eos_token, - 'tokenizer_type': self.__class__.__name__ + "vocab": self.vocab, + "vocab_size": self.vocab_size, + "pad_token": self.pad_token, + "unk_token": self.unk_token, + "bos_token": self.bos_token, + "eos_token": self.eos_token, + "tokenizer_type": self.__class__.__name__, } - - with open(filepath, 'w', encoding='utf-8') as f: + + with open(filepath, "w", encoding="utf-8") as f: json.dump(config, f, ensure_ascii=False, indent=2) - + @classmethod def load(cls, filepath: str): """ Загружает токенизатор из файла. - + Args: filepath: Путь к файлу - + Returns: BaseTokenizer: Загруженный токенизатор """ - with open(filepath, 'r', encoding='utf-8') as f: + with open(filepath, "r", encoding="utf-8") as f: config = json.load(f) - + # Создаем экземпляр токенизатора tokenizer = cls() - tokenizer.vocab = config['vocab'] - tokenizer.vocab_size = config['vocab_size'] - tokenizer.pad_token = config['pad_token'] - tokenizer.unk_token = config['unk_token'] - tokenizer.bos_token = config['bos_token'] - tokenizer.eos_token = config['eos_token'] - + tokenizer.vocab = config["vocab"] + tokenizer.vocab_size = config["vocab_size"] + tokenizer.pad_token = config["pad_token"] + tokenizer.unk_token = config["unk_token"] + tokenizer.bos_token = config["bos_token"] + tokenizer.eos_token = config["eos_token"] + # Создаем обратный словарь tokenizer.inverse_vocab = {v: k for k, v in tokenizer.vocab.items()} - + # Обновляем ID специальных токенов tokenizer.pad_token_id = tokenizer.vocab.get(tokenizer.pad_token) tokenizer.unk_token_id = tokenizer.vocab.get(tokenizer.unk_token) tokenizer.bos_token_id = tokenizer.vocab.get(tokenizer.bos_token) tokenizer.eos_token_id = tokenizer.vocab.get(tokenizer.eos_token) - + return tokenizer - + def __len__(self) -> int: """Возвращает размер словаря.""" return self.vocab_size - + def __repr__(self) -> str: return f"{self.__class__.__name__}(vocab_size={self.vocab_size})" diff --git a/llm/src/llm/tokenizers/bpe_tokenizer copy.py b/llm/src/llm/tokenizers/bpe_tokenizer copy.py index 0754560..95daf41 100644 --- a/llm/src/llm/tokenizers/bpe_tokenizer copy.py +++ b/llm/src/llm/tokenizers/bpe_tokenizer copy.py @@ -13,26 +13,26 @@ from .base_tokenizer import BaseTokenizer class BPETokenizer(BaseTokenizer): """ BPE токенизатор для обработки текста. - + Реализует алгоритм Byte Pair Encoding для создания субсловных токенов. - + Примеры использования: >>> tokenizer = BPETokenizer() >>> tokenizer.train(["пример текста для обучения"], vocab_size=1000) >>> tokens = tokenizer.encode("новый текст") >>> text = tokenizer.decode(tokens) """ - + def __init__(self): super().__init__() self.merges: Dict[Tuple[str, str], int] = {} self.pattern = r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""" self.compiled_pattern = re.compile(self.pattern, re.UNICODE) - + def train(self, texts: List[str], vocab_size: int = 1000, **kwargs): """ Обучение BPE токенизатора на текстах. - + Args: texts: Список текстов для обучения vocab_size: Желаемый размер словаря @@ -42,37 +42,40 @@ class BPETokenizer(BaseTokenizer): """ # Инициализация базового словаря self._initialize_vocab() - + # Добавляем специальные токены если указаны - special_tokens = kwargs.get('special_tokens', [self.pad_token, self.unk_token, self.bos_token, self.eos_token]) + special_tokens = kwargs.get( + "special_tokens", + [self.pad_token, self.unk_token, self.bos_token, self.eos_token], + ) self.add_special_tokens(special_tokens) - + # Предобработка текстов words = self._preprocess_texts(texts) - + # Получаем начальные токены vocab = self._get_initial_vocab(words) - + # Выполняем BPE мерджи - self._perform_merges(vocab, vocab_size, kwargs.get('min_frequency', 2)) - + self._perform_merges(vocab, vocab_size, kwargs.get("min_frequency", 2)) + # Строим финальный словарь self._build_final_vocab() - + def _initialize_vocab(self): """Инициализирует базовый словарь.""" self.vocab.clear() self.inverse_vocab.clear() self.merges.clear() self.vocab_size = 0 - + def _preprocess_texts(self, texts: List[str]) -> List[List[str]]: """ Предобработка текстов для обучения. - + Args: texts: Список текстов - + Returns: List[List[str]]: Предобработанные слова """ @@ -84,14 +87,14 @@ class BPETokenizer(BaseTokenizer): tokens = self.compiled_pattern.findall(text) words.append(tokens) return words - + def _get_initial_vocab(self, words: List[List[str]]) -> Dict[str, int]: """ Создает начальный словарь из символов. - + Args: words: Список токенизированных текстов - + Returns: Dict[str, int]: Начальный словарь частот """ @@ -99,43 +102,45 @@ class BPETokenizer(BaseTokenizer): for word_list in words: for word in word_list: # Разбиваем слово на символы и добавляем специальный символ конца слова - chars = list(word) + [''] - vocab.update([''.join(chars[i:i+1]) for i in range(len(chars))]) + chars = list(word) + [""] + vocab.update(["".join(chars[i : i + 1]) for i in range(len(chars))]) return vocab - - def _perform_merges(self, vocab: Dict[str, int], target_vocab_size: int, min_frequency: int): + + def _perform_merges( + self, vocab: Dict[str, int], target_vocab_size: int, min_frequency: int + ): """ Выполняет BPE мерджи до достижения целевого размера словаря. - + Args: vocab: Начальный словарь target_vocab_size: Целевой размер словаря min_frequency: Минимальная частота для мерджа """ current_vocab_size = len(vocab) + len(self.vocab) - + while current_vocab_size < target_vocab_size: # Находим наиболее частую пару pairs = self._get_stats(vocab) if not pairs: break - + best_pair = max(pairs, key=pairs.get) if pairs[best_pair] < min_frequency: break - + # Выполняем мердж vocab = self._merge_vocab(vocab, best_pair) self.merges[best_pair] = len(self.merges) current_vocab_size += 1 - + def _get_stats(self, vocab: Dict[str, int]) -> Dict[Tuple[str, str], int]: """ Собирает статистику по парам символов. - + Args: vocab: Словарь токенов - + Returns: Dict[Tuple[str, str], int]: Частоты пар """ @@ -145,107 +150,113 @@ class BPETokenizer(BaseTokenizer): for i in range(len(symbols) - 1): pairs[symbols[i], symbols[i + 1]] += freq return pairs - - def _merge_vocab(self, vocab: Dict[str, int], pair: Tuple[str, str]) -> Dict[str, int]: + + def _merge_vocab( + self, vocab: Dict[str, int], pair: Tuple[str, str] + ) -> Dict[str, int]: """ Объединяет пару символов в словаре. - + Args: vocab: Исходный словарь pair: Пара для объединения - + Returns: Dict[str, int]: Обновленный словарь """ new_vocab = {} - bigram = re.compile(r'(? List[int]: """ Кодирует текст в последовательность токенов. - + Args: text: Входной текст **kwargs: Дополнительные параметры - add_special_tokens: Добавлять специальные токены - + Returns: List[int]: Список идентификаторов токенов """ - add_special_tokens = kwargs.get('add_special_tokens', False) - + add_special_tokens = kwargs.get("add_special_tokens", False) + # Токенизация текста tokens = self.compiled_pattern.findall(text) - + # Применяем BPE к каждому токену bpe_tokens = [] for token in tokens: # Преобразуем токен в BPE представление bpe_token = self._apply_bpe(token) bpe_tokens.extend(bpe_token) - + # Конвертируем в ID token_ids = [] for token in bpe_tokens: token_id = self.vocab.get(token, self.unk_token_id) if token_id is not None: token_ids.append(token_id) - + # Добавляем специальные токены если нужно if add_special_tokens: if self.bos_token_id is not None: token_ids.insert(0, self.bos_token_id) if self.eos_token_id is not None: token_ids.append(self.eos_token_id) - + return token_ids - + def _apply_bpe(self, token: str) -> List[str]: """ Применяет BPE к одному токену. - + Args: token: Входной токен - + Returns: List[str]: Список BPE токенов """ # Простая реализация - в реальной реализации нужно применять обученные мерджи - word = token + '' - tokens = [word[i:i+1] for i in range(len(word))] - + word = token + "" + tokens = [word[i : i + 1] for i in range(len(word))] + # Применяем мерджи (упрощенная версия) # В полной реализации нужно применять все обученные мерджи for pair in self.merges: @@ -256,109 +267,114 @@ class BPETokenizer(BaseTokenizer): del tokens[i + 1] else: i += 1 - + return tokens - + def decode(self, tokens: List[int], **kwargs) -> str: """ Декодирует последовательность токенов в текст. - + Args: tokens: Список идентификаторов токенов **kwargs: Дополнительные параметры - skip_special_tokens: Пропускать специальные токены - + Returns: str: Декодированный текст """ - skip_special_tokens = kwargs.get('skip_special_tokens', True) - + skip_special_tokens = kwargs.get("skip_special_tokens", True) + # Конвертируем ID в токены token_strings = [] for token_id in tokens: token = self.inverse_vocab.get(token_id, self.unk_token) - + # Пропускаем специальные токены если нужно - if skip_special_tokens and token in [self.pad_token, self.unk_token, self.bos_token, self.eos_token]: + if skip_special_tokens and token in [ + self.pad_token, + self.unk_token, + self.bos_token, + self.eos_token, + ]: continue - + token_strings.append(token) - + # Объединяем токены в текст - text = ''.join(token_strings) - + text = "".join(token_strings) + # Убираем маркер конца слова - text = text.replace('', ' ') - + text = text.replace("", " ") + return text.strip() - + def save(self, filepath: str): """ Сохраняет BPE токенизатор в файл. - + Args: filepath: Путь для сохранения """ import json - + config = { - 'vocab': self.vocab, - 'merges': {f"{k[0]} {k[1]}": v for k, v in self.merges.items()}, - 'vocab_size': self.vocab_size, - 'pad_token': self.pad_token, - 'unk_token': self.unk_token, - 'bos_token': self.bos_token, - 'eos_token': self.eos_token, - 'pattern': self.pattern, - 'tokenizer_type': self.__class__.__name__ + "vocab": self.vocab, + "merges": {f"{k[0]} {k[1]}": v for k, v in self.merges.items()}, + "vocab_size": self.vocab_size, + "pad_token": self.pad_token, + "unk_token": self.unk_token, + "bos_token": self.bos_token, + "eos_token": self.eos_token, + "pattern": self.pattern, + "tokenizer_type": self.__class__.__name__, } - - with open(filepath, 'w', encoding='utf-8') as f: + + with open(filepath, "w", encoding="utf-8") as f: json.dump(config, f, ensure_ascii=False, indent=2) - + @classmethod def load(cls, filepath: str): """ Загружает BPE токенизатор из файла. - + Args: filepath: Путь к файлу - + Returns: BPETokenizer: Загруженный токенизатор """ import json - - with open(filepath, 'r', encoding='utf-8') as f: + + with open(filepath, "r", encoding="utf-8") as f: config = json.load(f) - + tokenizer = cls() - tokenizer.vocab = config['vocab'] - tokenizer.vocab_size = config['vocab_size'] - tokenizer.pad_token = config['pad_token'] - tokenizer.unk_token = config['unk_token'] - tokenizer.bos_token = config['bos_token'] - tokenizer.eos_token = config['eos_token'] - tokenizer.pattern = config.get('pattern', tokenizer.pattern) + tokenizer.vocab = config["vocab"] + tokenizer.vocab_size = config["vocab_size"] + tokenizer.pad_token = config["pad_token"] + tokenizer.unk_token = config["unk_token"] + tokenizer.bos_token = config["bos_token"] + tokenizer.eos_token = config["eos_token"] + tokenizer.pattern = config.get("pattern", tokenizer.pattern) tokenizer.compiled_pattern = re.compile(tokenizer.pattern, re.UNICODE) - + # Восстанавливаем мерджи - merges = config.get('merges', {}) + merges = config.get("merges", {}) tokenizer.merges = {} for k, v in merges.items(): parts = k.split() if len(parts) == 2: tokenizer.merges[(parts[0], parts[1])] = v - + # Создаем обратный словарь tokenizer.inverse_vocab = {v: k for k, v in tokenizer.vocab.items()} - + # Обновляем ID специальных токенов tokenizer.pad_token_id = tokenizer.vocab.get(tokenizer.pad_token) tokenizer.unk_token_id = tokenizer.vocab.get(tokenizer.unk_token) tokenizer.bos_token_id = tokenizer.vocab.get(tokenizer.bos_token) tokenizer.eos_token_id = tokenizer.vocab.get(tokenizer.eos_token) - + return tokenizer @@ -367,62 +383,72 @@ class SimpleBPETokenizer(BPETokenizer): """ Упрощенная версия BPE токенизатора для демонстрации. """ - + def train(self, texts: List[str], vocab_size: int = 1000, **kwargs): """Упрощенное обучение для демонстрации.""" # Инициализация базового словаря self._initialize_vocab() - + # Добавляем базовые токены - special_tokens = [self.pad_token, self.unk_token, self.bos_token, self.eos_token] + special_tokens = [ + self.pad_token, + self.unk_token, + self.bos_token, + self.eos_token, + ] self.add_special_tokens(special_tokens) - + # Простая реализация - собираем все символы all_chars = set() for text in texts: all_chars.update(text) - + # Добавляем символы в словарь for char in sorted(all_chars): if char not in self.vocab: self.vocab[char] = len(self.vocab) - + self.inverse_vocab = {v: k for k, v in self.vocab.items()} self.vocab_size = len(self.vocab) - + # Обновляем ID специальных токенов self.pad_token_id = self.vocab.get(self.pad_token) self.unk_token_id = self.vocab.get(self.unk_token) self.bos_token_id = self.vocab.get(self.bos_token) self.eos_token_id = self.vocab.get(self.eos_token) - + def encode(self, text: str, **kwargs) -> List[int]: """Упрощенное кодирование - разбиваем на символы.""" - add_special_tokens = kwargs.get('add_special_tokens', False) - + add_special_tokens = kwargs.get("add_special_tokens", False) + token_ids = [] for char in text: token_id = self.vocab.get(char, self.unk_token_id) if token_id is not None: token_ids.append(token_id) - + if add_special_tokens: if self.bos_token_id is not None: token_ids.insert(0, self.bos_token_id) if self.eos_token_id is not None: token_ids.append(self.eos_token_id) - + return token_ids - + def decode(self, tokens: List[int], **kwargs) -> str: """Упрощенное декодирование.""" - skip_special_tokens = kwargs.get('skip_special_tokens', True) - + skip_special_tokens = kwargs.get("skip_special_tokens", True) + chars = [] for token_id in tokens: char = self.inverse_vocab.get(token_id, self.unk_token) - if skip_special_tokens and char in [self.pad_token, self.unk_token, self.bos_token, self.eos_token]: + if skip_special_tokens and char in [ + self.pad_token, + self.unk_token, + self.bos_token, + self.eos_token, + ]: continue chars.append(char) - - return ''.join(chars) + + return "".join(chars) diff --git a/llm/src/llm/tokenizers/bpe_tokenizer.py b/llm/src/llm/tokenizers/bpe_tokenizer.py index 7e55e3e..221454b 100644 --- a/llm/src/llm/tokenizers/bpe_tokenizer.py +++ b/llm/src/llm/tokenizers/bpe_tokenizer.py @@ -11,26 +11,26 @@ from .base_tokenizer import BaseTokenizer class BPETokenizer(BaseTokenizer): """ BPE токенизатор для обработки текста. - + Реализует алгоритм Byte Pair Encoding для создания субсловных токенов. Использует вашу реализацию BPE. - + Примеры использования: >>> tokenizer = BPETokenizer() >>> tokenizer.train(["пример текста для обучения"], vocab_size=1000) >>> tokens = tokenizer.encode("новый текст") >>> text = tokenizer.decode(tokens) """ - + def __init__(self): super().__init__() self.merges: Dict[Tuple[str, str], int] = {} self.vocab_list: List[str] = [] - + def train(self, texts: List[str], vocab_size: int = 1000, **kwargs): """ Обучение BPE токенизатора на текстах. - + Args: texts: Список текстов для обучения vocab_size: Желаемый размер словаря @@ -39,7 +39,7 @@ class BPETokenizer(BaseTokenizer): """ # Объединяем все тексты в одну строку для обучения combined_text = " ".join(texts) - + # 1. Получаем уникальные токены (символы) unique_tokens = sorted(set(combined_text)) tokens = unique_tokens.copy() @@ -61,7 +61,10 @@ class BPETokenizer(BaseTokenizer): break # нет пар — выходим # Находим самую частую пару (в случае равенства — та, что встретилась первой) - most_frequent_pair = max(pair_freq.items(), key=lambda x: (x[1], -self._pair_first_index(sequence, x[0])))[0] + most_frequent_pair = max( + pair_freq.items(), + key=lambda x: (x[1], -self._pair_first_index(sequence, x[0])), + )[0] # Создаем новый токен new_token = most_frequent_pair[0] + most_frequent_pair[1] @@ -71,45 +74,51 @@ class BPETokenizer(BaseTokenizer): new_sequence = [] while i < len(sequence): - if i < len(sequence) - 1 and (sequence[i], sequence[i + 1]) == most_frequent_pair: + if ( + i < len(sequence) - 1 + and (sequence[i], sequence[i + 1]) == most_frequent_pair + ): new_sequence.append(new_token) i += 2 # пропускаем два символа — заменённую пару else: new_sequence.append(sequence[i]) i += 1 sequence = new_sequence - + # 4. Создаем словари self.vocab_list = tokens.copy() self.vocab = dict(zip(tokens, range(vocab_size))) self.inverse_vocab = dict(zip(range(vocab_size), tokens)) self.vocab_size = len(self.vocab) - + # Добавляем специальные токены если указаны - special_tokens = kwargs.get('special_tokens', [self.pad_token, self.unk_token, self.bos_token, self.eos_token]) + special_tokens = kwargs.get( + "special_tokens", + [self.pad_token, self.unk_token, self.bos_token, self.eos_token], + ) self.add_special_tokens(special_tokens) - + def _pair_first_index(self, sequence, pair): """Находит первый индекс пары в последовательности.""" for i in range(len(sequence) - 1): if (sequence[i], sequence[i + 1]) == pair: return i - return float('inf') # если пара не найдена (в теории не должно случиться) + return float("inf") # если пара не найдена (в теории не должно случиться) def encode(self, text: str, **kwargs) -> List[int]: """ Кодирует текст в последовательность токенов. - + Args: text: Входной текст **kwargs: Дополнительные параметры - add_special_tokens: Добавлять специальные токены - + Returns: List[int]: Список идентификаторов токенов """ - add_special_tokens = kwargs.get('add_special_tokens', False) - + add_special_tokens = kwargs.get("add_special_tokens", False) + # 1. Разбиваем текст на токены-символы sequence = list(text) # 2. Инициализация пустого списка токенов @@ -119,7 +128,9 @@ class BPETokenizer(BaseTokenizer): while i < len(text): # 3.1 Найти все токены в словаре, начинающиеся с text[i] start_char = text[i] - result = [token for token in self.vocab_list if token.startswith(start_char)] + result = [ + token for token in self.vocab_list if token.startswith(start_char) + ] # 3.2 Выбрать самый длинный подходящий токен find_token = self._find_max_matching_token(text[i:], result) if find_token is None: @@ -134,19 +145,19 @@ class BPETokenizer(BaseTokenizer): # 4. Заменить токены на их ID token_ids = self._tokens_to_ids(tokens) - + # Заменяем -1 на unk_token_id token_ids = [tid if tid != -1 else self.unk_token_id for tid in token_ids] - + # Добавляем специальные токены если нужно if add_special_tokens: if self.bos_token_id is not None: token_ids.insert(0, self.bos_token_id) if self.eos_token_id is not None: token_ids.append(self.eos_token_id) - + return token_ids - + def _find_max_matching_token(self, text: str, tokens: list) -> Optional[str]: """Находит самый длинный токен из списка, с которого начинается текст""" matching = [token for token in tokens if text.startswith(token)] @@ -161,33 +172,41 @@ class BPETokenizer(BaseTokenizer): else: ids.append(-1) # Специальное значение return ids - + def decode(self, tokens: List[int], **kwargs) -> str: """ Декодирует последовательность токенов в текст. - + Args: tokens: Список идентификаторов токенов **kwargs: Дополнительные параметры - skip_special_tokens: Пропускать специальные токены - + Returns: str: Декодированный текст """ - skip_special_tokens = kwargs.get('skip_special_tokens', True) - + skip_special_tokens = kwargs.get("skip_special_tokens", True) + # Фильтруем специальные токены если нужно if skip_special_tokens: - tokens = [tid for tid in tokens if tid not in [ - self.pad_token_id, self.unk_token_id, self.bos_token_id, self.eos_token_id - ]] - + tokens = [ + tid + for tid in tokens + if tid + not in [ + self.pad_token_id, + self.unk_token_id, + self.bos_token_id, + self.eos_token_id, + ] + ] + # Конвертируем ID в токены token_strings = self._ids_to_tokens(tokens) - + # Объединяем токены в текст - return ''.join(token_strings) - + return "".join(token_strings) + def _ids_to_tokens(self, ids: List[int]) -> List[str]: """Конвертирует список Ids в их tokens""" tokens = [] @@ -197,76 +216,76 @@ class BPETokenizer(BaseTokenizer): else: tokens.append(self.unk_token) # Специальное значение return tokens - + def save(self, filepath: str): """ Сохраняет токенизатор в файл. - + Args: filepath: Путь для сохранения """ import json - + # Преобразуем кортежи в строки для JSON сериализации merges_serializable = {f"{k[0]},{k[1]}": v for k, v in self.merges.items()} - + config = { - 'vocab': self.vocab, - 'vocab_size': self.vocab_size, - 'pad_token': self.pad_token, - 'unk_token': self.unk_token, - 'bos_token': self.bos_token, - 'eos_token': self.eos_token, - 'tokenizer_type': self.__class__.__name__, - 'merges': merges_serializable, - 'vocab_list': self.vocab_list + "vocab": self.vocab, + "vocab_size": self.vocab_size, + "pad_token": self.pad_token, + "unk_token": self.unk_token, + "bos_token": self.bos_token, + "eos_token": self.eos_token, + "tokenizer_type": self.__class__.__name__, + "merges": merges_serializable, + "vocab_list": self.vocab_list, } - - with open(filepath, 'w', encoding='utf-8') as f: + + with open(filepath, "w", encoding="utf-8") as f: json.dump(config, f, ensure_ascii=False, indent=2) - + @classmethod def load(cls, filepath: str): """ Загружает токенизатор из файла. - + Args: filepath: Путь к файлу - + Returns: BPETokenizer: Загруженный токенизатор """ import json - - with open(filepath, 'r', encoding='utf-8') as f: + + with open(filepath, "r", encoding="utf-8") as f: config = json.load(f) - + # Создаем экземпляр токенизатора tokenizer = cls() - tokenizer.vocab = config['vocab'] - tokenizer.vocab_size = config['vocab_size'] - tokenizer.pad_token = config['pad_token'] - tokenizer.unk_token = config['unk_token'] - tokenizer.bos_token = config['bos_token'] - tokenizer.eos_token = config['eos_token'] - tokenizer.vocab_list = config['vocab_list'] - + tokenizer.vocab = config["vocab"] + tokenizer.vocab_size = config["vocab_size"] + tokenizer.pad_token = config["pad_token"] + tokenizer.unk_token = config["unk_token"] + tokenizer.bos_token = config["bos_token"] + tokenizer.eos_token = config["eos_token"] + tokenizer.vocab_list = config["vocab_list"] + # Восстанавливаем кортежи из строк tokenizer.merges = {} - for k, v in config['merges'].items(): - parts = k.split(',') + for k, v in config["merges"].items(): + parts = k.split(",") if len(parts) == 2: tokenizer.merges[(parts[0], parts[1])] = v - + # Создаем обратный словарь tokenizer.inverse_vocab = {v: k for k, v in tokenizer.vocab.items()} - + # Обновляем ID специальных токенов tokenizer.pad_token_id = tokenizer.vocab.get(tokenizer.pad_token) tokenizer.unk_token_id = tokenizer.vocab.get(tokenizer.unk_token) tokenizer.bos_token_id = tokenizer.vocab.get(tokenizer.bos_token) tokenizer.eos_token_id = tokenizer.vocab.get(tokenizer.eos_token) - + return tokenizer @@ -275,4 +294,5 @@ class SimpleBPETokenizer(BPETokenizer): Упрощенная версия BPE токенизатора для демонстрации. Наследует вашу реализацию, но может быть упрощена при необходимости. """ + pass diff --git a/llm/src/llm/training/dataset.py b/llm/src/llm/training/dataset.py index e3cbd0e..8877637 100644 --- a/llm/src/llm/training/dataset.py +++ b/llm/src/llm/training/dataset.py @@ -12,7 +12,7 @@ class TextDataset(Dataset): def __init__(self, texts: List[str], tokenizer: Any, block_size: int = 128): """ Инициализация датасета. - + Args: texts: Список текстов для обучения tokenizer: Токенизатор с методами encode/decode @@ -25,15 +25,15 @@ class TextDataset(Dataset): for text in texts: # Кодируем текст в токены input_ids = tokenizer.encode(text, add_special_tokens=False) - + # Обрезаем или дополняем до нужной длины if len(input_ids) > block_size: input_ids = input_ids[:block_size] else: # Дополняем pad_token_id - pad_token_id = getattr(tokenizer, 'pad_token_id', 0) + pad_token_id = getattr(tokenizer, "pad_token_id", 0) input_ids = input_ids + [pad_token_id] * (block_size - len(input_ids)) - + self.examples.append(input_ids) def __len__(self): @@ -50,33 +50,35 @@ class StreamingTextDataset(Dataset): Датасет для потоковой обработки больших текстов. Токенизация происходит на лету, что экономит память. """ - + def __init__(self, texts: List[str], tokenizer: Any, block_size: int = 128): self.texts = texts self.tokenizer = tokenizer self.block_size = block_size - + # Получаем pad_token_id из токенизатора - self.pad_token_id = getattr(tokenizer, 'pad_token_id', 0) + self.pad_token_id = getattr(tokenizer, "pad_token_id", 0) def __len__(self): return len(self.texts) def __getitem__(self, idx): text = self.texts[idx] - + # Токенизация на лету input_ids = self.tokenizer.encode(text, add_special_tokens=False) - + # Обрезаем или дополняем до нужной длины if len(input_ids) > self.block_size: - input_ids = input_ids[:self.block_size] + input_ids = input_ids[: self.block_size] else: - input_ids = input_ids + [self.pad_token_id] * (self.block_size - len(input_ids)) - + input_ids = input_ids + [self.pad_token_id] * ( + self.block_size - len(input_ids) + ) + input_ids = torch.tensor(input_ids, dtype=torch.long) labels = input_ids.clone() - + return {"input_ids": input_ids, "labels": labels} @@ -84,9 +86,15 @@ class TextDatasetWithSpecialTokens(TextDataset): """ Расширенная версия TextDataset с поддержкой специальных токенов. """ - - def __init__(self, texts: List[str], tokenizer: Any, block_size: int = 128, - add_bos: bool = False, add_eos: bool = False): + + def __init__( + self, + texts: List[str], + tokenizer: Any, + block_size: int = 128, + add_bos: bool = False, + add_eos: bool = False, + ): """ Args: texts: Список текстов @@ -104,33 +112,38 @@ class TextDatasetWithSpecialTokens(TextDataset): for text in texts: # Кодируем с специальными токенами input_ids = tokenizer.encode( - text, - add_special_tokens=True, - add_bos_token=add_bos, - add_eos_token=eos + text, add_special_tokens=True, add_bos_token=add_bos, add_eos_token=eos ) - + # Учитываем специальные токены при обрезке/дополнении effective_block_size = block_size if add_bos: effective_block_size -= 1 if add_eos: effective_block_size -= 1 - + if len(input_ids) > effective_block_size: input_ids = input_ids[:effective_block_size] - + # Добавляем специальные токены если нужно - if add_bos and hasattr(tokenizer, 'bos_token_id') and tokenizer.bos_token_id is not None: + if ( + add_bos + and hasattr(tokenizer, "bos_token_id") + and tokenizer.bos_token_id is not None + ): input_ids = [tokenizer.bos_token_id] + input_ids - if add_eos and hasattr(tokenizer, 'eos_token_id') and tokenizer.eos_token_id is not None: + if ( + add_eos + and hasattr(tokenizer, "eos_token_id") + and tokenizer.eos_token_id is not None + ): input_ids = input_ids + [tokenizer.eos_token_id] - + # Дополняем до полной длины - pad_token_id = getattr(tokenizer, 'pad_token_id', 0) + pad_token_id = getattr(tokenizer, "pad_token_id", 0) if len(input_ids) < block_size: input_ids = input_ids + [pad_token_id] * (block_size - len(input_ids)) - + self.examples.append(input_ids) def __len__(self): diff --git a/llm/src/llm/training/optimizer.py b/llm/src/llm/training/optimizer.py index b58b05a..0ee6359 100644 --- a/llm/src/llm/training/optimizer.py +++ b/llm/src/llm/training/optimizer.py @@ -1,5 +1,6 @@ import torch.optim as optim + def get_optimizer(model, lr=3e-4, weight_decay=0.01, optimizer_type="adamw"): """ Возвращает оптимизатор для обучения модели. diff --git a/llm/src/llm/training/scheduler.py b/llm/src/llm/training/scheduler.py index c0fb2be..9105a12 100644 --- a/llm/src/llm/training/scheduler.py +++ b/llm/src/llm/training/scheduler.py @@ -1,5 +1,6 @@ from torch.optim.lr_scheduler import LambdaLR + def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps): """ Линейный планировщик обучения с warmup. @@ -8,6 +9,10 @@ def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_st def lr_lambda(current_step): if current_step < num_warmup_steps: return float(current_step) / float(max(1, num_warmup_steps)) - return max(0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps))) + return max( + 0.0, + float(num_training_steps - current_step) + / float(max(1, num_training_steps - num_warmup_steps)), + ) return LambdaLR(optimizer, lr_lambda) diff --git a/llm/src/llm/training/trainer.py b/llm/src/llm/training/trainer.py index 966cdc0..5ed16e9 100644 --- a/llm/src/llm/training/trainer.py +++ b/llm/src/llm/training/trainer.py @@ -5,15 +5,29 @@ from tqdm import tqdm from llm.training.optimizer import get_optimizer from llm.training.scheduler import get_linear_schedule_with_warmup + class Trainer: """ Универсальный класс обучения LLM (GPT, LLaMA, Mistral и т.д.) """ - def __init__(self, model, train_dataset, val_dataset=None, lr=3e-4, batch_size=8, num_epochs=3, warmup_steps=100): + def __init__( + self, + model, + train_dataset, + val_dataset=None, + lr=3e-4, + batch_size=8, + num_epochs=3, + warmup_steps=100, + ): self.model = model - self.train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) - self.val_loader = DataLoader(val_dataset, batch_size=batch_size) if val_dataset else None + self.train_loader = DataLoader( + train_dataset, batch_size=batch_size, shuffle=True + ) + self.val_loader = ( + DataLoader(val_dataset, batch_size=batch_size) if val_dataset else None + ) self.optimizer = get_optimizer(model, lr=lr) self.scheduler = None self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") @@ -29,24 +43,28 @@ class Trainer: # Сдвигаем логиты и метки для языкового моделирования shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() - + # Вычисляем cross-entropy loss loss = F.cross_entropy( shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), - ignore_index=-100 # Игнорируем padding tokens + ignore_index=-100, # Игнорируем padding tokens ) return loss def train(self): total_steps = len(self.train_loader) * self.num_epochs - self.scheduler = get_linear_schedule_with_warmup(self.optimizer, self.warmup_steps, total_steps) + self.scheduler = get_linear_schedule_with_warmup( + self.optimizer, self.warmup_steps, total_steps + ) for epoch in range(self.num_epochs): self.model.train() total_loss = 0 - progress_bar = tqdm(self.train_loader, desc=f"Epoch {epoch+1}/{self.num_epochs}") + progress_bar = tqdm( + self.train_loader, desc=f"Epoch {epoch+1}/{self.num_epochs}" + ) for batch in progress_bar: self.optimizer.zero_grad() @@ -59,7 +77,7 @@ class Trainer: logits = outputs[0] else: logits = outputs - + # Trainer вычисляет loss loss = self.compute_lm_loss(logits, labels) loss.backward() @@ -85,7 +103,7 @@ class Trainer: for batch in self.val_loader: input_ids = batch["input_ids"].to(self.device) labels = batch["labels"].to(self.device) - + outputs = self.model(input_ids) if isinstance(outputs, tuple): logits = outputs[0] diff --git a/llm/tests/conftest.py b/llm/tests/conftest.py index b3bfbac..321dff4 100644 --- a/llm/tests/conftest.py +++ b/llm/tests/conftest.py @@ -58,7 +58,7 @@ def gpt_config(vocab_size, embed_dim, num_heads, num_layers): "num_heads": num_heads, "num_layers": num_layers, "max_position_embeddings": 1024, - "dropout": 0.1 + "dropout": 0.1, } @@ -68,12 +68,14 @@ def random_inputs(batch_size, seq_len, vocab_size): input_ids = torch.randint(0, vocab_size, (batch_size, seq_len)) return input_ids + @pytest.fixture def random_float_inputs(batch_size, seq_len, embed_dim): """Generate random floating point input tensors for testing feed forward.""" inputs = torch.randn(batch_size, seq_len, embed_dim) return inputs + @pytest.fixture def random_embeddings(batch_size, seq_len, embed_dim): """Generate random embedding tensors for testing attention modules.""" diff --git a/llm/tests/core/test_decoder.py b/llm/tests/core/test_decoder.py index a710632..8eae46f 100644 --- a/llm/tests/core/test_decoder.py +++ b/llm/tests/core/test_decoder.py @@ -9,180 +9,233 @@ from llm.core.decoder import Decoder class TestDecoder: """Test cases for Decoder.""" - + def test_initialization(self, embed_dim, num_heads): """Test that Decoder can be initialized.""" head_size = embed_dim // num_heads max_seq_len = 1024 - decoder = Decoder(num_heads=num_heads, emb_size=embed_dim, head_size=head_size, max_seq_len=max_seq_len) + decoder = Decoder( + num_heads=num_heads, + emb_size=embed_dim, + head_size=head_size, + max_seq_len=max_seq_len, + ) assert decoder is not None - + # Check internal components - assert hasattr(decoder, '_heads') - assert hasattr(decoder, '_ff') - assert hasattr(decoder, '_norm1') - assert hasattr(decoder, '_norm2') - + assert hasattr(decoder, "_heads") + assert hasattr(decoder, "_ff") + assert hasattr(decoder, "_norm1") + assert hasattr(decoder, "_norm2") + def test_forward_pass(self, embed_dim, num_heads, random_embeddings): """Test forward pass of Decoder.""" head_size = embed_dim // num_heads max_seq_len = 1024 - decoder = Decoder(num_heads=num_heads, emb_size=embed_dim, head_size=head_size, max_seq_len=max_seq_len) - + decoder = Decoder( + num_heads=num_heads, + emb_size=embed_dim, + head_size=head_size, + max_seq_len=max_seq_len, + ) + # Forward pass output = decoder(random_embeddings) - + # Check output shape assert output.shape == random_embeddings.shape assert isinstance(output, torch.Tensor) - + def test_forward_with_causal_mask(self, embed_dim, num_heads, random_embeddings): """Test forward pass with causal mask.""" head_size = embed_dim // num_heads max_seq_len = 1024 - decoder = Decoder(num_heads=num_heads, emb_size=embed_dim, head_size=head_size, max_seq_len=max_seq_len) - + decoder = Decoder( + num_heads=num_heads, + emb_size=embed_dim, + head_size=head_size, + max_seq_len=max_seq_len, + ) + batch_size, seq_len = random_embeddings.shape[:2] # Create causal mask mask = torch.tril(torch.ones(seq_len, seq_len)) - + # Forward pass with causal mask output = decoder(random_embeddings, mask=mask) - + # Check output shape assert output.shape == random_embeddings.shape - + def test_residual_connections(self, embed_dim, num_heads, random_embeddings): """Test that residual connections are properly applied.""" head_size = embed_dim // num_heads max_seq_len = 1024 - decoder = Decoder(num_heads=num_heads, emb_size=embed_dim, head_size=head_size, max_seq_len=max_seq_len) - + decoder = Decoder( + num_heads=num_heads, + emb_size=embed_dim, + head_size=head_size, + max_seq_len=max_seq_len, + ) + output = decoder(random_embeddings) - + # With residual connections and layer norm, the output shouldn't be # too different from input (in terms of scale/distribution) input_norm = random_embeddings.norm(dim=-1).mean() output_norm = output.norm(dim=-1).mean() - + # Norms should be of similar magnitude (not exact due to transformations) assert 0.1 < (output_norm / input_norm) < 10.0 - + def test_layer_norm(self, embed_dim, num_heads, random_embeddings): """Test that layer normalization is applied.""" head_size = embed_dim // num_heads max_seq_len = 1024 - decoder = Decoder(num_heads=num_heads, emb_size=embed_dim, head_size=head_size, max_seq_len=max_seq_len) - + decoder = Decoder( + num_heads=num_heads, + emb_size=embed_dim, + head_size=head_size, + max_seq_len=max_seq_len, + ) + output = decoder(random_embeddings) - + # Check that output has reasonable statistics (due to layer norm) # Mean should be close to 0, std close to 1 for each sequence position output_mean = output.mean(dim=-1) output_std = output.std(dim=-1) - + # These are approximate checks since the data goes through multiple transformations assert torch.allclose(output_mean, torch.zeros_like(output_mean), atol=1.0) assert torch.allclose(output_std, torch.ones_like(output_std), atol=2.0) - + def test_gradient_flow(self, embed_dim, num_heads, random_embeddings): """Test that gradients flow through Decoder.""" head_size = embed_dim // num_heads max_seq_len = 1024 - decoder = Decoder(num_heads=num_heads, emb_size=embed_dim, head_size=head_size, max_seq_len=max_seq_len) - + decoder = Decoder( + num_heads=num_heads, + emb_size=embed_dim, + head_size=head_size, + max_seq_len=max_seq_len, + ) + # Forward pass output = decoder(random_embeddings) - + # Create a dummy loss and backward pass loss = output.sum() loss.backward() - + # Check that gradients are computed for learnable parameters # in attention and feed forward components assert decoder._heads._layer.weight.grad is not None assert decoder._ff._layer1.weight.grad is not None assert decoder._norm1.weight.grad is not None assert decoder._norm2.weight.grad is not None - + def test_device_consistency(self, embed_dim, num_heads, random_embeddings, device): """Test that Decoder works on correct device.""" head_size = embed_dim // num_heads max_seq_len = 1024 - decoder = Decoder(num_heads=num_heads, emb_size=embed_dim, head_size=head_size, max_seq_len=max_seq_len).to(device) + decoder = Decoder( + num_heads=num_heads, + emb_size=embed_dim, + head_size=head_size, + max_seq_len=max_seq_len, + ).to(device) inputs = random_embeddings.to(device) - + # Forward pass output = decoder(inputs) - + # Check device consistency assert output.device == device assert decoder._heads._layer.weight.device == device - + def test_different_configurations(self): """Test Decoder with different configurations.""" test_cases = [ - (64, 2), # embed_dim=64, num_heads=2 + (64, 2), # embed_dim=64, num_heads=2 (128, 4), # embed_dim=128, num_heads=4 (256, 8), # embed_dim=256, num_heads=8 ] - + for embed_dim, num_heads in test_cases: head_size = embed_dim // num_heads max_seq_len = 1024 - decoder = Decoder(num_heads=num_heads, emb_size=embed_dim, head_size=head_size, max_seq_len=max_seq_len) + decoder = Decoder( + num_heads=num_heads, + emb_size=embed_dim, + head_size=head_size, + max_seq_len=max_seq_len, + ) batch_size, seq_len = 2, 16 inputs = torch.randn(batch_size, seq_len, embed_dim) - + output = decoder(inputs) - + assert output.shape == inputs.shape - + @pytest.mark.parametrize("batch_size,seq_len", [(1, 8), (2, 16), (4, 32)]) def test_different_input_shapes(self, embed_dim, num_heads, batch_size, seq_len): """Test Decoder with different input shapes.""" head_size = embed_dim // num_heads max_seq_len = 1024 - decoder = Decoder(num_heads=num_heads, emb_size=embed_dim, head_size=head_size, max_seq_len=max_seq_len) - + decoder = Decoder( + num_heads=num_heads, + emb_size=embed_dim, + head_size=head_size, + max_seq_len=max_seq_len, + ) + inputs = torch.randn(batch_size, seq_len, embed_dim) output = decoder(inputs) - + assert output.shape == (batch_size, seq_len, embed_dim) - + def test_training_vs_evaluation(self, embed_dim, num_heads, random_embeddings): """Test that Decoder behaves differently in train vs eval mode.""" head_size = embed_dim // num_heads max_seq_len = 1024 - decoder = Decoder(num_heads=num_heads, emb_size=embed_dim, head_size=head_size, max_seq_len=max_seq_len, dropout=0.5) - + decoder = Decoder( + num_heads=num_heads, + emb_size=embed_dim, + head_size=head_size, + max_seq_len=max_seq_len, + dropout=0.5, + ) + # Training mode decoder.train() output_train = decoder(random_embeddings) - + # Evaluation mode decoder.eval() output_eval = decoder(random_embeddings) - + # Outputs should be different due to dropout assert not torch.allclose(output_train, output_eval) - + def test_parameter_initialization(self, embed_dim, num_heads): """Test that parameters are properly initialized.""" head_size = embed_dim // num_heads max_seq_len = 1024 - decoder = Decoder(num_heads=num_heads, emb_size=embed_dim, head_size=head_size, max_seq_len=max_seq_len) - + decoder = Decoder( + num_heads=num_heads, + emb_size=embed_dim, + head_size=head_size, + max_seq_len=max_seq_len, + ) + # Check that various components have non-zero parameters assert not torch.allclose( - decoder._heads._layer.weight, - torch.zeros_like(decoder._heads._layer.weight) + decoder._heads._layer.weight, torch.zeros_like(decoder._heads._layer.weight) ) assert not torch.allclose( - decoder._ff._layer1.weight, - torch.zeros_like(decoder._ff._layer1.weight) + decoder._ff._layer1.weight, torch.zeros_like(decoder._ff._layer1.weight) ) assert not torch.allclose( - decoder._norm1.weight, - torch.zeros_like(decoder._norm1.weight) + decoder._norm1.weight, torch.zeros_like(decoder._norm1.weight) ) diff --git a/llm/tests/core/test_feed_forward.py b/llm/tests/core/test_feed_forward.py index c746195..6416da5 100644 --- a/llm/tests/core/test_feed_forward.py +++ b/llm/tests/core/test_feed_forward.py @@ -10,168 +10,178 @@ from llm.core.feed_forward import FeedForward class TestFeedForward: """Test cases for FeedForward.""" - + def test_initialization(self, embed_dim): """Test that FeedForward can be initialized.""" ff = FeedForward(embed_dim) assert ff is not None - + # Check internal layers - assert hasattr(ff, '_layer1') - assert hasattr(ff, '_layer2') - assert hasattr(ff, '_activation') - assert hasattr(ff, '_dropout') - + assert hasattr(ff, "_layer1") + assert hasattr(ff, "_layer2") + assert hasattr(ff, "_activation") + assert hasattr(ff, "_dropout") + # Check layer dimensions expected_hidden_dim = embed_dim * 4 # Default expansion factor assert ff._layer1.weight.shape == (expected_hidden_dim, embed_dim) assert ff._layer2.weight.shape == (embed_dim, expected_hidden_dim) - + def test_forward_pass(self, embed_dim, random_float_inputs): """Test forward pass of FeedForward.""" ff = FeedForward(embed_dim) - + # Forward pass output = ff(random_float_inputs) - + # Check output shape assert output.shape == random_float_inputs.shape assert isinstance(output, torch.Tensor) - + def test_custom_hidden_dim(self, embed_dim): """Test FeedForward with custom hidden dimension.""" # FeedForward doesn't support custom hidden_dim in current implementation # This test is not applicable ff = FeedForward(embed_dim) - + # Check layer dimensions (fixed 4x expansion) expected_hidden_dim = embed_dim * 4 assert ff._layer1.weight.shape == (expected_hidden_dim, embed_dim) assert ff._layer2.weight.shape == (embed_dim, expected_hidden_dim) - + def test_dropout(self, embed_dim, random_float_inputs): """Test that dropout is applied during training.""" ff = FeedForward(embed_dim, dropout=0.5) ff.train() # Set to training mode - + output = ff(random_float_inputs) - + # In training mode with dropout, some values should be zeroed # This is probabilistic, so we can't assert exact zeros, # but we can check the structure is preserved assert output.shape == random_float_inputs.shape - + def test_no_dropout_in_eval(self, embed_dim, random_float_inputs): """Test that dropout is not applied during evaluation.""" ff = FeedForward(embed_dim, dropout=0.5) ff.eval() # Set to evaluation mode - + # Run forward pass multiple times - outputs should be identical output1 = ff(random_float_inputs) output2 = ff(random_float_inputs) - + assert torch.allclose(output1, output2) - + def test_activation_function(self, embed_dim, random_float_inputs): """Test that activation function is applied.""" ff = FeedForward(embed_dim) - + # Manually compute expected output without dropout for deterministic comparison hidden = ff._layer1(random_float_inputs) activated = ff._activation(hidden) expected_output = ff._layer2(activated) - + # Compare with forward pass in eval mode (no dropout) ff.eval() actual_output = ff(random_float_inputs) - + assert torch.allclose(actual_output, expected_output, rtol=1e-4) - + def test_gradient_flow(self, embed_dim, random_float_inputs): """Test that gradients flow through FeedForward.""" ff = FeedForward(embed_dim) - + # Forward pass output = ff(random_float_inputs) - + # Create a dummy loss and backward pass loss = output.sum() loss.backward() - + # Check that gradients are computed for learnable parameters assert ff._layer1.weight.grad is not None assert ff._layer2.weight.grad is not None - assert not torch.allclose(ff._layer1.weight.grad, - torch.zeros_like(ff._layer1.weight.grad)) - assert not torch.allclose(ff._layer2.weight.grad, - torch.zeros_like(ff._layer2.weight.grad)) - + assert not torch.allclose( + ff._layer1.weight.grad, torch.zeros_like(ff._layer1.weight.grad) + ) + assert not torch.allclose( + ff._layer2.weight.grad, torch.zeros_like(ff._layer2.weight.grad) + ) + def test_device_consistency(self, embed_dim, random_float_inputs, device): """Test that FeedForward works on correct device.""" ff = FeedForward(embed_dim).to(device) inputs = random_float_inputs.to(device) - + # Forward pass output = ff(inputs) - + # Check device consistency assert output.device == device assert ff._layer1.weight.device == device assert ff._layer2.weight.device == device - + def test_different_embed_dims(self): """Test FeedForward with different embedding dimensions.""" test_cases = [64, 128, 256, 512] - + for embed_dim in test_cases: ff = FeedForward(embed_dim) batch_size, seq_len = 2, 16 inputs = torch.randn(batch_size, seq_len, embed_dim) - + output = ff(inputs) - + assert output.shape == inputs.shape - + @pytest.mark.parametrize("batch_size,seq_len", [(1, 8), (2, 16), (4, 32)]) def test_different_input_shapes(self, embed_dim, batch_size, seq_len): """Test FeedForward with different input shapes.""" ff = FeedForward(embed_dim) - + inputs = torch.randn(batch_size, seq_len, embed_dim) output = ff(inputs) - + assert output.shape == (batch_size, seq_len, embed_dim) - + def test_non_linearity(self, embed_dim, random_float_inputs): """Test that FeedForward introduces non-linearity.""" ff = FeedForward(embed_dim) - + # Create a simple linear transformation for comparison linear_layer = nn.Linear(embed_dim, embed_dim) - + # Copy weights to make comparison fair with torch.no_grad(): linear_layer.weight.copy_(ff._layer2.weight @ ff._layer1.weight) if linear_layer.bias is not None: linear_layer.bias.zero_() - + linear_output = linear_layer(random_float_inputs) ff_output = ff(random_float_inputs) - + # FeedForward output should be different from pure linear transformation # due to activation function assert not torch.allclose(ff_output, linear_output, rtol=1e-4) - + def test_parameter_initialization(self, embed_dim): """Test that parameters are properly initialized.""" ff = FeedForward(embed_dim) - + # Check that weights are not all zeros - assert not torch.allclose(ff._layer1.weight, torch.zeros_like(ff._layer1.weight)) - assert not torch.allclose(ff._layer2.weight, torch.zeros_like(ff._layer2.weight)) - + assert not torch.allclose( + ff._layer1.weight, torch.zeros_like(ff._layer1.weight) + ) + assert not torch.allclose( + ff._layer2.weight, torch.zeros_like(ff._layer2.weight) + ) + # Check that biases are not all zeros (they should be initialized with some values) if ff._layer1.bias is not None: - assert not torch.allclose(ff._layer1.bias, torch.zeros_like(ff._layer1.bias)) + assert not torch.allclose( + ff._layer1.bias, torch.zeros_like(ff._layer1.bias) + ) if ff._layer2.bias is not None: - assert not torch.allclose(ff._layer2.bias, torch.zeros_like(ff._layer2.bias)) + assert not torch.allclose( + ff._layer2.bias, torch.zeros_like(ff._layer2.bias) + ) diff --git a/llm/tests/core/test_multi_head_attention.py b/llm/tests/core/test_multi_head_attention.py index 892ef68..cd4246d 100644 --- a/llm/tests/core/test_multi_head_attention.py +++ b/llm/tests/core/test_multi_head_attention.py @@ -9,157 +9,181 @@ from llm.core.multi_head_attention import MultiHeadAttention class TestMultiHeadAttention: """Test cases for MultiHeadAttention.""" - + def test_initialization(self, embed_dim, num_heads): """Test that MultiHeadAttention can be initialized.""" head_size = embed_dim // num_heads - attention = MultiHeadAttention(num_heads, embed_dim, head_size, max_seq_len=1024) + attention = MultiHeadAttention( + num_heads, embed_dim, head_size, max_seq_len=1024 + ) assert attention is not None - + # Check internal attributes assert len(attention._heads) == num_heads assert attention._layer.in_features == embed_dim assert attention._layer.out_features == embed_dim - + def test_forward_pass(self, embed_dim, num_heads, random_embeddings): """Test forward pass of MultiHeadAttention.""" head_size = embed_dim // num_heads - attention = MultiHeadAttention(num_heads, embed_dim, head_size, max_seq_len=1024) - + attention = MultiHeadAttention( + num_heads, embed_dim, head_size, max_seq_len=1024 + ) + # Forward pass output, _ = attention(random_embeddings) - + # Check output shape assert output.shape == random_embeddings.shape assert isinstance(output, torch.Tensor) - + def test_forward_with_mask(self, embed_dim, num_heads, random_embeddings): """Test forward pass with attention mask.""" head_size = embed_dim // num_heads - attention = MultiHeadAttention(num_heads, embed_dim, head_size, max_seq_len=1024) - + attention = MultiHeadAttention( + num_heads, embed_dim, head_size, max_seq_len=1024 + ) + # Create a simple mask seq_len = random_embeddings.shape[1] mask = torch.tril(torch.ones(seq_len, seq_len)) # Causal mask - + # Forward pass with mask output, _ = attention(random_embeddings, mask=mask) - + # Check output shape assert output.shape == random_embeddings.shape - + def test_causal_mask(self, embed_dim, num_heads, random_embeddings): """Test that causal mask prevents attending to future positions.""" head_size = embed_dim // num_heads - attention = MultiHeadAttention(num_heads, embed_dim, head_size, max_seq_len=1024) - + attention = MultiHeadAttention( + num_heads, embed_dim, head_size, max_seq_len=1024 + ) + # Create causal mask seq_len = random_embeddings.shape[1] causal_mask = torch.tril(torch.ones(seq_len, seq_len)) - + # Forward pass with causal mask output, _ = attention(random_embeddings, mask=causal_mask) - + # Check output shape assert output.shape == random_embeddings.shape - - def test_attention_weights_normalization(self, embed_dim, num_heads, random_embeddings): + + def test_attention_weights_normalization( + self, embed_dim, num_heads, random_embeddings + ): """Test that attention weights are properly normalized.""" head_size = embed_dim // num_heads - attention = MultiHeadAttention(num_heads, embed_dim, head_size, max_seq_len=1024) - + attention = MultiHeadAttention( + num_heads, embed_dim, head_size, max_seq_len=1024 + ) + # Forward pass output, _ = attention(random_embeddings) - + # Check output shape assert output.shape == random_embeddings.shape - + def test_gradient_flow(self, embed_dim, num_heads, random_embeddings): """Test that gradients flow through MultiHeadAttention.""" head_size = embed_dim // num_heads - attention = MultiHeadAttention(num_heads, embed_dim, head_size, max_seq_len=1024) - + attention = MultiHeadAttention( + num_heads, embed_dim, head_size, max_seq_len=1024 + ) + # Forward pass output, _ = attention(random_embeddings) - + # Create a dummy loss and backward pass loss = output.sum() loss.backward() - + # Check that gradients are computed for learnable parameters assert attention._layer.weight.grad is not None if len(attention._heads) > 0: assert attention._heads[0]._q.weight.grad is not None - + def test_device_consistency(self, embed_dim, num_heads, random_embeddings, device): """Test that MultiHeadAttention works on correct device.""" head_size = embed_dim // num_heads - attention = MultiHeadAttention(num_heads, embed_dim, head_size, max_seq_len=1024).to(device) + attention = MultiHeadAttention( + num_heads, embed_dim, head_size, max_seq_len=1024 + ).to(device) inputs = random_embeddings.to(device) - + # Forward pass output, _ = attention(inputs) - + # Check device consistency assert output.device == device assert attention._layer.weight.device == device - + def test_different_embed_dim_and_heads(self): """Test MultiHeadAttention with different embed_dim and num_heads combinations.""" test_cases = [ - (64, 2), # embed_dim=64, num_heads=2 + (64, 2), # embed_dim=64, num_heads=2 (128, 4), # embed_dim=128, num_heads=4 (256, 8), # embed_dim=256, num_heads=8 - (512, 16), # embed_dim=512, num_heads=16 + (512, 16), # embed_dim=512, num_heads=16 ] - + for embed_dim, num_heads in test_cases: head_size = embed_dim // num_heads - attention = MultiHeadAttention(num_heads, embed_dim, head_size, max_seq_len=1024) + attention = MultiHeadAttention( + num_heads, embed_dim, head_size, max_seq_len=1024 + ) batch_size, seq_len = 2, 16 inputs = torch.randn(batch_size, seq_len, embed_dim) - + output, _ = attention(inputs) - + assert output.shape == inputs.shape - + def test_attention_output_range(self, embed_dim, num_heads, random_embeddings): """Test that attention output is in reasonable range.""" head_size = embed_dim // num_heads - attention = MultiHeadAttention(num_heads, embed_dim, head_size, max_seq_len=1024) - + attention = MultiHeadAttention( + num_heads, embed_dim, head_size, max_seq_len=1024 + ) + output, _ = attention(random_embeddings) - + # Output shouldn't have extreme values assert output.abs().max() < 100 # Reasonable upper bound - + @pytest.mark.parametrize("batch_size,seq_len", [(1, 8), (2, 16), (4, 32)]) def test_different_input_shapes(self, embed_dim, num_heads, batch_size, seq_len): """Test MultiHeadAttention with different input shapes.""" head_size = embed_dim // num_heads - attention = MultiHeadAttention(num_heads, embed_dim, head_size, max_seq_len=1024) - + attention = MultiHeadAttention( + num_heads, embed_dim, head_size, max_seq_len=1024 + ) + inputs = torch.randn(batch_size, seq_len, embed_dim) output, _ = attention(inputs) - + assert output.shape == (batch_size, seq_len, embed_dim) - + def test_parameter_sharing(self, embed_dim, num_heads): """Test that parameters are properly shared across the sequence.""" head_size = embed_dim // num_heads - attention = MultiHeadAttention(num_heads, embed_dim, head_size, max_seq_len=1024, dropout=0.0) # No dropout for deterministic test - + attention = MultiHeadAttention( + num_heads, embed_dim, head_size, max_seq_len=1024, dropout=0.0 + ) # No dropout for deterministic test + # Create two identical sequences seq_len = 10 base_sequence = torch.randn(1, seq_len, embed_dim) identical_sequence = base_sequence.clone() - + # Set to eval mode to disable dropout attention.eval() - + with torch.no_grad(): output1, _ = attention(base_sequence) output2, _ = attention(identical_sequence) - + # With identical inputs and same parameters, outputs should be identical assert torch.allclose(output1, output2, rtol=1e-5) diff --git a/llm/tests/core/test_positional_embeddings.py b/llm/tests/core/test_positional_embeddings.py index a31df24..d9f3f59 100644 --- a/llm/tests/core/test_positional_embeddings.py +++ b/llm/tests/core/test_positional_embeddings.py @@ -10,127 +10,134 @@ from llm.core.positional_embeddings import PositionalEmbeddings class TestPositionalEmbeddings: """Test cases for PositionalEmbeddings.""" - + def test_initialization(self, embed_dim): """Test that PositionalEmbeddings can be initialized.""" max_seq_len = 1024 embeddings = PositionalEmbeddings(max_seq_len, embed_dim) assert embeddings is not None - + # Check that positional embeddings are created - assert hasattr(embeddings, 'embedding') + assert hasattr(embeddings, "embedding") assert embeddings.embedding.weight.shape == (max_seq_len, embed_dim) - + def test_forward_pass(self, embed_dim): """Test forward pass of PositionalEmbeddings.""" max_seq_len = 1024 seq_len = 64 embeddings = PositionalEmbeddings(max_seq_len, embed_dim) - + # Forward pass - takes sequence length, not input tensor output = embeddings(seq_len) - + # Check output shape expected_shape = (seq_len, embed_dim) assert output.shape == expected_shape assert isinstance(output, torch.Tensor) - + def test_positional_encoding_values(self, embed_dim): """Test that positional encoding values are computed correctly.""" max_seq_len = 10 embeddings = PositionalEmbeddings(max_seq_len, embed_dim) - + # Get embeddings for all positions pe = embeddings(max_seq_len) # Shape: [max_seq_len, embed_dim] - + # Check that different positions have different embeddings # (since these are learnable embeddings, not fixed sine/cosine) for pos in range(max_seq_len): for i in range(pos + 1, max_seq_len): assert not torch.allclose(pe[pos], pe[i], rtol=1e-4) - + def test_different_sequence_lengths(self, embed_dim): """Test PositionalEmbeddings with different sequence lengths.""" test_cases = [ - (10, 5), # seq_len < max_seq_len + (10, 5), # seq_len < max_seq_len (10, 10), # seq_len == max_seq_len ] - + for max_seq_len, seq_len in test_cases: embeddings = PositionalEmbeddings(max_seq_len, embed_dim) - + # Get embeddings for specific sequence length output = embeddings(seq_len) - + # Output should have shape [seq_len, embed_dim] assert output.shape == (seq_len, embed_dim) - + def test_gradient_flow(self, embed_dim): """Test that gradients flow through PositionalEmbeddings.""" max_seq_len = 64 seq_len = 32 embeddings = PositionalEmbeddings(max_seq_len, embed_dim) - + # Forward pass output = embeddings(seq_len) - + # Create a dummy loss and backward pass loss = output.sum() loss.backward() - + # Positional embeddings should have gradients (they're learnable) assert embeddings.embedding.weight.grad is not None - assert not torch.allclose(embeddings.embedding.weight.grad, - torch.zeros_like(embeddings.embedding.weight.grad)) - + assert not torch.allclose( + embeddings.embedding.weight.grad, + torch.zeros_like(embeddings.embedding.weight.grad), + ) + def test_device_consistency(self, embed_dim, device): """Test that PositionalEmbeddings works on correct device.""" max_seq_len = 64 seq_len = 32 embeddings = PositionalEmbeddings(max_seq_len, embed_dim).to(device) - + # Forward pass output = embeddings(seq_len) - + # Check device consistency assert output.device == device assert embeddings.embedding.weight.device == device - + def test_reproducibility(self, embed_dim): """Test that positional embeddings are reproducible.""" max_seq_len = 100 embeddings1 = PositionalEmbeddings(max_seq_len, embed_dim) embeddings2 = PositionalEmbeddings(max_seq_len, embed_dim) - + # Different instances should have different embeddings (random initialization) - assert not torch.allclose(embeddings1.embedding.weight, embeddings2.embedding.weight) - + assert not torch.allclose( + embeddings1.embedding.weight, embeddings2.embedding.weight + ) + # But same instance should produce same output for same input seq_len = 50 output1 = embeddings1(seq_len) output2 = embeddings1(seq_len) # Same instance, same input assert torch.allclose(output1, output2) - + def test_positional_pattern(self, embed_dim): """Test that positional embeddings create a meaningful pattern.""" max_seq_len = 50 embeddings = PositionalEmbeddings(max_seq_len, embed_dim) pe = embeddings(max_seq_len) # Get all positional embeddings - + # Check that different positions have different embeddings # (with high probability due to random initialization) assert not torch.allclose(pe[0], pe[1], rtol=1e-4) assert not torch.allclose(pe[10], pe[20], rtol=1e-4) - - @pytest.mark.parametrize("max_seq_len,seq_len,embed_dim", [ - (64, 10, 64), - (128, 50, 128), - (256, 100, 256), - ]) + + @pytest.mark.parametrize( + "max_seq_len,seq_len,embed_dim", + [ + (64, 10, 64), + (128, 50, 128), + (256, 100, 256), + ], + ) def test_different_configurations(self, max_seq_len, seq_len, embed_dim): """Test PositionalEmbeddings with different configurations.""" embeddings = PositionalEmbeddings(max_seq_len, embed_dim) - + output = embeddings(seq_len) - + assert output.shape == (seq_len, embed_dim) diff --git a/llm/tests/core/test_token_embeddings.py b/llm/tests/core/test_token_embeddings.py index c613aae..82bdc06 100644 --- a/llm/tests/core/test_token_embeddings.py +++ b/llm/tests/core/test_token_embeddings.py @@ -9,99 +9,103 @@ from llm.core.token_embeddings import TokenEmbeddings class TestTokenEmbeddings: """Test cases for TokenEmbeddings.""" - + def test_initialization(self, vocab_size, embed_dim): """Test that TokenEmbeddings can be initialized.""" embeddings = TokenEmbeddings(vocab_size, embed_dim) assert embeddings is not None - + # Check embedding layer - assert hasattr(embeddings, '_embedding') + assert hasattr(embeddings, "_embedding") assert embeddings._embedding.weight.shape == (vocab_size, embed_dim) - + def test_forward_pass(self, vocab_size, embed_dim, random_inputs): """Test forward pass of TokenEmbeddings.""" embeddings = TokenEmbeddings(vocab_size, embed_dim) - + # Forward pass output = embeddings(random_inputs) - + # Check output shape - assert output.shape == (random_inputs.shape[0], random_inputs.shape[1], embed_dim) + assert output.shape == ( + random_inputs.shape[0], + random_inputs.shape[1], + embed_dim, + ) assert isinstance(output, torch.Tensor) - + def test_embedding_weights(self, vocab_size, embed_dim): """Test that embedding weights are properly initialized.""" embeddings = TokenEmbeddings(vocab_size, embed_dim) - + weights = embeddings._embedding.weight assert weights.requires_grad is True - + # Check that weights are not all zeros assert not torch.allclose(weights, torch.zeros_like(weights)) - + def test_different_vocab_sizes(self): """Test TokenEmbeddings with different vocabulary sizes.""" - test_cases = [ - (100, 128), - (1000, 256), - (50000, 512) - ] - + test_cases = [(100, 128), (1000, 256), (50000, 512)] + for vocab_size, embed_dim in test_cases: embeddings = TokenEmbeddings(vocab_size, embed_dim) assert embeddings._embedding.weight.shape == (vocab_size, embed_dim) - + def test_gradient_flow(self, vocab_size, embed_dim, random_inputs): """Test that gradients flow through TokenEmbeddings.""" embeddings = TokenEmbeddings(vocab_size, embed_dim) - + # Forward pass output = embeddings(random_inputs) - + # Create a dummy loss and backward pass loss = output.sum() loss.backward() - + # Check that gradients are computed assert embeddings._embedding.weight.grad is not None - assert not torch.allclose(embeddings._embedding.weight.grad, - torch.zeros_like(embeddings._embedding.weight.grad)) - + assert not torch.allclose( + embeddings._embedding.weight.grad, + torch.zeros_like(embeddings._embedding.weight.grad), + ) + def test_device_consistency(self, vocab_size, embed_dim, random_inputs, device): """Test that TokenEmbeddings works on correct device.""" embeddings = TokenEmbeddings(vocab_size, embed_dim).to(device) inputs = random_inputs.to(device) - + # Forward pass output = embeddings(inputs) - + # Check device consistency assert output.device == device assert embeddings._embedding.weight.device == device - + def test_embedding_lookup(self, vocab_size, embed_dim): """Test specific embedding lookups.""" embeddings = TokenEmbeddings(vocab_size, embed_dim) - + # Test lookup for specific tokens - test_tokens = torch.tensor([[0, 1, 2], [vocab_size - 1, vocab_size - 2, vocab_size - 3]]) - + test_tokens = torch.tensor( + [[0, 1, 2], [vocab_size - 1, vocab_size - 2, vocab_size - 3]] + ) + output = embeddings(test_tokens) - + # Check shape assert output.shape == (2, 3, embed_dim) - + # Check that different tokens have different embeddings # (with high probability due to random initialization) assert not torch.allclose(output[0, 0], output[0, 1], rtol=1e-4) - + @pytest.mark.parametrize("batch_size,seq_len", [(1, 1), (2, 10), (8, 64)]) def test_different_input_shapes(self, vocab_size, embed_dim, batch_size, seq_len): """Test TokenEmbeddings with different input shapes.""" embeddings = TokenEmbeddings(vocab_size, embed_dim) - + inputs = torch.randint(0, vocab_size, (batch_size, seq_len)) output = embeddings(inputs) - + assert output.shape == (batch_size, seq_len, embed_dim) diff --git a/llm/tests/models/test_gpt.py b/llm/tests/models/test_gpt.py index 7ae2d9d..90040ce 100644 --- a/llm/tests/models/test_gpt.py +++ b/llm/tests/models/test_gpt.py @@ -9,162 +9,156 @@ from llm.models.gpt import GPT class TestGPT: """Test cases for GPT model.""" - + def test_initialization(self, gpt_config): """Test that GPT can be initialized.""" model = GPT(gpt_config) assert model is not None - + # Check that model has required components - assert hasattr(model, '_token_embeddings') - assert hasattr(model, '_position_embeddings') - assert hasattr(model, '_decoders') - assert hasattr(model, '_linear') - assert hasattr(model, '_dropout') - + assert hasattr(model, "_token_embeddings") + assert hasattr(model, "_position_embeddings") + assert hasattr(model, "_decoders") + assert hasattr(model, "_linear") + assert hasattr(model, "_dropout") + # Check number of decoder layers - assert len(model._decoders) == gpt_config['num_layers'] - + assert len(model._decoders) == gpt_config["num_layers"] + def test_forward_pass(self, gpt_config, random_inputs): """Test forward pass of GPT.""" model = GPT(gpt_config) - + # Forward pass logits = model(random_inputs) - + # Check output shape batch_size, seq_len = random_inputs.shape - vocab_size = gpt_config['vocab_size'] + vocab_size = gpt_config["vocab_size"] assert logits.shape == (batch_size, seq_len, vocab_size) assert isinstance(logits, torch.Tensor) - - def test_forward_with_attention_mask(self, gpt_config, random_inputs, attention_mask): + + def test_forward_with_attention_mask( + self, gpt_config, random_inputs, attention_mask + ): """Test forward pass with attention mask.""" model = GPT(gpt_config) - + # Forward pass with mask logits = model(random_inputs, attention_mask=attention_mask) - + # Check output shape batch_size, seq_len = random_inputs.shape - vocab_size = gpt_config['vocab_size'] + vocab_size = gpt_config["vocab_size"] assert logits.shape == (batch_size, seq_len, vocab_size) - + def test_generate_text(self, gpt_config): """Test text generation.""" model = GPT(gpt_config) model.eval() # Set to evaluation mode for generation - + # Create initial input batch_size = 2 initial_seq_len = 5 - input_ids = torch.randint(0, gpt_config['vocab_size'], (batch_size, initial_seq_len)) - + input_ids = torch.randint( + 0, gpt_config["vocab_size"], (batch_size, initial_seq_len) + ) + # Generate text with torch.no_grad(): generated = model.generate( x=input_ids, max_new_tokens=10, - do_sample=False # Use greedy for deterministic testing + do_sample=False, # Use greedy for deterministic testing ) - + # Check output shape expected_seq_len = initial_seq_len + 10 assert generated.shape == (batch_size, expected_seq_len) - + # Check that initial sequence is preserved assert torch.allclose(generated[:, :initial_seq_len], input_ids) - + def test_generate_with_temperature(self, gpt_config): """Test text generation with temperature sampling.""" model = GPT(gpt_config) model.eval() - + # Create initial input - input_ids = torch.randint(0, gpt_config['vocab_size'], (1, 3)) - + input_ids = torch.randint(0, gpt_config["vocab_size"], (1, 3)) + # Generate with temperature with torch.no_grad(): generated = model.generate( - x=input_ids, - max_new_tokens=5, - do_sample=True, - temperature=0.8 + x=input_ids, max_new_tokens=5, do_sample=True, temperature=0.8 ) - + assert generated.shape == (1, 8) # 3 initial + 5 new tokens - + def test_generate_with_top_k(self, gpt_config): """Test text generation with top-k sampling.""" model = GPT(gpt_config) model.eval() - + # Create initial input - input_ids = torch.randint(0, gpt_config['vocab_size'], (1, 3)) - + input_ids = torch.randint(0, gpt_config["vocab_size"], (1, 3)) + # Generate with top-k with torch.no_grad(): generated = model.generate( - x=input_ids, - max_new_tokens=5, - do_sample=True, - top_k=10 + x=input_ids, max_new_tokens=5, do_sample=True, top_k=10 ) - + assert generated.shape == (1, 8) - + def test_generate_with_top_p(self, gpt_config): """Test text generation with top-p (nucleus) sampling.""" model = GPT(gpt_config) model.eval() - + # Create initial input - input_ids = torch.randint(0, gpt_config['vocab_size'], (1, 3)) - + input_ids = torch.randint(0, gpt_config["vocab_size"], (1, 3)) + # Generate with top-p with torch.no_grad(): generated = model.generate( - x=input_ids, - max_new_tokens=5, - do_sample=True, - top_p=0.9 + x=input_ids, max_new_tokens=5, do_sample=True, top_p=0.9 ) - + assert generated.shape == (1, 8) - + def test_gradient_flow(self, gpt_config, random_inputs): """Test that gradients flow through GPT.""" model = GPT(gpt_config) - + # Forward pass logits = model(random_inputs) - + # Create a dummy loss and backward pass - targets = torch.randint(0, gpt_config['vocab_size'], random_inputs.shape) + targets = torch.randint(0, gpt_config["vocab_size"], random_inputs.shape) loss = torch.nn.functional.cross_entropy( - logits.view(-1, logits.size(-1)), - targets.view(-1) + logits.view(-1, logits.size(-1)), targets.view(-1) ) loss.backward() - + # Check that gradients are computed for various components assert model._token_embeddings._embedding.weight.grad is not None assert model._linear.weight.grad is not None if len(model._decoders) > 0: assert model._decoders[0]._heads._heads[0]._q.weight.grad is not None - + def test_device_consistency(self, gpt_config, random_inputs, device): """Test that GPT works on correct device.""" model = GPT(gpt_config).to(device) inputs = random_inputs.to(device) - + # Forward pass logits = model(inputs) - + # Check device consistency assert logits.device == device assert model._token_embeddings._embedding.weight.device == device - + def test_different_configurations(self): """Test GPT with different configurations.""" test_configs = [ @@ -174,7 +168,7 @@ class TestGPT: "num_heads": 2, "num_layers": 2, "max_position_embeddings": 256, - "dropout": 0.1 + "dropout": 0.1, }, { "vocab_size": 5000, @@ -182,7 +176,7 @@ class TestGPT: "num_heads": 4, "num_layers": 4, "max_position_embeddings": 512, - "dropout": 0.1 + "dropout": 0.1, }, { "vocab_size": 10000, @@ -190,98 +184,94 @@ class TestGPT: "num_heads": 8, "num_layers": 6, "max_position_embeddings": 1024, - "dropout": 0.1 - } + "dropout": 0.1, + }, ] - + for config in test_configs: model = GPT(config) batch_size, seq_len = 2, 16 - inputs = torch.randint(0, config['vocab_size'], (batch_size, seq_len)) - + inputs = torch.randint(0, config["vocab_size"], (batch_size, seq_len)) + logits = model(inputs) - - expected_shape = (batch_size, seq_len, config['vocab_size']) + + expected_shape = (batch_size, seq_len, config["vocab_size"]) assert logits.shape == expected_shape - + @pytest.mark.parametrize("batch_size,seq_len", [(1, 8), (2, 16), (4, 32)]) def test_different_input_shapes(self, gpt_config, batch_size, seq_len): """Test GPT with different input shapes.""" model = GPT(gpt_config) - - inputs = torch.randint(0, gpt_config['vocab_size'], (batch_size, seq_len)) + + inputs = torch.randint(0, gpt_config["vocab_size"], (batch_size, seq_len)) logits = model(inputs) - - expected_shape = (batch_size, seq_len, gpt_config['vocab_size']) + + expected_shape = (batch_size, seq_len, gpt_config["vocab_size"]) assert logits.shape == expected_shape - + def test_training_vs_evaluation(self, gpt_config, random_inputs): """Test that GPT behaves differently in train vs eval mode.""" model = GPT(gpt_config) - + # Training mode model.train() output_train = model(random_inputs) - + # Evaluation mode model.eval() output_eval = model(random_inputs) - + # Outputs should be different due to dropout assert not torch.allclose(output_train, output_eval) - + def test_parameter_count(self, gpt_config): """Test that GPT has reasonable number of parameters.""" model = GPT(gpt_config) - + total_params = sum(p.numel() for p in model.parameters()) - + # For a small GPT model, parameters should be in reasonable range - vocab_size = gpt_config['vocab_size'] - embed_dim = gpt_config['embed_dim'] - num_layers = gpt_config['num_layers'] - num_heads = gpt_config['num_heads'] - + vocab_size = gpt_config["vocab_size"] + embed_dim = gpt_config["embed_dim"] + num_layers = gpt_config["num_layers"] + num_heads = gpt_config["num_heads"] + # Rough estimate: token_embeddings + output_layer + (attention + ff) * layers expected_min = vocab_size * embed_dim * 2 # embeddings and output expected_max = expected_min * 10 # Allow for decoder parameters - + assert expected_min < total_params < expected_max - + def test_causal_attention(self, gpt_config): """Test that GPT uses causal attention during generation.""" model = GPT(gpt_config) model.eval() - + # Create input with known pattern input_ids = torch.tensor([[1, 2, 3]]).long() - + with torch.no_grad(): # Get logits for next token prediction logits = model(input_ids) - + # The model should only attend to previous tokens (causal) # We can't directly test attention masks in the public API, # but we can verify the generation works correctly - - generated = model.generate( - x=input_ids, - max_new_tokens=3, - do_sample=False - ) - + + generated = model.generate(x=input_ids, max_new_tokens=3, do_sample=False) + # Generated sequence should be longer than input assert generated.shape[1] == input_ids.shape[1] + 3 - + def test_output_distribution(self, gpt_config, random_inputs): """Test that GPT output has proper distribution.""" model = GPT(gpt_config) - + logits = model(random_inputs) - + # Logits should not have extreme values assert logits.abs().max() < 100 - + # Softmax should produce valid probabilities probs = torch.softmax(logits, dim=-1) assert torch.allclose(probs.sum(dim=-1), torch.ones_like(probs.sum(dim=-1))) diff --git a/llm/tests/test_basic.py b/llm/tests/test_basic.py index 3733c5f..8d18689 100644 --- a/llm/tests/test_basic.py +++ b/llm/tests/test_basic.py @@ -11,25 +11,25 @@ import os def test_gpt_model_creation(): """Test that GPT model can be created and forward pass works.""" from llm.models.gpt import GPT - + config = { "vocab_size": 1000, "embed_dim": 128, "num_heads": 4, "num_layers": 2, "max_position_embeddings": 256, - "dropout": 0.1 + "dropout": 0.1, } - + model = GPT(config) - + # Test forward pass batch_size, seq_len = 2, 16 input_ids = torch.randint(0, config["vocab_size"], (batch_size, seq_len)) - + with torch.no_grad(): logits = model(input_ids) - + assert logits.shape == (batch_size, seq_len, config["vocab_size"]) print("✅ GPT model creation and forward pass test passed") @@ -37,27 +37,21 @@ def test_gpt_model_creation(): def test_bpe_tokenizer_basic(): """Test basic BPE tokenizer functionality.""" from llm.tokenizers import BPETokenizer - + tokenizer = BPETokenizer() - + # Train on simple texts - texts = [ - "hello world", - "test tokenization", - "simple example" - ] - + texts = ["hello world", "test tokenization", "simple example"] + tokenizer.train( - texts=texts, - vocab_size=50, - special_tokens=["", "", "", ""] + texts=texts, vocab_size=50, special_tokens=["", "", "", ""] ) - + # Test encoding/decoding text = "hello world" tokens = tokenizer.encode(text) decoded = tokenizer.decode(tokens) - + assert isinstance(tokens, list) assert isinstance(decoded, str) assert len(tokens) > 0 @@ -67,18 +61,18 @@ def test_bpe_tokenizer_basic(): def test_token_embeddings(): """Test token embeddings.""" from llm.core.token_embeddings import TokenEmbeddings - + vocab_size = 1000 embed_dim = 128 - + embeddings = TokenEmbeddings(vocab_size, embed_dim) - + # Test forward pass batch_size, seq_len = 2, 16 input_ids = torch.randint(0, vocab_size, (batch_size, seq_len)) - + output = embeddings(input_ids) - + assert output.shape == (batch_size, seq_len, embed_dim) print("✅ Token embeddings test passed") @@ -86,20 +80,20 @@ def test_token_embeddings(): def test_multi_head_attention(): """Test multi-head attention.""" from llm.core.multi_head_attention import MultiHeadAttention - + num_heads = 4 emb_size = 128 head_size = emb_size // num_heads max_seq_len = 256 - + attention = MultiHeadAttention(num_heads, emb_size, head_size, max_seq_len) - + # Test forward pass batch_size, seq_len = 2, 16 inputs = torch.randn(batch_size, seq_len, emb_size) - + output, _ = attention(inputs) - + assert output.shape == inputs.shape print("✅ Multi-head attention test passed") @@ -107,17 +101,17 @@ def test_multi_head_attention(): def test_feed_forward(): """Test feed forward network.""" from llm.core.feed_forward import FeedForward - + embed_dim = 128 - + ff = FeedForward(embed_dim) - + # Test forward pass batch_size, seq_len = 2, 16 inputs = torch.randn(batch_size, seq_len, embed_dim) - + output = ff(inputs) - + assert output.shape == inputs.shape print("✅ Feed forward test passed") @@ -125,29 +119,25 @@ def test_feed_forward(): def test_gpt_generation(): """Test GPT text generation.""" from llm.models.gpt import GPT - + config = { "vocab_size": 1000, "embed_dim": 128, "num_heads": 4, "num_layers": 2, "max_position_embeddings": 256, - "dropout": 0.1 + "dropout": 0.1, } - + model = GPT(config) model.eval() - + # Test greedy generation input_ids = torch.randint(0, config["vocab_size"], (1, 5)) - + with torch.no_grad(): - generated = model.generate( - x=input_ids, - max_new_tokens=3, - do_sample=False - ) - + generated = model.generate(x=input_ids, max_new_tokens=3, do_sample=False) + assert generated.shape == (1, 8) # 5 initial + 3 new tokens print("✅ GPT generation test passed") @@ -155,50 +145,48 @@ def test_gpt_generation(): def test_bpe_tokenizer_save_load(): """Test BPE tokenizer save/load functionality.""" from llm.tokenizers import BPETokenizer - + tokenizer = BPETokenizer() - + # Train on simple texts texts = ["hello world", "test save load"] tokenizer.train( - texts=texts, - vocab_size=30, - special_tokens=["", "", "", ""] + texts=texts, vocab_size=30, special_tokens=["", "", "", ""] ) - + with tempfile.TemporaryDirectory() as temp_dir: save_path = os.path.join(temp_dir, "test_tokenizer.json") - + # Save tokenizer tokenizer.save(save_path) assert os.path.exists(save_path) - + # Load tokenizer loaded_tokenizer = BPETokenizer.load(save_path) - + # Test that vocab size is the same assert tokenizer.get_vocab_size() == loaded_tokenizer.get_vocab_size() - + # Test that vocabularies are the same assert tokenizer.get_vocab() == loaded_tokenizer.get_vocab() - + # Test that both can encode/decode (even if tokens differ due to BPE state) text = "hello world" original_tokens = tokenizer.encode(text) loaded_tokens = loaded_tokenizer.encode(text) - + # Both should produce valid token lists assert isinstance(original_tokens, list) assert isinstance(loaded_tokens, list) assert len(original_tokens) > 0 assert len(loaded_tokens) > 0 - + # Both should be able to decode original_decoded = tokenizer.decode(original_tokens) loaded_decoded = loaded_tokenizer.decode(loaded_tokens) assert isinstance(original_decoded, str) assert isinstance(loaded_decoded, str) - + print("✅ BPE tokenizer save/load test passed") @@ -206,18 +194,16 @@ def test_gpt_with_tokenizer(): """Test GPT model with tokenizer integration.""" from llm.models.gpt import GPT from llm.tokenizers import BPETokenizer - + # Create and train tokenizer tokenizer = BPETokenizer() texts = ["hello world", "test integration"] tokenizer.train( - texts=texts, - vocab_size=50, - special_tokens=["", "", "", ""] + texts=texts, vocab_size=50, special_tokens=["", "", "", ""] ) - + vocab_size = tokenizer.get_vocab_size() - + # Create GPT model with tokenizer's vocab size config = { "vocab_size": vocab_size, @@ -225,19 +211,19 @@ def test_gpt_with_tokenizer(): "num_heads": 4, "num_layers": 2, "max_position_embeddings": 256, - "dropout": 0.1 + "dropout": 0.1, } - + model = GPT(config) - + # Test with tokenized input text = "hello world" tokens = tokenizer.encode(text, add_special_tokens=False) input_ids = torch.tensor([tokens]) - + with torch.no_grad(): logits = model(input_ids) - + assert logits.shape == (1, len(tokens), vocab_size) print("✅ GPT with tokenizer integration test passed") @@ -245,7 +231,7 @@ def test_gpt_with_tokenizer(): def run_all_tests(): """Run all basic tests.""" print("🧪 Running basic tests for llm library...") - + test_gpt_model_creation() test_bpe_tokenizer_basic() test_token_embeddings() @@ -254,7 +240,7 @@ def run_all_tests(): test_gpt_generation() test_bpe_tokenizer_save_load() test_gpt_with_tokenizer() - + print("🎉 All basic tests passed!") diff --git a/llm/tests/tokenizers/test_base_tokenizer.py b/llm/tests/tokenizers/test_base_tokenizer.py index 629aca3..32bb77d 100644 --- a/llm/tests/tokenizers/test_base_tokenizer.py +++ b/llm/tests/tokenizers/test_base_tokenizer.py @@ -8,15 +8,15 @@ from llm.tokenizers import BaseTokenizer class ConcreteTokenizer(BaseTokenizer): """Concrete implementation for testing BaseTokenizer.""" - + def train(self, texts: list, vocab_size: int = 1000, **kwargs): """Dummy implementation for testing.""" pass - + def encode(self, text: str, **kwargs) -> list: """Dummy implementation for testing.""" return [1, 2, 3] - + def decode(self, tokens: list, **kwargs) -> str: """Dummy implementation for testing.""" return "decoded text" @@ -24,33 +24,33 @@ class ConcreteTokenizer(BaseTokenizer): class TestBaseTokenizer: """Test cases for BaseTokenizer.""" - + def test_initialization(self): """Test that BaseTokenizer can be initialized through concrete class.""" tokenizer = ConcreteTokenizer() assert tokenizer is not None assert tokenizer.vocab == {} assert tokenizer.vocab_size == 0 - + def test_encode_implemented(self): """Test that encode method works in concrete implementation.""" tokenizer = ConcreteTokenizer() result = tokenizer.encode("test text") assert result == [1, 2, 3] - + def test_decode_implemented(self): """Test that decode method works in concrete implementation.""" tokenizer = ConcreteTokenizer() result = tokenizer.decode([1, 2, 3]) assert result == "decoded text" - + def test_get_vocab_size(self): """Test that get_vocab_size method works.""" tokenizer = ConcreteTokenizer() tokenizer.vocab = {"a": 0, "b": 1, "c": 2} tokenizer.vocab_size = 3 assert tokenizer.get_vocab_size() == 3 - + def test_get_vocab(self): """Test that get_vocab method works.""" tokenizer = ConcreteTokenizer() diff --git a/llm/tests/tokenizers/test_bpe_tokenizer.py b/llm/tests/tokenizers/test_bpe_tokenizer.py index 1644ecd..57ba765 100644 --- a/llm/tests/tokenizers/test_bpe_tokenizer.py +++ b/llm/tests/tokenizers/test_bpe_tokenizer.py @@ -10,18 +10,18 @@ from llm.tokenizers import BPETokenizer class TestBPETokenizer: """Test cases for BPETokenizer.""" - + @pytest.fixture def sample_texts(self): """Sample texts for training tokenizer.""" return [ "Искусственный интеллект", - "Нейронные сети", + "Нейронные сети", "Машинное обучение", "Глубокое обучение", - "Трансформеры" + "Трансформеры", ] - + @pytest.fixture def trained_tokenizer(self, sample_texts): """Create and train a BPE tokenizer.""" @@ -29,128 +29,130 @@ class TestBPETokenizer: tokenizer.train( texts=sample_texts, vocab_size=100, - special_tokens=["", "", "", ""] + special_tokens=["", "", "", ""], ) return tokenizer - + def test_initialization(self): """Test that BPETokenizer can be initialized.""" tokenizer = BPETokenizer() assert tokenizer is not None - + def test_train_tokenizer(self, sample_texts): """Test that tokenizer can be trained.""" tokenizer = BPETokenizer() tokenizer.train( texts=sample_texts, vocab_size=50, - special_tokens=["", "", "", ""] + special_tokens=["", "", "", ""], ) - + assert tokenizer.get_vocab_size() > 0 assert len(tokenizer.get_vocab()) == tokenizer.get_vocab_size() - + def test_encode_decode(self, trained_tokenizer): """Test encoding and decoding text.""" text = "Искусственный интеллект" - + # Encode text tokens = trained_tokenizer.encode(text) assert isinstance(tokens, list) assert len(tokens) > 0 assert all(isinstance(token, int) for token in tokens) - + # Decode tokens decoded_text = trained_tokenizer.decode(tokens) assert isinstance(decoded_text, str) # Decoded text should be similar to original (may have special tokens) assert len(decoded_text) > 0 - + def test_encode_with_special_tokens(self, trained_tokenizer): """Test encoding with special tokens.""" text = "Нейронные сети" - + # Without special tokens tokens_no_special = trained_tokenizer.encode(text, add_special_tokens=False) - + # With special tokens tokens_with_special = trained_tokenizer.encode(text, add_special_tokens=True) - + # Should have more tokens when special tokens are added assert len(tokens_with_special) >= len(tokens_no_special) - + def test_vocab_size(self, trained_tokenizer): """Test vocabulary size.""" vocab_size = trained_tokenizer.get_vocab_size() assert isinstance(vocab_size, int) assert vocab_size > 0 - + vocab = trained_tokenizer.get_vocab() assert isinstance(vocab, dict) assert len(vocab) == vocab_size - + def test_special_tokens(self, trained_tokenizer): """Test that special tokens are in vocabulary.""" vocab = trained_tokenizer.get_vocab() - + # Check that special tokens are in vocabulary special_tokens = ["", "", "", ""] for token in special_tokens: assert token in vocab assert isinstance(vocab[token], int) - + def test_save_load(self, trained_tokenizer, sample_texts): """Test saving and loading tokenizer.""" with tempfile.TemporaryDirectory() as temp_dir: save_path = os.path.join(temp_dir, "test_tokenizer.json") - + # Save tokenizer trained_tokenizer.save(save_path) assert os.path.exists(save_path) - + # Load tokenizer loaded_tokenizer = BPETokenizer.load(save_path) assert loaded_tokenizer is not None - + # Check that loaded tokenizer works the same original_vocab = trained_tokenizer.get_vocab() loaded_vocab = loaded_tokenizer.get_vocab() - + assert original_vocab == loaded_vocab - assert trained_tokenizer.get_vocab_size() == loaded_tokenizer.get_vocab_size() - + assert ( + trained_tokenizer.get_vocab_size() == loaded_tokenizer.get_vocab_size() + ) + # Test encoding consistency text = sample_texts[0] original_tokens = trained_tokenizer.encode(text) loaded_tokens = loaded_tokenizer.encode(text) - + assert original_tokens == loaded_tokens - + def test_unknown_tokens(self, trained_tokenizer): """Test handling of unknown tokens.""" # Use text that likely contains unknown subwords text = "xyzabc123" # Random text that shouldn't be in training data - + tokens = trained_tokenizer.encode(text) assert len(tokens) > 0 - + # Should be able to decode back (even if it's mostly unk tokens) decoded = trained_tokenizer.decode(tokens) assert isinstance(decoded, str) - + def test_empty_text(self, trained_tokenizer): """Test encoding and decoding empty text.""" tokens = trained_tokenizer.encode("") assert isinstance(tokens, list) - + decoded = trained_tokenizer.decode([]) assert decoded == "" - + def test_tokenize_method(self, trained_tokenizer): """Test the tokenize method.""" text = "Искусственный интеллект" tokens = trained_tokenizer.tokenize(text) - + assert isinstance(tokens, list) assert len(tokens) > 0 assert all(isinstance(token, str) for token in tokens)