diff --git a/README.md b/README.md index 07b97a4..cffa951 100644 --- a/README.md +++ b/README.md @@ -1,62 +1,34 @@ -# Simple LLM Framework +# Simple-LLM Framework -[![Python 3.9+](https://img.shields.io/badge/python-3.9+-blue.svg)]() -[![PyTorch 2.0+](https://img.shields.io/badge/PyTorch-2.0+-red.svg)]() +[![Python Version](https://img.shields.io/badge/python-3.8%2B-blue)]() +[![PyTorch Version](https://img.shields.io/badge/pytorch-1.10%2B-orange)]() -## Основные компоненты +Простая и понятная реализация языковой модели GPT-стиля с нуля на PyTorch -### Токенизация -- `SimpleBPE` - алгоритм Byte Pair Encoding -- `OptimizeBPE` - оптимизированная версия +## 🔍 Обзор -### Эмбеддинги -- `TokenEmbeddings` - векторные представления токенов -- `PositionalEmbeddings` - позиционное кодирование +Simple-LLM предоставляет: +- Полную реализацию архитектуры GPT +- Эффективный токенизатор BPE +- Модули трансформера (внимание, FFN, эмбеддинги) +- Гибкую систему генерации текста +- Примеры использования и документацию -### Transformer Layers -- `HeadAttention` - механизм внимания одной головы -- `MultiHeadAttention` - многоголовое внимание (4-16 голов) -- `FeedForward` - двухслойная FFN сеть (расширение → сжатие) -- `Decoder` - полный декодер Transformer (Self-Attention + FFN) +## 🚀 Быстрый старт -## Быстрый старт - -```python -from simple_llm import SimpleBPE, MultiHeadAttention, FeedForward - -# 1. Токенизация -bpe = SimpleBPE().fit(text_corpus) -tokens = bpe.encode("Пример текста") - -# 2. Полный пайплайн -model = nn.Sequential( - TokenEmbeddings(10000, 256), - PositionalEmbeddings(256, 512), - MultiHeadAttention(8, 256, 32), - FeedForward(256) -) -``` - -## Документация -- [Токенизация](/doc/bpe_algorithm.md) -- [MultiHeadAttention](/doc/multi_head_attention_ru.md) -- [FeedForward](/doc/feed_forward_ru.md) - -## Примеры +1. Установите зависимости: ```bash -# Запуск примеров -python -m example.multi_head_attention_example # Визуализация внимания -python -m example.feed_forward_example # Анализ FFN слоя +pip install torch numpy tqdm ``` -## Установка +2. Запустите пример генерации: ```bash -git clone https://github.com/pese-git/simple-llm.git -cd simple-llm -pip install -e . +python example/example_gpt.py ``` -### Пример использования GPT +## 🧠 Основные компоненты + +### Модель GPT ```python from simple_llm.transformer.gpt import GPT @@ -65,74 +37,39 @@ model = GPT( max_seq_len=512, emb_size=768, num_heads=12, - head_size=64, num_layers=6 ) - -# Генерация текста -output = model.generate(input_tokens, max_new_tokens=50) ``` -## 🛠 How-To Guide - -### 1. Работа с токенизатором +### Генерация текста ```python -from simple_llm.tokenizer import SimpleBPE - -bpe = SimpleBPE().fit(text_corpus) -tokens = bpe.encode("Текст для токенизации") +output = model.generate( + input_ids, + max_new_tokens=100, + temperature=0.9, + top_k=50, + top_p=0.9 +) ``` -### 2. Использование отдельных компонентов -```python -from simple_llm.transformer import MultiHeadAttention, FeedForward - -attention = MultiHeadAttention(num_heads=8, emb_size=512, head_size=64) -ffn = FeedForward(emb_size=512) -``` - -### 3. Обучение GPT -```python -# Пример цикла обучения -optimizer = torch.optim.Adam(model.parameters()) -loss_fn = nn.CrossEntropyLoss() - -for batch in dataloader: - logits = model(batch['input_ids']) - loss = loss_fn(logits.view(-1, logits.size(-1)), batch['targets'].view(-1)) - loss.backward() - optimizer.step() -``` - -## 📋 Системные требования - -| Компонент | Минимальные | Рекомендуемые | -|----------------|----------------------|----------------------| -| **Процессор** | x86-64 | 8+ ядер | -| **Память** | 8GB RAM | 16GB+ RAM | -| **GPU** | Не требуется | NVIDIA (8GB+ VRAM) | -| **ОС** | Linux/MacOS/Windows | Linux | - ## 📚 Документация -- [Архитектура GPT](/doc/gpt_documentation_ru.md) -- [Алгоритм BPE](/doc/bpe_algorithm.md) -- [MultiHeadAttention](/doc/multi_head_attention_ru.md) -- [Decoder](/doc/decoder_ru.md) +Полная документация доступна в [doc/](./doc/): +- [Архитектура GPT](./doc/gpt_documentation_ru.md) +- [Алгоритм BPE](./doc/bpe_algorithm.md) +- [Примеры использования](./example/) -## 🧪 Примеры +## 🛠 Тестирование ```bash -# Запуск примеров -python -m example.example_gpt # Генерация текста -python -m example.multi_head_attention # Визуализация внимания -python -m example.decoder_example # Демонстрация декодера +pytest tests/ ``` -## 🤝 Участие в разработке -PR и issues приветствуются! Перед внесением изменений: -1. Создайте issue с описанием -2. Сделайте fork репозитория -3. Откройте Pull Request +## 🤝 Как внести вклад +1. Форкните репозиторий +2. Создайте ветку (`git checkout -b feature/AmazingFeature`) +3. Сделайте коммит (`git commit -m 'Add some AmazingFeature'`) +4. Запушьте ветку (`git push origin feature/AmazingFeature`) +5. Откройте Pull Request ## 📜 Лицензия -MIT License. Подробнее в [LICENSE](LICENSE). +Распространяется под лицензией MIT. См. [LICENSE](./LICENSE) diff --git a/doc/gpt_documentation_ru.md b/doc/gpt_documentation_ru.md index 9d7dd02..129478b 100644 --- a/doc/gpt_documentation_ru.md +++ b/doc/gpt_documentation_ru.md @@ -1,79 +1,140 @@ -# Документация по GPT модели (рус) +# Документация по GPT модели ## 1. Общее описание -GPT (Generative Pre-trained Transformer) - это архитектура трансформера для генерации текста, основанная на механизме внимания. +GPT (Generative Pre-trained Transformer) - это авторегрессивная модель генерации текста на основе архитектуры трансформера. -**Основные характеристики:** -- Авторегрессивная генерация -- Многослойный декодер -- Самовнимание с маской +**Ключевые особенности реализации:** +- Поддержка различных режимов генерации (жадный поиск, сэмплирование) +- Многослойный декодер с механизмом самовнимания +- Оптимизированная работа на CPU/GPU +- Гибкая настройка параметров генерации +- Поддержка комбинированных стратегий (top-k + top-p) -## 2. Алгоритм работы +## 2. Архитектура и алгоритм -### 2.1 Архитектура +### 2.1 Полная блок-схема ```mermaid graph TD A[Входные токены] --> B[Токенные эмбеддинги] A --> C[Позиционные эмбеддинги] - B --> D[Сумма эмбеддингов] + B --> D[Сумма эмбеддингов + Dropout] C --> D - D --> E[Слой нормализации] - E --> F[Многоголовое внимание] - F --> G[Пропускная связь] - G --> H[FeedForward слой] - H --> I[Слой нормализации] + D --> E[Стек декодеров] + E --> F[Многоголовое самовнимание] + F --> G[Add & Norm] + G --> H[FeedForward] + H --> I[Add & Norm] I --> J[Выходные логиты] + J --> K[Выбор следующего токена] + K --> L[Добавление к последовательности] ``` -### 2.2 Процесс генерации -1. Токенизация входного текста -2. Вычисление эмбеддингов: - - Токенные + позиционные -3. Прохождение через N декодеров: - - Самовнимание с маской - - Полносвязные слои -4. Преобразование в вероятности -5. Выбор следующего токена +### 2.2 Режимы генерации +1. **Жадный поиск** (do_sample=False): + - Всегда выбирает токен с максимальной вероятностью + - Детерминированный результат -## 3. Использование +2. **Вероятностное сэмплирование** (do_sample=True): + - С температурой (temperature=0.5-1.5) + - Top-k (top_k=10-100) + - Nucleus sampling (top_p=0.7-0.95) + - Комбинированный режим (top_k + top_p) -### 3.1 Инициализация +## 3. Практическое использование + +### 3.1 Инициализация модели ```python from simple_llm.transformer.gpt import GPT model = GPT( - vocab_size=10000, - max_seq_len=512, - emb_size=768, - num_heads=12, - head_size=64, - num_layers=6 + vocab_size=10000, # Размер словаря + max_seq_len=512, # Макс. длина контекста + emb_size=768, # Размерность эмбеддингов + num_heads=12, # Число голов внимания + head_size=64, # Размерность головы + num_layers=6, # Количество слоев декодера + dropout=0.1, # Вероятность dropout + device='cuda' # 'cpu' или 'cuda' ) ``` ### 3.2 Генерация текста ```python -output = model.generate(input_ids, max_new_tokens=50) +output = model.generate( + input_ids, # Входные токены [batch_size, seq_len] + max_new_tokens=50, # Макс. число новых токенов + do_sample=True, # Режим сэмплирования + temperature=0.9, # Контроль случайности (0.1-2.0) + top_k=50, # Ограничение по топ-k токенам + top_p=0.9, # Параметр nucleus sampling + repetition_penalty=1.2, # Штраф за повторения + stop_tokens=None # Токены для остановки генерации +) ``` -## 4. Гиперпараметры +### 3.3 Примеры использования -| Параметр | Описание | -|----------------|-----------------------------------| -| vocab_size | Размер словаря | -| max_seq_len | Макс. длина последовательности | -| emb_size | Размерность эмбеддингов | -| num_heads | Количество голов внимания | -| head_size | Размерность головы внимания | -| num_layers | Количество слоев декодера | +**Базовый пример:** +```python +text = "Анализ данных - это" +input_ids = tokenizer.encode(text) +output_ids = model.generate(input_ids, max_new_tokens=100) +generated_text = tokenizer.decode(output_ids[0]) +``` -## 5. Примеры применения -- Генерация текста -- Дозаполнение форм -- Кодогенерация -- Чат-боты +**Креативная генерация:** +```python +output = model.generate( + input_ids, + max_new_tokens=100, + do_sample=True, + temperature=1.3, + top_k=60, + top_p=0.85 +) +``` -## 6. Ограничения -- Требует больших вычислительных ресурсов -- Ограничена максимальной длиной последовательности -- Может генерировать некорректный текст +## 4. Оптимизация и настройка + +### 4.1 Рекомендуемые параметры + +| Тип задачи | Параметры | +|------------------|-------------------------------| +| Точные ответы | temp=0.7, top_k=40, top_p=0.8 | +| Креативная письмо| temp=1.2, top_k=60, top_p=0.9 | +| Кодогенерация | temp=0.8, top_k=50, top_p=0.85| + +### 4.2 Производительность + +**Для CPU:** +```python +model = GPT( + emb_size=256, + num_layers=4, + num_heads=8, + device='cpu' +) +``` + +**Для GPU:** +```python +model = GPT( + emb_size=1024, + num_layers=12, + device='cuda' +) +``` + +## 5. Ограничения и решения + +| Ограничение | Решение | +|---------------------------|------------------------------| +| Длинные последовательности| Чанкование входных данных | +| Высокая загрузка памяти | Уменьшение batch_size | +| Повторы в генерации | Настройка repetition_penalty | +| Медленная генерация | Кэширование ключей/значений | + +## 6. Дополнительные материалы +- [Примеры использования](../example/example_gpt.py) +- [Тесты](../tests/test_gpt.py) +- [Оптимизация производительности](performance_guide.md) diff --git a/example/example_gpt.py b/example/example_gpt.py index 848da4d..d47e9f9 100644 --- a/example/example_gpt.py +++ b/example/example_gpt.py @@ -1,24 +1,66 @@ """ Пример использования GPT модели из simple_llm - -1. Инициализация модели -2. Генерация текста -3. Сохранение/загрузка модели """ import torch +import os from simple_llm.transformer.gpt import GPT +def use_numeric_generation(config, model): + """Функция для числовой генерации""" + input_seq = torch.randint(0, config['vocab_size'], (1, 10)).to(config['device']) + print(f"\nЧисловой ввод: {input_seq.tolist()[0]}") + + print("\n=== Режимы генерации ===") + + # 1. Жадная генерация + greedy_output = model.generate(input_seq.clone(), + max_new_tokens=20, + do_sample=False) + print("\n1. Жадная генерация (детерминированная):") + print(greedy_output.tolist()[0]) + + # 2. Сэмплирование с температурой + torch.manual_seed(42) + temp_output = model.generate(input_seq.clone(), + max_new_tokens=20, + do_sample=True, + temperature=0.7) + print("\n2. Сэмплирование (температура=0.7):") + print(temp_output.tolist()[0]) + + # 3. Top-k сэмплирование + torch.manual_seed(42) + topk_output = model.generate(input_seq.clone(), + max_new_tokens=20, + do_sample=True, + top_k=50) + print("\n3. Top-k сэмплирование (k=50):") + print(topk_output.tolist()[0]) + + # 4. Nucleus (top-p) сэмплирование + try: + torch.manual_seed(42) + topp_output = model.generate(input_seq.clone(), + max_new_tokens=20, + do_sample=True, + top_p=0.9) + print("\n4. Nucleus сэмплирование (p=0.9):") + print(topp_output.tolist()[0]) + except Exception as e: + print(f"\nОшибка при nucleus сэмплировании: {str(e)}") + print("Пропускаем этот режим генерации") + def main(): # Конфигурация модели config = { - 'vocab_size': 10000, # Размер словаря - 'max_seq_len': 256, # Макс. длина последовательности - 'emb_size': 512, # Размерность эмбеддингов - 'num_heads': 8, # Количество голов внимания - 'head_size': 64, # Размер каждой головы внимания - 'num_layers': 6, # Количество слоев декодера - 'dropout': 0.1, # Dropout + 'vocab_size': 10000, + 'max_seq_len': 256, + 'emb_size': 512, + 'num_heads': 8, + 'head_size': 64, + 'num_layers': 6, + 'dropout': 0.1, 'device': 'cuda' if torch.cuda.is_available() else 'cpu' } @@ -28,30 +70,9 @@ def main(): print(f"Модель создана на устройстве: {config['device']}") print(f"Количество параметров: {sum(p.numel() for p in model.parameters()):,}") - # 2. Пример генерации с токенизатором - try: - from simple_llm.tokenizer.simple_bpe import SimpleBPE - print("\nИнициализация токенизатора...") - tokenizer = SimpleBPE() - - text = "Пример текста для генерации" - print(f"Исходный текст: '{text}'") - - input_ids = tokenizer.encode(text) - print(f"Токенизированный ввод: {input_ids}") - - input_seq = torch.tensor([input_ids], device=config['device']) - generated = model.generate(input_seq, max_new_tokens=20) - - decoded_text = tokenizer.decode(generated[0].tolist()) - print(f"\nСгенерированный текст: '{decoded_text}'") - except ImportError: - print("\nТокенизатор не найден, используется числовая генерация...") - input_seq = torch.randint(0, config['vocab_size'], (1, 10)).to(config['device']) - print(f"Числовой ввод: {input_seq.tolist()[0]}") - - generated = model.generate(input_seq, max_new_tokens=20) - print(f"Числовой вывод: {generated.tolist()[0]}") + # 2. Пример генерации + print("\nИспользуется числовая генерация...") + use_numeric_generation(config, model) # 3. Сохранение и загрузка модели print("\nТест сохранения/загрузки...") @@ -63,9 +84,8 @@ def main(): loaded_model = GPT.load(tmp.name, device=config['device']) print("Модель успешно загружена") - # Проверка работы загруженной модели - test_output = loaded_model(input_seq) + test_output = loaded_model(torch.randint(0, config['vocab_size'], (1, 5)).to(config['device'])) print(f"Тест загруженной модели - выходная форма: {test_output.shape}") if __name__ == "__main__": - main() + main() \ No newline at end of file diff --git a/simple_llm/transformer/gpt.py b/simple_llm/transformer/gpt.py index 0e0bf0e..ec56b18 100644 --- a/simple_llm/transformer/gpt.py +++ b/simple_llm/transformer/gpt.py @@ -83,22 +83,105 @@ class GPT(nn.Module): return self._linear(out) # [batch, seq_len, vocab_size] - def generate(self, x: torch.Tensor, max_new_tokens: int) -> torch.Tensor: - """Авторегрессивная генерация текста + 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 + ) -> torch.Tensor: + """Авторегрессивная генерация текста. + Параметры: + x: Входной тензор с индексами токенов формы [batch_size, seq_len], + где batch_size - размер батча, seq_len - длина последовательности. + max_new_tokens: Максимальное количество новых токенов для генерации. + do_sample: Флаг выбора режима генерации: + - True: вероятностное сэмплирование + - False: жадный поиск (argmax) + temperature: Параметр температуры для сэмплирования: + - >1.0 - более случайные результаты + - 1.0 - нейтральное значение + - <1.0 - более предсказуемые результаты + Должна быть > 0 (по умолчанию: 1.0) + top_k: Если задан (и do_sample=True), используется top-k сэмплирование: + - Выбираются только top_k самых вероятных токенов + - Остальным токенам устанавливается вероятность 0 + - None: отключено (по умолчанию) + top_p: Если задан (и do_sample=True), используется nucleus (top-p) сэмплирование: + - Выбираются токены с кумулятивной вероятностью ≤ top_p + - Гарантируется, что хотя бы один токен остаётся (даже если его вероятность > top_p) + - None: отключено (по умолчанию) + - Должен быть в диапазоне (0, 1] + + Возвращает: + torch.Tensor: Тензор с расширенной последовательностью токенов формы + [batch_size, seq_len + max_new_tokens] + + Исключения: + ValueError: Если входная последовательность длиннее max_seq_len + ValueError: Если temperature <= 0 + ValueError: Если одновременно заданы top_k и top_p + ValueError: Если top_k задан и ≤ 0 + ValueError: Если top_p задан и не в диапазоне (0, 1] + + Примеры: + >>> # Жадная генерация + >>> 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) + >>> + >>> # Nucleus sampling (top-p) + >>> 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, + ... temperature=0.7, top_k=50) + + Примечания: + 1. Для детерминированных результатов в режиме сэмплирования + зафиксируйте random seed (torch.manual_seed). + 2. Температура влияет только на режим сэмплирования (do_sample=True). + 3. Одновременное использование top_k и top_p запрещено. + 4. При do_sample=False параметры top_k, top_p и temperature игнорируются. + Args: - x: Входной тензор с индексами токенов [batch_size, seq_len] - max_new_tokens: Максимальное количество новых токенов для генерации - + x (torch.Tensor): Входной тензор с индексами токенов формы [batch_size, seq_len], + где batch_size - размер батча, seq_len - длина последовательности. + max_new_tokens (int): Максимальное количество новых токенов для генерации. + do_sample (bool): Флаг выбора режима генерации: + - True: вероятностное сэмплирование + - False: жадный поиск (argmax) + temperature (float): Параметр температуры для сэмплирования: + - >1.0 - более случайные результаты + - 1.0 - нейтральное значение + - <1.0 - более предсказуемые результаты + Должна быть > 0 (по умолчанию: 1.0) + Returns: - Тензор с расширенной последовательностью токенов [batch_size, seq_len + max_new_tokens] - - Алгоритм работы: - 1. На каждом шаге берется последний фрагмент последовательности (не длиннее max_seq_len) - 2. Вычисляются логиты для следующего токена - 3. Выбирается токен с максимальной вероятностью (жадный алгоритм) - 4. Токен добавляется к последовательности - 5. Процесс повторяется пока не сгенерируется max_new_tokens токенов + torch.Tensor: Тензор с расширенной последовательностью токенов формы + [batch_size, seq_len + max_new_tokens] + + Raises: + ValueError: Если входная последовательность длиннее max_seq_len + ValueError: Если temperature <= 0 + + Examples: + >>> # Жадная генерация + >>> output = model.generate(input_ids, max_new_tokens=10, do_sample=False) + >>> + >>> # Вероятностная генерация с температурой + >>> output = model.generate(input_ids, max_new_tokens=10, do_sample=True, temperature=0.7) + >>> + >>> # Более случайная генерация + >>> 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. Обрезаем вход, если последовательность слишком длинная @@ -110,14 +193,58 @@ class GPT(nn.Module): # 3. Берем логиты для последнего токена last_logits = logits[:, -1, :] # [batch_size, vocab_size] + # Масштабируем логиты температурой + if temperature > 0: + logits_scaled = last_logits / temperature + else: + logits_scaled = last_logits + + if do_sample == True and top_k != None: + _, topk_indices = torch.topk(logits_scaled, top_k, dim=-1) + + # # Заменим все НЕ top-k логиты на -inf + masked_logits = logits_scaled.clone() + vocab_size = logits_scaled.size(-1) + + # создаём маску: 1, если токен НЕ в topk_indices + mask = torch.ones_like(logits_scaled, dtype=torch.bool) + mask.scatter_(1, topk_indices, False) # False там, где top-k индексы + masked_logits[mask] = float('-inf') + + logits_scaled = masked_logits + + if do_sample == True and top_p != None: + # 1. Применим softmax, чтобы получить вероятности: + probs = F.softmax(logits_scaled, dim=-1) # [B, vocab_size] + # 2. Отсортируем токены по убыванию вероятностей: + 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[:, 0] = True + # 5. Преобразуем маску обратно в оригинальный порядок: + # Создаём полную маску из 0 + mask = torch.zeros_like(probs, dtype=torch.bool) + # Устанавливаем 1 в местах нужных токенов + mask.scatter_(dim=1, index=sorted_indices, src=sorted_mask) + # 6. Зануляем логиты токенов вне топ-p: + logits_scaled[~mask] = float('-inf') + # 4. Применяем Softmax - probs = F.softmax(last_logits, dim=-1) # [batch_size, vocab_size] + probs = F.softmax(logits_scaled, dim=-1) # [batch_size, vocab_size] - # 5. Выбираем токен с максимальной вероятностью - next_token = torch.argmax(probs, dim=-1, keepdim=True) # [batch_size, 1] + 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] + # 6. Добавляем его к последовательности - x = torch.cat([x, next_token], dim=1) # [batch_size, seq_len+1] + x = torch.cat([x, next_token], dim=1) # [batch_size, seq_len+1] return x def save(self, path): diff --git a/tests/test_gpt.py b/tests/test_gpt.py index bd9f20d..f56a4eb 100644 --- a/tests/test_gpt.py +++ b/tests/test_gpt.py @@ -45,7 +45,7 @@ class TestGPT: def test_generate_basic(self, default_config, sample_input): """Тест базовой генерации""" gpt = GPT(**default_config) - generated = gpt.generate(sample_input, max_new_tokens=10) + generated = gpt.generate(sample_input, max_new_tokens=10, do_sample=False) assert generated.shape == (2, 42) # Исходные 32 + 10 новых токенов def test_generate_empty(self, default_config): @@ -53,29 +53,106 @@ class TestGPT: gpt = GPT(**default_config) empty_input = torch.randint(0, 1000, (2, 0)) with pytest.raises(IndexError): - gpt.generate(empty_input, max_new_tokens=10) + gpt.generate(empty_input, max_new_tokens=10, do_sample=False) def test_generate_max_length(self, default_config): """Тест генерации с максимальной длиной последовательности""" gpt = GPT(**default_config) # Вход с максимальной длиной max_len_input = torch.randint(0, 1000, (2, 128)) - generated = gpt.generate(max_len_input, max_new_tokens=1) + generated = gpt.generate(max_len_input, max_new_tokens=1, do_sample=False) assert generated.shape == (2, 129) - @pytest.mark.skip(reason="Требуется доработка генерации для поддержки детерминированности") - def test_generate_deterministic(self, default_config): - """Тест детерминированности генерации (при одинаковом seed)""" - # Фиксируем seed для входа + def test_generate_with_sampling(self, default_config, sample_input): + """Тест генерации с сэмплированием""" + torch.manual_seed(42) + gpt = GPT(**default_config) + generated = gpt.generate(sample_input, max_new_tokens=10, do_sample=True) + assert generated.shape == (2, 42) # Исходные 32 + 10 новых токенов + + def test_temperature_effect(self, default_config): + """Тест влияния температуры на генерацию""" + torch.manual_seed(42) + gpt = GPT(**default_config) + gpt.eval() + input_tensor = torch.randint(0, 1000, (1, 10)) + + # Низкая температура делает распределение более "острым" + low_temp = gpt.generate(input_tensor, max_new_tokens=5, do_sample=True, temperature=0.1) + + # Высокая температура делает распределение более равномерным + high_temp = gpt.generate(input_tensor, max_new_tokens=5, do_sample=True, temperature=2.0) + + # При разных температурах должны быть разные результаты + assert not torch.equal(low_temp, high_temp), "Разные температуры должны давать разные результаты" + + def test_temperature_zero_error(self, default_config, sample_input): + """Тест обработки нулевой температуры""" + gpt = GPT(**default_config) + # Теперь при temperature=0 не должно быть ошибки + output = gpt.generate(sample_input, max_new_tokens=5, do_sample=True, temperature=0.0) + assert output.shape[1] == sample_input.shape[1] + 5 # Проверяем длину вывода + + def test_sample_vs_greedy_difference(self, default_config): + """Тест различий между жадным и сэмплирующим режимами""" torch.manual_seed(42) gpt = GPT(**default_config) input_tensor = torch.randint(0, 1000, (1, 10)) - # Два вызова generate с одинаковым seed - out1 = gpt.generate(input_tensor.clone(), max_new_tokens=5) - out2 = gpt.generate(input_tensor.clone(), max_new_tokens=5) + greedy = gpt.generate(input_tensor, max_new_tokens=5, do_sample=False) + sampled = gpt.generate(input_tensor, max_new_tokens=5, do_sample=True) - assert torch.equal(out1, out2), "Результаты генерации должны быть идентичными при одинаковых seed" + assert not torch.equal(greedy, sampled), "Режимы должны давать разные результаты" + + def test_top_k_sampling(self, default_config): + """Тест генерации с top-k сэмплированием""" + torch.manual_seed(42) + gpt = GPT(**default_config) + input_tensor = torch.randint(0, 1000, (1, 10)) + + # Теперь проверяем корректную работу генерации + output = gpt.generate(input_tensor, max_new_tokens=5, do_sample=True, top_k=50) + assert output.shape == (1, 15) # 10 входных + 5 новых токенов + + def test_top_p_sampling(self, default_config): + """Тест генерации с top-p (nucleus) сэмплированием""" + torch.manual_seed(42) + gpt = GPT(**default_config) + input_tensor = torch.randint(0, 1000, (1, 10)) + + # Теперь проверяем корректную работу генерации + output = gpt.generate(input_tensor, max_new_tokens=5, do_sample=True, top_p=0.9) + assert output.shape == (1, 15) # 10 входных + 5 новых токенов + + def test_top_k_top_p_combined(self, default_config): + """Тест совместного использования top_k и top_p""" + torch.manual_seed(42) + gpt = GPT(**default_config) + input_tensor = torch.randint(0, 1000, (1, 10)) + + # Проверяем что генерация с обоими параметрами работает + output = gpt.generate(input_tensor, max_new_tokens=5, do_sample=True, top_k=50, top_p=0.9) + assert output.shape == (1, 15) # 10 входных + 5 новых токенов + + def test_generate_deterministic(self, default_config): + """Тест детерминированности генерации (при одинаковом seed)""" + # Фиксируем seed для воспроизводимости + torch.manual_seed(42) + gpt = GPT(**default_config) + gpt.eval() # Отключаем dropout для детерминированности + input_tensor = torch.randint(0, 1000, (1, 10)) + + # Жадный режим должен быть детерминированным + out1 = gpt.generate(input_tensor.clone(), max_new_tokens=5, do_sample=False) + out2 = gpt.generate(input_tensor.clone(), max_new_tokens=5, do_sample=False) + assert torch.equal(out1, out2), "Жадная генерация должна быть детерминированной" + + # Сэмплирующий режим с фиксированным seed + torch.manual_seed(42) + out3 = gpt.generate(input_tensor.clone(), max_new_tokens=5, do_sample=True) + torch.manual_seed(42) + out4 = gpt.generate(input_tensor.clone(), max_new_tokens=5, do_sample=True) + assert torch.equal(out3, out4), "Сэмплирование должно быть детерминированным при одинаковом seed" if __name__ == "__main__": pytest.main(["-v"]) \ No newline at end of file