Project structure:

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


Project files:

model.py
```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
        )

        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: (seq_length, batch_size)
        seq_length = src.size(0)
        embed = self.embedding(src) * torch.sqrt(torch.tensor(self.embed_size, dtype=torch.float32, device=src.device))
        embed = self.positional_encoding(embed)

        mask = self.generate_square_subsequent_mask(seq_length)
        output = self.transformer(src=embed, tgt=embed, src_mask=mask, tgt_mask=mask)
        logits = self.fc_out(output)
        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(1)  # Shape: (max_len, 1, embed_size)
        self.register_buffer('pe', pe)

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


predict.py
```py
import os

import torch

from MyLLM.train import seq_length
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']

# Загрузка последнего чекпоинта
import glob
checkpoint_list = glob.glob('checkpoints/model_epoch_1.pth')
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):
    model.eval()
    input_ids = [stoi.get(ch, stoi[' ']) for ch in start_text]
    input_ids = torch.tensor(input_ids, dtype=torch.long).unsqueeze(1).to(device)
    generated_text = start_text

    with torch.no_grad():
        for _ in range(max_length):
            input_seq = input_ids[-seq_length:] if len(input_ids) > seq_length else input_ids
            output = model(input_seq)
            logits = output[-1, 0, :] / temperature
            probs = torch.softmax(logits, dim=0)
            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=0)

    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
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.cuda import amp

# Фиксация сидов для воспроизводимости
torch.manual_seed(42)
np.random.seed(42)

# Параметры модели и обучения
embed_size = 256  # Увеличенный размер эмбеддингов
num_heads = 8
hidden_dim = 512  # Увеличена размерность скрытого слоя
num_layers = 4  # Уменьшено количество слоев для ускорения обучения
dropout = 0.1
batch_size = 256  # Увеличен размер батча для лучшей загрузки GPU
seq_length = 128  # Увеличена длина последовательности
num_epochs = 50
learning_rate = 1e-4

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


# Создаем словарь символов
def build_vocab(dataset_path):
    text = ""
    for filename in os.listdir(dataset_path):
        if filename.endswith(".txt"):
            with open(os.path.join(dataset_path, filename), 'r', encoding='utf-8') as f:
                text += f.read()
    chars = sorted(list(set(text)))
    stoi = {ch: i for i, ch in enumerate(chars)}
    itos = {i: ch for ch, i in stoi.items()}
    return stoi, itos, chars


stoi, itos, chars = build_vocab('dataset')
vocab_size = len(chars)
print(f"Размер словаря: {vocab_size}")


# Класс Dataset
class TextDataset(Dataset):
    def __init__(self, dataset_path, stoi, seq_length):
        data = ""
        for filename in os.listdir(dataset_path):
            if filename.endswith(".txt"):
                with open(os.path.join(dataset_path, filename), 'r', encoding='utf-8') as f:
                    data += f.read()
        self.data = [stoi[ch] for ch in data if ch in stoi]
        self.seq_length = seq_length

    def __len__(self):
        return len(self.data) - self.seq_length

    def __getitem__(self, idx):
        input_seq = torch.tensor(self.data[idx:idx + self.seq_length], dtype=torch.long)
        target_seq = torch.tensor(self.data[idx + 1:idx + self.seq_length + 1], dtype=torch.long)
        return input_seq, target_seq


# Создаем DataLoader с несколькими работниками для ускорения загрузки данных
dataset = TextDataset('dataset', stoi, seq_length)
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)

# Используем смешанную точность обучения
scaler = amp.GradScaler()


# Функция обучения
def train(model, dataloader, criterion, optimizer, num_epochs):
    model.train()
    for epoch in range(num_epochs):
        epoch_loss = 0
        for batch_idx, (input_seq, target_seq) in enumerate(dataloader):
            input_seq = input_seq.transpose(0, 1).to(device, non_blocking=True)
            target_seq = target_seq.transpose(0, 1).to(device, non_blocking=True)

            optimizer.zero_grad()
            with amp.autocast():
                output = model(input_seq)
                loss = criterion(output.view(-1, vocab_size), target_seq.reshape(-1))

            scaler.scale(loss).backward()
            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')


# Запускаем обучение
if __name__ == '__main__':
    train(model, dataloader, criterion, optimizer, num_epochs)
```


Это код моей LLM.
Размер датасета относительно большой и у меня не хватает ОЗУ. Что делать?
Вообще странно, зачем мне ОЗУ, не должна ли она записываться в память видеокарты?