Project structure:

MyLLM/
├── model.py
├── predict.py
├── process_dataset.py
└── train.py


Project files:

model.py
```py
# model.py
import torch
import torch.nn as nn


class CharTransformerModel(nn.Module):
    def __init__(self, vocab_size, embed_size, num_heads, hidden_dim, num_layers, dropout=0.1):
        super(CharTransformerModel, self).__init__()
        self.embed_size = embed_size
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.positional_encoding = PositionalEncoding(embed_size, dropout)

        self.transformer = nn.Transformer(
            d_model=embed_size,
            nhead=num_heads,
            num_encoder_layers=num_layers,
            num_decoder_layers=num_layers,
            dim_feedforward=hidden_dim,
            dropout=dropout,
            batch_first=True  # Устанавливаем batch_first=True
        )

        self.fc_out = nn.Linear(embed_size, vocab_size)
        self.src_mask = None

    def generate_square_subsequent_mask(self, sz):
        if self.src_mask is None or self.src_mask.size(0) != sz:
            mask = torch.triu(torch.ones(sz, sz), 1)
            self.src_mask = mask.masked_fill(mask == 1, float('-inf')).to(next(self.parameters()).device)
        return self.src_mask

    def forward(self, src):
        """
        src shape: (batch_size, seq_length)
        """
        seq_length = src.size(1)
        embed = self.embedding(src) * torch.sqrt(torch.tensor(self.embed_size, dtype=torch.float32, device=src.device))
        embed = self.positional_encoding(embed)  # Shape: (batch_size, seq_length, embed_size)

        mask = self.generate_square_subsequent_mask(seq_length)
        # Transformer expects src and tgt; для генерации используем src=embed и tgt=embed
        output = self.transformer(src=embed, tgt=embed, src_mask=mask, tgt_mask=mask)
        logits = self.fc_out(output)  # Shape: (batch_size, seq_length, vocab_size)
        return logits


class PositionalEncoding(nn.Module):
    def __init__(self, embed_size, dropout=0.1, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, embed_size, device='cpu')
        position = torch.arange(0, max_len, dtype=torch.float32, device='cpu').unsqueeze(1)
        div_term = torch.exp(torch.arange(0, embed_size, 2, dtype=torch.float32, device='cpu') * (
                    -torch.log(torch.tensor(10000.0)) / embed_size))
        pe[:, 0::2] = torch.sin(position * div_term)
        if embed_size % 2 == 1:
            pe[:, 1::2] = torch.cos(position * div_term[:-1])
        else:
            pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)  # Shape: (1, max_len, embed_size)
        self.register_buffer('pe', pe)

    def forward(self, x):
        # x shape: (batch_size, seq_length, embed_size)
        x = x + self.pe[:, :x.size(1)]
        return self.dropout(x)
```


predict.py
```py
# predict.py
import os
import glob
import torch
from model import CharTransformerModel

# Определение устройства
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Функция загрузки модели
def load_checkpoint(checkpoint_path):
    checkpoint = torch.load(checkpoint_path, map_location=device)
    model_params = {
        'vocab_size': len(checkpoint['stoi']),
        'embed_size': checkpoint['embed_size'],
        'num_heads': checkpoint['num_heads'],
        'hidden_dim': checkpoint['hidden_dim'],
        'num_layers': checkpoint['num_layers'],
        'dropout': checkpoint['dropout']
    }
    model = CharTransformerModel(**model_params).to(device)
    model.load_state_dict(checkpoint['model_state_dict'])
    model.eval()
    return model, checkpoint['stoi'], checkpoint['itos']

# Загрузка последнего чекпоинта
checkpoint_list = glob.glob('checkpoints/model_epoch_*.pth')
if not checkpoint_list:
    raise FileNotFoundError("Чекпоинты не найдены в директории 'checkpoints'.")
latest_checkpoint = max(checkpoint_list, key=os.path.getctime)
model, stoi, itos = load_checkpoint(latest_checkpoint)
vocab_size = len(stoi)

# Генерация текста
def generate_text(model, start_text, stoi, itos, max_length=200, temperature=1.0, seq_length=128):
    model.eval()
    input_ids = [stoi.get(ch, stoi[' ']) for ch in start_text]
    input_ids = torch.tensor(input_ids, dtype=torch.long).unsqueeze(0).to(device)  # Shape: (1, seq_length)
    generated_text = start_text

    with torch.no_grad():
        for _ in range(max_length):
            input_seq = input_ids[:, -seq_length:] if input_ids.size(1) > seq_length else input_ids
            output = model(input_seq)  # Shape: (1, seq_length, vocab_size)
            logits = output[:, -1, :] / temperature  # Берем последний шаг для каждой последовательности
            probs = torch.softmax(logits, dim=-1)
            next_char_id = torch.multinomial(probs, num_samples=1).item()
            next_char = itos[next_char_id]
            generated_text += next_char

            # Добавляем следующий символ в последовательность
            input_ids = torch.cat([input_ids, torch.tensor([[next_char_id]], device=device)], dim=1)

    return generated_text

# Пример использования
if __name__ == '__main__':
    start_text = input("Введите начальный текст: ")
    generated_text = generate_text(model, start_text, stoi, itos, max_length=500, temperature=0.8)
    print("Сгенерированный текст:")
    print(generated_text)
```


