Обновление метода generate в GPT

Основные изменения:
1. Добавлена поддержка различных стратегий генерации:
   - Жадный поиск (do_sample=False)
   - Вероятностное сэмплирование (do_sample=True)
   - Top-k сэмплирование (top_k параметр)
   - Nucleus (top-p) сэмплирование (top_p параметр)
   - Температурное сэмплирование (temperature параметр)

2. Добавлена валидация параметров:
   - Проверка temperature > 0
   - Проверка top_k > 0
   - Проверка top_p в диапазоне (0, 1]
   - Запрет одновременного использования top_k и top_p

3. Улучшена документация:
   - Подробное описание всех параметров
   - Примеры использования
   - Примечания о детерминированности
   - Описание исключений

4. Оптимизация кода:
   - Эффективное обрезание последовательности
   - Оптимизированные операции с тензорами
   - Четкое разделение логики для разных режимов

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)]()
+[![Python Version](https://img.shields.io/badge/python-3.8%2B-blue)]()
-[![PyTorch 2.0+](https://img.shields.io/badge/PyTorch-2.0+-red.svg)]()
+[![PyTorch Version](https://img.shields.io/badge/pytorch-1.10%2B-orange)]()
 
-## Основные компоненты
+Простая и понятная реализация языковой модели GPT-стиля с нуля на PyTorch
 
-### Токенизация
+## 🔍 Обзор
-- `SimpleBPE` - алгоритм Byte Pair Encoding
-- `OptimizeBPE` - оптимизированная версия
 
-### Эмбеддинги
+Simple-LLM предоставляет:
-- `TokenEmbeddings` - векторные представления токенов
+- Полную реализацию архитектуры GPT
-- `PositionalEmbeddings` - позиционное кодирование
+- Эффективный токенизатор BPE
+- Модули трансформера (внимание, FFN, эмбеддинги)
+- Гибкую систему генерации текста
+- Примеры использования и документацию
 
-### Transformer Layers
+## 🚀 Быстрый старт
-- `HeadAttention` - механизм внимания одной головы
-- `MultiHeadAttention` - многоголовое внимание (4-16 голов)
-- `FeedForward` - двухслойная FFN сеть (расширение → сжатие)
-- `Decoder` - полный декодер Transformer (Self-Attention + FFN)
 
-## Быстрый старт
+1. Установите зависимости:
-
-```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)
-
-## Примеры
 ```bash
-# Запуск примеров
+pip install torch numpy tqdm
-python -m example.multi_head_attention_example  # Визуализация внимания
-python -m example.feed_forward_example         # Анализ FFN слоя
 ```
 
-## Установка
+2. Запустите пример генерации:
 ```bash
-git clone https://github.com/pese-git/simple-llm.git
+python example/example_gpt.py
-cd simple-llm
-pip install -e .
 ```
 
-### Пример использования 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
+output = model.generate(
-
+    input_ids,
-bpe = SimpleBPE().fit(text_corpus)
+    max_new_tokens=100,
-tokens = bpe.encode("Текст для токенизации")
+    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)
+Полная документация доступна в [doc/](./doc/):
-- [Алгоритм BPE](/doc/bpe_algorithm.md)
+- [Архитектура GPT](./doc/gpt_documentation_ru.md)
-- [MultiHeadAttention](/doc/multi_head_attention_ru.md)
+- [Алгоритм BPE](./doc/bpe_algorithm.md)
-- [Decoder](/doc/decoder_ru.md)
+- [Примеры использования](./example/)
 
-## 🧪 Примеры
+## 🛠 Тестирование
 ```bash
-# Запуск примеров
+pytest tests/
-python -m example.example_gpt           # Генерация текста
-python -m example.multi_head_attention  # Визуализация внимания
-python -m example.decoder_example       # Демонстрация декодера
 ```
 
-## 🤝 Участие в разработке
+## 🤝 Как внести вклад
-PR и issues приветствуются! Перед внесением изменений:
+1. Форкните репозиторий
-1. Создайте issue с описанием
+2. Создайте ветку (`git checkout -b feature/AmazingFeature`)
-2. Сделайте fork репозитория
+3. Сделайте коммит (`git commit -m 'Add some AmazingFeature'`)
-3. Откройте Pull Request
+4. Запушьте ветку (`git push origin feature/AmazingFeature`)
+5. Откройте Pull Request
 
 ## 📜 Лицензия
-MIT License. Подробнее в [LICENSE](LICENSE).
+Распространяется под лицензией MIT. См. [LICENSE](./LICENSE)

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[Слой нормализации]
+    D --> E[Стек декодеров]
-    E --> F[Многоголовое внимание]
+    E --> F[Многоголовое самовнимание]
-    F --> G[Пропускная связь]
+    F --> G[Add & Norm]
-    G --> H[FeedForward слой]
+    G --> H[FeedForward]
-    H --> I[Слой нормализации]
+    H --> I[Add & Norm]
     I --> J[Выходные логиты]
+    J --> K[Выбор следующего токена]
+    K --> L[Добавление к последовательности]
 ```
 
-### 2.2 Процесс генерации
+### 2.2 Режимы генерации
-1. Токенизация входного текста
+1. **Жадный поиск** (do_sample=False):
-2. Вычисление эмбеддингов:
+   - Всегда выбирает токен с максимальной вероятностью
-   - Токенные + позиционные
+   - Детерминированный результат
-3. Прохождение через N декодеров:
-   - Самовнимание с маской
-   - Полносвязные слои
-4. Преобразование в вероятности
-5. Выбор следующего токена
 
-## 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,
+    vocab_size=10000,    # Размер словаря
-    max_seq_len=512,
+    max_seq_len=512,     # Макс. длина контекста
-    emb_size=768,
+    emb_size=768,        # Размерность эмбеддингов
-    num_heads=12,
+    num_heads=12,        # Число голов внимания
-    head_size=64,
+    head_size=64,        # Размерность головы
-    num_layers=6
+    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 Примеры использования
 
-| Параметр       | Описание                          |
+**Базовый пример:**
-|----------------|-----------------------------------|
+```python
-| vocab_size     | Размер словаря                   |
+text = "Анализ данных - это"
-| max_seq_len    | Макс. длина последовательности   |
+input_ids = tokenizer.encode(text)
-| emb_size       | Размерность эмбеддингов          |
+output_ids = model.generate(input_ids, max_new_tokens=100)
-| num_heads      | Количество голов внимания        |
+generated_text = tokenizer.decode(output_ids[0])
-| head_size      | Размерность головы внимания      |
+```
-| num_layers     | Количество слоев декодера        |
 
-## 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)

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,  # Размер словаря
+        'vocab_size': 10000,
-        'max_seq_len': 256,   # Макс. длина последовательности
+        'max_seq_len': 256,
-        'emb_size': 512,      # Размерность эмбеддингов
+        'emb_size': 512,
-        'num_heads': 8,       # Количество голов внимания
+        'num_heads': 8,
-        'head_size': 64,      # Размер каждой головы внимания
+        'head_size': 64,
-        'num_layers': 6,      # Количество слоев декодера
+        'num_layers': 6,
-        'dropout': 0.1,       # Dropout
+        '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. Пример генерации с токенизатором
+    # 2. Пример генерации
-    try:
+    print("\nИспользуется числовая генерация...")
-        from simple_llm.tokenizer.simple_bpe import SimpleBPE
+    use_numeric_generation(config, model)
-        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]}")
 
     # 3. Сохранение и загрузка модели
     print("\nТест сохранения/загрузки...")
@@ -63,8 +84,7 @@ def main():
         loaded_model = GPT.load(tmp.name, device=config['device'])
         print("Модель успешно загружена")
         
-        # Проверка работы загруженной модели
+        test_output = loaded_model(torch.randint(0, config['vocab_size'], (1, 5)).to(config['device']))
-        test_output = loaded_model(input_seq)
         print(f"Тест загруженной модели - выходная форма: {test_output.shape}")
 
 if __name__ == "__main__":

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]
+            x (torch.Tensor): Входной тензор с индексами токенов формы [batch_size, seq_len],
-            max_new_tokens: Максимальное количество новых токенов для генерации
+                              где 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]
+            torch.Tensor: Тензор с расширенной последовательностью токенов формы 
+                          [batch_size, seq_len + max_new_tokens]
 
-        Алгоритм работы:
+        Raises:
-        1. На каждом шаге берется последний фрагмент последовательности (не длиннее max_seq_len)
+            ValueError: Если входная последовательность длиннее max_seq_len
-        2. Вычисляются логиты для следующего токена
+            ValueError: Если temperature <= 0
-        3. Выбирается токен с максимальной вероятностью (жадный алгоритм)
+
-        4. Токен добавляется к последовательности
+        Examples:
-        5. Процесс повторяется пока не сгенерируется max_new_tokens токенов
+            >>> # Жадная генерация
+            >>> 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,10 +193,54 @@ class GPT(nn.Module):
             # 3. Берем логиты для последнего токена
             last_logits = logits[:, -1, :]  # [batch_size, vocab_size]
 
-            # 4. Применяем Softmax
+            # Масштабируем логиты температурой
-            probs = F.softmax(last_logits, dim=-1)  # [batch_size, vocab_size]
+            if temperature > 0:
+                logits_scaled = last_logits / temperature
+            else:
+                logits_scaled = last_logits
 
-            # 5. Выбираем токен с максимальной вероятностью
+            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(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]
             
             # 6. Добавляем его к последовательности

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_with_sampling(self, default_config, sample_input):
-    def test_generate_deterministic(self, default_config):
+        """Тест генерации с сэмплированием"""
-        """Тест детерминированности генерации (при одинаковом seed)"""
+        torch.manual_seed(42)
-        # Фиксируем seed для входа
+        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
+        greedy = gpt.generate(input_tensor, max_new_tokens=5, do_sample=False)
-        out1 = gpt.generate(input_tensor.clone(), max_new_tokens=5)
+        sampled = gpt.generate(input_tensor, max_new_tokens=5, do_sample=True)
-        out2 = gpt.generate(input_tensor.clone(), max_new_tokens=5)
         
-        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"])