import torch import pytest from simple_llm.transformer.gpt import GPT class TestGPT: @pytest.fixture def default_config(self): return { 'vocab_size': 1000, 'max_seq_len': 128, 'emb_size': 256, 'num_heads': 4, 'head_size': 64, 'num_layers': 2, 'dropout': 0.1 } @pytest.fixture def sample_input(self): return torch.randint(0, 1000, (2, 32)) # batch_size=2, seq_len=32 def test_initialization(self, default_config): """Проверка создания модели""" gpt = GPT(**default_config) assert isinstance(gpt, torch.nn.Module) assert len(gpt._decoders) == default_config['num_layers'] def test_forward_pass(self, default_config, sample_input): """Тест прямого прохода""" gpt = GPT(**default_config) output = gpt(sample_input) assert output.shape == (2, 32, 1000) # batch, seq_len, vocab_size def test_max_length(self, default_config): """Проверка обработки максимальной длины""" gpt = GPT(**default_config) # Корректная длина x = torch.randint(0, 1000, (1, 128)) output = gpt(x) # Слишком длинная последовательность with pytest.raises(ValueError): x = torch.randint(0, 1000, (1, 129)) gpt(x) def test_generate_basic(self, default_config, sample_input): """Тест базовой генерации""" gpt = GPT(**default_config) 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): """Тест генерации с пустым входом""" gpt = GPT(**default_config) empty_input = torch.randint(0, 1000, (2, 0)) with pytest.raises(IndexError): 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, do_sample=False) assert generated.shape == (2, 129) 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)) 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 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"])