| import torch
|
| import torch.nn as nn
|
| import torch.optim as optim
|
| import psutil
|
| from tqdm import tqdm
|
| import time
|
|
|
|
|
| def load_data(file_path):
|
| with open(file_path, 'r', encoding='utf-8') as f:
|
| return f.read()
|
|
|
|
|
| class SimpleTokenizer:
|
| def __init__(self, vocab_path):
|
| self.char_to_idx = torch.load(vocab_path)
|
| self.idx_to_char = {i: c for c, i in self.char_to_idx.items()}
|
|
|
| def encode(self, text):
|
| return [self.char_to_idx.get(c, self.char_to_idx.get('<unk>', 0)) for c in text]
|
|
|
| def decode(self, indices):
|
| return ''.join([self.idx_to_char.get(i, '') for i in indices])
|
|
|
|
|
| class TransformerModel(nn.Module):
|
| def __init__(self, vocab_size, emb_size=256, num_heads=4, num_layers=4, ff_hid_dim=1024):
|
| super().__init__()
|
| self.embedding = nn.Embedding(vocab_size, emb_size)
|
| self.pos_embedding = nn.Parameter(torch.zeros(1, 512, emb_size))
|
| self.transformer_blocks = nn.ModuleList([
|
| nn.TransformerEncoderLayer(d_model=emb_size, nhead=num_heads, dim_feedforward=ff_hid_dim)
|
| for _ in range(num_layers)
|
| ])
|
| self.output = nn.Linear(emb_size, vocab_size)
|
|
|
| def forward(self, x):
|
| x = self.embedding(x) + self.pos_embedding[:, :x.size(1), :]
|
| for block in self.transformer_blocks:
|
| x = block(x)
|
| return self.output(x)
|
|
|
|
|
| def get_batches(data, batch_size, seq_length):
|
| inputs, targets = [], []
|
| for i in range(0, len(data) - seq_length - 1, seq_length):
|
| x = data[i:i + seq_length]
|
| y = data[i + 1:i + 1 + seq_length]
|
| if len(x) == seq_length and len(y) == seq_length:
|
| inputs.append(x)
|
| targets.append(y)
|
| if len(inputs) == batch_size:
|
| yield (
|
| torch.tensor(inputs, dtype=torch.long),
|
| torch.tensor(targets, dtype=torch.long)
|
| )
|
| inputs, targets = [], []
|
|
|
|
|
| def show_memory():
|
| process = psutil.Process()
|
| mem_info = process.memory_info()
|
| return f"{mem_info.rss / 1024**2:.2f} MB"
|
|
|
|
|
| def train():
|
| vocab_size = 30000
|
| batch_size = 64
|
| seq_length = 64
|
| num_epochs = 3
|
| lr = 0.001
|
| vocab_path = 'vocab.pth'
|
| data_path = 'data.txt'
|
|
|
| text = load_data(data_path)
|
| tokenizer = SimpleTokenizer(vocab_path)
|
| tokens = tokenizer.encode(text)
|
| model = TransformerModel(vocab_size)
|
| optimizer = optim.Adam(model.parameters(), lr=lr)
|
| criterion = nn.CrossEntropyLoss()
|
|
|
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
| model.to(device)
|
|
|
| for epoch in range(num_epochs):
|
| batches = list(get_batches(tokens, batch_size, seq_length))
|
| total = len(batches)
|
| total_loss = 0
|
| print(f"\n🧠 Epoch {epoch+1}/{num_epochs} — {total} batches")
|
|
|
| with tqdm(total=total, desc="Training", bar_format="{l_bar}{bar} [ time left: {remaining} ]") as pbar:
|
| for step, (x, y) in enumerate(batches):
|
| x, y = x.to(device), y.to(device)
|
| optimizer.zero_grad()
|
| output = model(x)
|
| loss = criterion(output.view(-1, vocab_size), y.view(-1))
|
| loss.backward()
|
| optimizer.step()
|
|
|
| total_loss += loss.item()
|
| avg_loss = total_loss / (step + 1)
|
|
|
| if step % 10 == 0:
|
| pbar.set_description(f"Loss: {loss.item():.4f} | RAM: {show_memory()}")
|
| pbar.update(1)
|
|
|
| torch.save(model.state_dict(), f"model_epoch_{epoch+1}.pth")
|
| print(f"💾 Model saved: model_epoch_{epoch+1}.pth")
|
|
|
| if __name__ == "__main__":
|
| train()
|
|
|
