import json
import torch
from transformers import (
GPT2LMHeadModel,
GPT2Tokenizer,
Trainer,
TrainingArguments,
default_data_collator,
EarlyStoppingCallback,
TrainerCallback
)
from torch.utils.data import Dataset, random_split
from tqdm import tqdm

Конфигурация
MODEL_NAME = "gpt2"
TRAIN_FILE = 'bd.json'
MAX_LENGTH = 512
NUM_EPOCHS = 200
BATCH_SIZE = 8
GRADIENT_ACCUMULATION_STEPS = 2
LEARNING_RATE = 5e-5
WARMUP_RATIO = 0.1
EARLY_STOPPING_PATIENCE = 2

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

Инициализация модели и токенизатора
tokenizer = GPT2Tokenizer.from_pretrained(MODEL_NAME)
model = GPT2LMHeadModel.from_pretrained(MODEL_NAME)

Добавление специальных токенов (включая токены для цепочки рассуждений)
special_tokens = {
'pad_token': '<pad>',
'additional_special_tokens': [
'<question>', '</question>',
'<chain_of_thought>', '</chain_of_thought>',
'<answer>', '</answer>'
]
}
tokenizer.add_special_tokens(special_tokens)
model.resize_token_embeddings(len(tokenizer))
model.to(device)

def load_data(file_path):
with open(file_path, 'r', encoding='utf-8') as f:
data = json.load(f)
# Если в исходных данных нет цепочки рассуждений, формируем строки в формате:
# <question>... </question><answer>... </answer>{eos_token}
return [
f"<question>{item['question'].strip()}</question><answer>{item['answer'].strip()}</answer>{tokenizer.eos_token}"
for item in tqdm(data)
]

class QADataset(Dataset):
def init(self, texts, tokenizer, max_length=MAX_LENGTH):
self.encodings = tokenizer(
texts,
max_length=max_length,
truncation=True,
padding=True,
return_tensors="pt"
)

text
    answer_start_id = tokenizer.convert_tokens_to_ids('<answer>')
    answer_end_id = tokenizer.convert_tokens_to_ids('</answer>')
    eos_id = tokenizer.eos_token_id

    labels = []
    for i in range(len(self.encodings['input_ids'])):
        input_ids = self.encodings['input_ids'][i]
        label = input_ids.clone()

        start_idx = (input_ids == answer_start_id).nonzero()
        end_idx = (input_ids == answer_end_id).nonzero()

        if start_idx.size(0) > 0 and end_idx.size(0) > 0:
            start = start_idx[0] + 1
            end = end_idx[0]
            label[:start] = -100
            label[end:] = -100
            if input_ids[-1] == eos_id:
                label[-1] = eos_id
        else:
            label[:] = -100

        labels.append(label)

    self.encodings['labels'] = torch.stack(labels)

def __len__(self):
    return len(self.encodings['input_ids'])

def __getitem__(self, idx):
    return {key: val[idx] for key, val in self.encodings.items()}
Подготовка данных
all_data = load_data(TRAIN_FILE)
train_size = int(0.9 * len(all_data))
val_size = len(all_data) - train_size
train_data, val_data = random_split(all_data, [train_size, val_size])

Аргументы обучения
training_args = TrainingArguments(
output_dir="./results",
num_train_epochs=NUM_EPOCHS,
per_device_train_batch_size=BATCH_SIZE,
per_device_eval_batch_size=BATCH_SIZE,
gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS,
learning_rate=LEARNING_RATE,
warmup_ratio=WARMUP_RATIO,
evaluation_strategy="epoch",
save_strategy="epoch",
load_best_model_at_end=True,
metric_for_best_model="eval_loss",
greater_is_better=False,
fp16=torch.cuda.is_available(),
dataloader_pin_memory=False,
fp16_full_eval=False,
logging_steps=50,
weight_decay=0.01,
lr_scheduler_type="cosine",
gradient_checkpointing=True
)

Функция генерации ответа (при chain=True дополнительно генерируются рассуждения)
def generate_response(question, max_new_tokens=200, chain=False):
if chain:
# Запрос с цепочкой рассуждений: сначала просим описать рассуждения, потом дать ответ
prompt = (f"<question>{question.strip()}</question>"
f"<chain_of_thought>Пожалуйста, опиши свои размышления:</chain_of_thought>"
f"<answer>")
else:
prompt = f"<question>{question.strip()}</question><answer>"

text
inputs = tokenizer(prompt, return_tensors="pt", max_length=MAX_LENGTH - 50, truncation=True).to(device)

outputs = model.generate(
    **inputs,
    max_new_tokens=max_new_tokens,
    temperature=0.7,
    top_k=40,
    top_p=0.9,
    repetition_penalty=1.2,
    no_repeat_ngram_size=3,
    pad_token_id=tokenizer.pad_token_id,
    eos_token_id=tokenizer.eos_token_id,
    do_sample=True,
    num_beams=3,
    early_stopping=True
)

decoded = tokenizer.decode(outputs[0], skip_special_tokens=False)

if chain:
    # Извлекаем рассуждения
    cot_start = decoded.find("<chain_of_thought>")
    cot_end = decoded.find("</chain_of_thought>")
    if cot_start != -1 and cot_end != -1:
        chain_text = decoded[cot_start + len("<chain_of_thought>"):cot_end].strip()
    else:
        chain_text = "[Не удалось извлечь рассуждения]"
else:
    chain_text = None

# Извлекаем ответ
# Ищем <answer> и </answer>
ans_start = decoded.find("<answer>")
ans_end = decoded.find("</answer>")
if ans_start != -1:
    answer = decoded[ans_start + len("<answer>"): ans_end if ans_end != -1 else None]
else:
    answer = decoded
answer = answer.replace(tokenizer.pad_token, "").replace(tokenizer.eos_token, "").strip()

if chain:
    return {"chain_of_thought": chain_text, "answer": answer}
else:
    return answer
Кастомный Callback, который каждые 5 эпох выводит тестовый пример (вопрос, цепочку рассуждений и ответ)
class ChainOfThoughtCallback(TrainerCallback):
def on_epoch_end(self, args, state, control, **kwargs):
# state.epoch может быть float, поэтому округляем до целого числа
current_epoch = int(state.epoch)
if current_epoch > 0 and current_epoch % 5 == 0:
test_question = "Как, по-твоему, будет развиваться искусственный интеллект в ближайшие 10 лет?"
output = generate_response(test_question, chain=True)
print("\n=== Эпоха {} ===".format(current_epoch))
print("Вопрос: {}".format(test_question))
print("Размышления: {}".format(output["chain_of_thought"]))
print("Ответ: {}".format(output["answer"]))
print("====================\n")
return control

Инициализация Trainer
trainer = Trainer(
model=model,
args=training_args,
train_dataset=QADataset([all_data[i] for i in train_data.indices], tokenizer),
eval_dataset=QADataset([all_data[i] for i in val_data.indices], tokenizer),
data_collator=default_data_collator,
callbacks=[EarlyStoppingCallback(early_stopping_patience=EARLY_STOPPING_PATIENCE),
ChainOfThoughtCallback()]
)

Обучение
print("Начало обучения...")
trainer.train()
print("Обучение завершено!")

Сохранение модели
model.save_pretrained("./trained_model")
tokenizer.save_pretrained("./trained_model")

Интерактивный режим для генерации ответов
if name == "main":
print("Введите вопрос (или 'exit' для выхода):")
while True:
question = input("> ")
if question.lower() in ["exit", "quit"]:
break
# Здесь в интерактивном режиме можно вызвать генерацию без цепочки рассуждений:
print(f"Ответ: {generate_response(question)}\n")