import torch
from transformers import T5Config, T5ForConditionalGeneration

def create_t5_3b_model(vocab_size: int = 32128) -> T5ForConditionalGeneration:
    """
    Создаёт encoder-decoder модель T5 с ~3 миллиардами параметров.
    Параметры взяты из официальной спецификации T5-3B (v1.1).
    """
    config = T5Config(
        vocab_size=vocab_size,          # Размер словаря (как в оригинальном T5)
        d_model=1024,                    # Размерность скрытого слоя
        d_kv=128,                         # Размерность ключей/значений в attention
        num_heads=32,                      # Количество голов внимания
        num_layers=24,                      # Количество слоёв в encoder и decoder
        d_ff=16384,                         # Размерность feed‑forward слоя
        dropout_rate=0.1,                   # Вероятность dropout
        feed_forward_proj="gated-gelu",      # Тип проекции (T5 v1.1 использует gated‑GELU)
        tie_word_embeddings=False,           # Не связываем входные и выходные эмбеддинги (обычно False для T5)
        # Дополнительные параметры можно оставить по умолчанию
    )
    model = T5ForConditionalGeneration(config)
    return model

if __name__ == "__main__":
    # Создаём модель
    model = create_t5_3b_model()
    
    # Выводим число параметров
    num_params = model.num_parameters()
    print(f"Модель содержит {num_params:,} параметров.")
    
    # Сохраняем конфигурацию (и веса, если нужно)
    model.config.save_pretrained("./t5-3b-config")
    # model.save_pretrained("./t5-3b-random")  # раскомментировать для сохранения весов (около 12 ГБ)
    print("Конфигурация сохранена в папку ./t5-3b-config")