train.py
```py
# train.py
import os
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from model import CharTransformerModel
import numpy as np
from torch import amp
import random


# Класс Dataset, который читает данные по мере необходимости
class TextDataset(Dataset):
    def __init__(self, dataset_path, stoi, seq_length):
        self.file_paths = [os.path.join(dataset_path, filename)
                           for filename in os.listdir(dataset_path) if filename.endswith(".txt")]
        self.stoi = stoi
        self.seq_length = seq_length

        # Подсчет общего количества последовательностей во всех файлах
        self.total_sequences = 0
        self.file_sequences = []
        for file_path in self.file_paths:
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    f.seek(0, os.SEEK_END)
                    file_size = f.tell()
            except UnicodeDecodeError:
                print(f"Файл {file_path} содержит некорректные символы и будет пропущен.")
                self.file_sequences.append(0)
                continue
            num_sequences = max(0, file_size - self.seq_length - 1)
            self.file_sequences.append(num_sequences)
            self.total_sequences += num_sequences

    def __len__(self):
        return self.total_sequences

    def __getitem__(self, idx):
        # Определение, в какой файл попадает текущий индекс
        cumulative_sequences = 0
        for file_idx, num_sequences in enumerate(self.file_sequences):
            if cumulative_sequences + num_sequences > idx:
                break
            cumulative_sequences += num_sequences

        # Преобразование индекса в локальный индекс внутри файла
        local_idx = idx - cumulative_sequences
        file_path = self.file_paths[file_idx]

        with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
            # Переход к стартовой позиции
            f.seek(local_idx)
            data = f.read(self.seq_length + 1)
            # Если данных недостаточно, дополнить пробелами
            if len(data) < self.seq_length + 1:
                data += ' ' * (self.seq_length + 1 - len(data))
            data_indices = [self.stoi.get(ch, self.stoi[' ']) for ch in data]
            input_seq = torch.tensor(data_indices[:-1], dtype=torch.long)
            target_seq = torch.tensor(data_indices[1:], dtype=torch.long)
            return input_seq, target_seq


def build_vocab(dataset_path):
    chars = set()
    for filename in os.listdir(dataset_path):
        if filename.endswith(".txt"):
            file_path = os.path.join(dataset_path, filename)
            with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                while True:
                    chunk = f.read(1024 * 1024)  # Чтение чанками по 1МБ
                    if not chunk:
                        break
                    chars.update(chunk)
    chars = sorted(chars)
    stoi = {ch: i for i, ch in enumerate(chars)}
    itos = {i: ch for i, ch in enumerate(chars)}
    return stoi, itos, chars


def train(model, dataloader, criterion, optimizer, num_epochs, device, stoi, itos, embed_size, num_heads, hidden_dim, num_layers, dropout):
    model.train()
    scaler = amp.GradScaler()  # Исправлено согласно FutureWarning

    for epoch in range(num_epochs):
        epoch_loss = 0
        for batch_idx, (input_seq, target_seq) in enumerate(dataloader):
            # Перестановка размерностей: (batch_size, seq_length) уже соответствует batch_first=True
            input_seq = input_seq.to(device, non_blocking=True)
            target_seq = target_seq.to(device, non_blocking=True).long()  # Убедимся, что тип Long

            optimizer.zero_grad()
            str_device = torch.device("cuda" if torch.cuda.is_available() else "cpu").__str__()
            with amp.autocast(device_type=str_device):
                output = model(input_seq)  # output: (batch_size, seq_length, vocab_size)
                # Изменение формы для CrossEntropyLoss: (batch_size * seq_length, vocab_size)
                output = output.view(-1, output.size(-1))
                # Цели: (batch_size * seq_length)
                target = target_seq.view(-1)

                loss = criterion(output, target)

            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)  # Для градиентного клиппинга
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # Градиентный клиппинг
            scaler.step(optimizer)
            scaler.update()
            epoch_loss += loss.item()

            if batch_idx % 50 == 0:
                print(
                    f"Эпоха [{epoch + 1}/{num_epochs}], Батч [{batch_idx}/{len(dataloader)}], Потеря: {loss.item():.4f}")

        avg_loss = epoch_loss / len(dataloader)
        print(f"Эпоха [{epoch + 1}/{num_epochs}], Средняя потеря: {avg_loss:.4f}")

        # Сохранение чекпоинта после каждой эпохи
        checkpoint_path = f'checkpoints/model_epoch_{epoch + 1}.pth'
        os.makedirs('checkpoints', exist_ok=True)
        torch.save({
            'epoch': epoch + 1,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'scaler_state_dict': scaler.state_dict(),
            'loss': avg_loss,
            'stoi': stoi,
            'itos': itos,
            'embed_size': embed_size,
            'num_heads': num_heads,
            'hidden_dim': hidden_dim,
            'num_layers': num_layers,
            'dropout': dropout
        }, checkpoint_path)
        print(f"Чекпоинт сохранен: {checkpoint_path}")

    # Сохранение финальной модели
    torch.save(model.state_dict(), 'model_final.pth')
    print("Финальная модель сохранена: model_final.pth")


def main():
    # Фиксация случайных зерен для воспроизводимости
    torch.manual_seed(42)
    np.random.seed(42)
    random.seed(42)

    # Параметры модели и обучения
    embed_size = 256
    num_heads = 8
    hidden_dim = 512
    num_layers = 4
    dropout = 0.1
    batch_size = 256
    seq_length = 128
    num_epochs = 50
    learning_rate = 1e-4

    # Определение устройства
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Используется устройство: {device}")

    # Построение словаря символов без загрузки всего датасета в оперативную память
    stoi, itos, chars = build_vocab('dataset')
    vocab_size = len(chars)
    print(f"Размер словаря: {vocab_size}")

    if vocab_size == 0:
        raise ValueError("Словарь пуст. Проверьте содержимое директории 'dataset'.")

    # Создание DataLoader
    dataset = TextDataset('dataset', stoi, seq_length)
    if len(dataset) == 0:
        raise ValueError("Датасет не содержит последовательностей для обучения.")
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True,
                            drop_last=True, pin_memory=True, num_workers=4)

    # Инициализация модели
    model = CharTransformerModel(vocab_size, embed_size, num_heads, hidden_dim, num_layers, dropout).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

    # Запуск обучения
    train(model, dataloader, criterion, optimizer, num_epochs, device,
          stoi, itos, embed_size, num_heads, hidden_dim, num_layers, dropout)


if __name__ == '__main__':
    main()
```


Слушай, ну он ОЧЕНЬ МЕДЛЕННО обучается. Вот что получилось за 5 минут:
```
Эпоха [1/50], Батч [0/819631], Потеря: 7.2273
Эпоха [1/50], Батч [50/819631], Потеря: 3.9903
Эпоха [1/50], Батч [100/819631], Потеря: 3.0844
Эпоха [1/50], Батч [150/819631], Потеря: 2.8335
```
Я подсчитал, получается, одна эпоха будет обучаться 11 дней. Это никуда не годится. Что не так?