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

1import torch
2import torch.nn as nn
3from torch.utils.data import Dataset, DataLoader
4import random
5
6# Hyperparameters
7SEQ_LENGTH = 100  # Length of input sequences
8BATCH_SIZE = 64
9HIDDEN_SIZE = 256
10NUM_LAYERS = 2
11LEARNING_RATE = 0.001
12NUM_EPOCHS = 50
13DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
14
15# Dataset Preparation
16class CharDataset(Dataset):
17    def __init__(self, text, seq_length):
18        self.text = text
19        self.seq_length = seq_length
20        self.chars = sorted(list(set(text)))
21        self.char_to_idx = {c: i for i, c in enumerate(self.chars)}
22        self.idx_to_char = {i: c for i, c in enumerate(self.chars)}
23        self.encoded_text = [self.char_to_idx[c] for c in text]
24        
25    def __len__(self):
26        return len(self.text) - self.seq_length
27    
28    def __getitem__(self, idx):
29        inputs = torch.tensor(self.encoded_text[idx:idx+self.seq_length])
30        targets = torch.tensor(self.encoded_text[idx+1:idx+self.seq_length+1])
31        return inputs, targets
32
33# Model Definition
34class CharRNN(nn.Module):
35    def __init__(self, input_size, hidden_size, output_size, num_layers):
36        super().__init__()
37        self.embedding = nn.Embedding(input_size, hidden_size)
38        self.lstm = nn.LSTM(hidden_size, hidden_size, num_layers, batch_first=True)
39        self.fc = nn.Linear(hidden_size, output_size)
40        
41    def forward(self, x, hidden=None):
42        x = self.embedding(x)
43        out, hidden = self.lstm(x, hidden)
44        out = self.fc(out)
45        return out, hidden
46
47# Load your text data (replace with your own text file)
48with open('dataset.txt', 'r', encoding='utf-8') as f:
49    text = f.read()
50
51# Create dataset and dataloader
52dataset = CharDataset(text, SEQ_LENGTH)
53dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
54
55# Initialize model
56model = CharRNN(
57    input_size=len(dataset.chars),
58    hidden_size=HIDDEN_SIZE,
59    output_size=len(dataset.chars),
60    num_layers=NUM_LAYERS
61).to(DEVICE)
62
63criterion = nn.CrossEntropyLoss()
64optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
65
66# Training loop
67for epoch in range(NUM_EPOCHS):
68    model.train()
69    total_loss = 0
70    
71    for inputs, targets in dataloader:
72        inputs, targets = inputs.to(DEVICE), targets.to(DEVICE)
73        
74        optimizer.zero_grad()
75        outputs, _ = model(inputs)
76        loss = criterion(outputs.view(-1, len(dataset.chars)), targets.view(-1))
77        loss.backward()
78        optimizer.step()
79        
80        total_loss += loss.item()
81    
82    avg_loss = total_loss / len(dataloader)
83    print(f'Epoch {epoch+1}/{NUM_EPOCHS}, Loss: {avg_loss:.4f}')
84
85# Text generation function
86def generate(model, start_str, length=100, temperature=0.8):
87    model.eval()
88    chars = [c for c in start_str]
89    hidden = None
90    
91    with torch.no_grad():
92        # Initialize hidden state with starting string
93        for char in chars[:-1]:
94            x = torch.tensor([[dataset.char_to_idx[char]]]).to(DEVICE)
95            _, hidden = model(x, hidden)
96        
97        # Generate remaining characters
98        x = torch.tensor([[dataset.char_to_idx[chars[-1]]]]).to(DEVICE)
99        
100        for _ in range(length):
101            output, hidden = model(x, hidden)
102            probs = torch.softmax(output / temperature, dim=-1).cpu()
103            char_idx = torch.multinomial(probs.view(-1), 1).item()
104            chars.append(dataset.idx_to_char[char_idx])
105            x = torch.tensor([[char_idx]]).to(DEVICE)
106    
107    return ''.join(chars)
108
109# Generate sample text
110print(generate(model, start_str="The ", length=500))