ysn-rfd/text-dataset-tiny-code-script-py-format

1import torch
2import torch.nn as nn
3import torch.optim as optim
4import numpy as np
5import random
6# خواندن داده‌ها از فایل
7with open("data.txt", "r", encoding="utf-8") as f:
8    text = f.read()
9
10# ایجاد دیکشنری برای تبدیل کاراکترها به اندیس و برعکس
11chars = sorted(list(set(text)))
12char_to_idx = {ch: i for i, ch in enumerate(chars)}
13idx_to_char = {i: ch for i, ch in enumerate(chars)}
14
15# تبدیل متن به لیست از اندیس‌ها
16data = [char_to_idx[ch] for ch in text]
17
18# تنظیم پارامترهای آموزشی
19seq_length = 50  # طول دنباله ورودی
20batch_size = 64
21hidden_size = 128
22num_layers = 2
23num_epochs = 100
24learning_rate = 0.01
25class TextDataset(torch.utils.data.Dataset):
26    def __init__(self, data, seq_length):
27        self.data = data
28        self.seq_length = seq_length
29    
30    def __len__(self):
31        return len(self.data) - self.seq_length
32    
33    def __getitem__(self, idx):
34        return (
35            torch.tensor(self.data[idx:idx+self.seq_length], dtype=torch.long),
36            torch.tensor(self.data[idx+1:idx+self.seq_length+1], dtype=torch.long)
37        )
38
39dataset = TextDataset(data, seq_length)
40dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)
41class LSTMModel(nn.Module):
42    def __init__(self, vocab_size, hidden_size, num_layers):
43        super(LSTMModel, self).__init__()
44        self.embedding = nn.Embedding(vocab_size, hidden_size)
45        self.lstm = nn.LSTM(hidden_size, hidden_size, num_layers, batch_first=True)
46        self.fc = nn.Linear(hidden_size, vocab_size)
47
48    def forward(self, x, hidden=None):
49        x = self.embedding(x)
50        output, hidden = self.lstm(x, hidden)
51        output = self.fc(output)
52        return output, hidden
53
54vocab_size = len(chars)
55model = LSTMModel(vocab_size, hidden_size, num_layers)
56criterion = nn.CrossEntropyLoss()
57optimizer = optim.Adam(model.parameters(), lr=learning_rate)
58device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
59model.to(device)
60
61for epoch in range(num_epochs):
62    hidden = None  # مقدار اولیه hidden
63
64    for inputs, targets in dataloader:
65        inputs, targets = inputs.to(device), targets.to(device)
66        optimizer.zero_grad()
67
68        # Forward pass
69        outputs, hidden = model(inputs, hidden)
70
71        # Detach hidden state to avoid graph dependency issues
72        hidden = (hidden[0].detach(), hidden[1].detach())
73
74        # Compute loss
75        loss = criterion(outputs.view(-1, vocab_size), targets.view(-1))
76
77        # Backpropagation
78        loss.backward()
79        optimizer.step()
80
81    print(f"Epoch {epoch+1}/{num_epochs}, Loss: {loss.item():.4f}")
82
83    print(f"Epoch {epoch+1}/{num_epochs}, Loss: {total_loss / len(dataloader):.4f}")
84def generate_text(model, start_text, length=200):
85    model.eval()
86    input_seq = torch.tensor([char_to_idx[ch] for ch in start_text], dtype=torch.long).unsqueeze(0).to(device)
87    hidden = None
88    generated_text = start_text
89
90    for _ in range(length):
91        output, hidden = model(input_seq, hidden)
92        next_char_idx = torch.argmax(output[:, -1, :]).item()
93        generated_text += idx_to_char[next_char_idx]
94        input_seq = torch.cat([input_seq[:, 1:], torch.tensor([[next_char_idx]], dtype=torch.long).to(device)], dim=1)
95
96    return generated_text
97
98# تست تولید متن
99start_text = "Once upon a time"
100print(generate_text(model, start_text, 200))
101