File size: 2,170 Bytes
cf6354c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 | 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") |