import logging
import torch
from transformers import (
GPT2LMHeadModel,
GPT2TokenizerFast,
TextDataset,
DataCollatorForLanguageModeling,
Trainer,
TrainingArguments,
AutoConfig,
)
from datasets import load_dataset
import os
Настройка логирования
logging.basicConfig(level=logging.INFO)
Выбор устройства
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info(f"Используется устройство: {device}")
Определение имени модели (выберите подходящую модель!)
MODEL_NAME = "sberbank-ai/rugpt3large_based_on_gpt2" # или другой вариант
Путь к файлу с данными для обучения
TRAIN_FILE = "data.txt" #Замените на свой
EVAL_FILE = "data.txt" #Замените на свой или уберите, чтобы не делать валидацию.
Создание директории для сохранения модели
OUTPUT_DIR = "my_gpt2_model"
os.makedirs(OUTPUT_DIR, exist_ok=True)
1. Загрузка токенизатора и модели
tokenizer = GPT2TokenizerFast.from_pretrained(MODEL_NAME)
tokenizer.pad_token = tokenizer.eos_token # Важно!
model = GPT2LMHeadModel.from_pretrained(MODEL_NAME).to(device)
2. Подготовка данных
dataset = TextDataset(
tokenizer=tokenizer,
file_path=TRAIN_FILE,
block_size=128 # Размер блока текста
)
dataset = load_dataset("text", data_files={"train": TRAIN_FILE,"test":EVAL_FILE})
def tokenize_function(examples):
return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=128) #truncation=True and padding="max_length"
tokenized_datasets = dataset.map(tokenize_function, batched=True, num_proc=4, remove_columns=["text"])
train_dataset = tokenized_datasets["train"]
eval_dataset = tokenized_datasets["test"]
3. Data Collator
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False) # mlm=False важно для GPT-2
4. Параметры обучения
training_args = TrainingArguments(
output_dir=OUTPUT_DIR,
overwrite_output_dir=True,
num_train_epochs=3, # Количество эпох обучения
per_device_train_batch_size=4, # Размер батча
save_steps=1000, # Как часто сохранять модель
save_total_limit=2, # Максимальное количество сохраняемых моделей
prediction_loss_only=True, # Вычислять только loss
evaluation_strategy="steps",
eval_steps=500,
learning_rate = 5e-5,
weight_decay = 0.01,
push_to_hub=False,
)
5. Trainer
trainer = Trainer(
model=model,
args=training_args,
data_collator=data_collator,
train_dataset=train_dataset,
eval_dataset=eval_dataset
)
6. Обучение
trainer.train()
7. Сохранение модели
trainer.save_model(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)
logging.info(f"Модель сохранена в {OUTPUT_DIR}")
#Загрузка и пример использования
config = AutoConfig.from_pretrained(OUTPUT_DIR)
tokenizer = GPT2TokenizerFast.from_pretrained(OUTPUT_DIR)
model = GPT2LMHeadModel.from_pretrained(OUTPUT_DIR).to(device)
prompt = "Привет, мир! Как дела?"
input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
with torch.no_grad():
output = model.generate(input_ids, max_length=200, num_return_sequences=1)
generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
logging.info(f"Сгенерированный текст: {generated_text}")
