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
5
6# Hyperparameters
7sequence_length = 100
8batch_size = 64
9hidden_size = 128
10num_layers = 1
11learning_rate = 0.01
12num_epochs = 100
13device = torch.device('cpu')
14
15# Sample text data
16text = "Hello, this is a simple language model for text generation using PyTorch."
17chars = sorted(list(set(text)))
18vocab_size = len(chars)
19char_to_idx = {ch: i for i, ch in enumerate(chars)}
20idx_to_char = {i: ch for i, ch in enumerate(chars)}
21
22# Prepare data
23data = [char_to_idx[ch] for ch in text]
24data = torch.tensor(data, dtype=torch.long).to(device)
25
26# Create batches
27def create_batches(data, batch_size, sequence_length):
28    num_batches = len(data) // (batch_size * sequence_length)
29    data = data[:num_batches * batch_size * sequence_length]
30    data = data.view(batch_size, -1)
31    for i in range(0, data.size(1), sequence_length):
32        x = data[:, i:i+sequence_length]
33        y = data[:, (i+1):(i+1)+sequence_length]
34        yield x, y
35
36# Define the model
37class CharRNN(nn.Module):
38    def __init__(self, vocab_size, hidden_size, num_layers):
39        super(CharRNN, self).__init__()
40        self.hidden_size = hidden_size
41        self.num_layers = num_layers
42        self.embedding = nn.Embedding(vocab_size, hidden_size)
43        self.lstm = nn.LSTM(hidden_size, hidden_size, num_layers, batch_first=True)
44        self.fc = nn.Linear(hidden_size, vocab_size)
45
46    def forward(self, x, hidden):
47        x = self.embedding(x)
48        out, hidden = self.lstm(x, hidden)
49        out = self.fc(out)
50        return out, hidden
51
52    def init_hidden(self, batch_size):
53        return (torch.zeros(self.num_layers, batch_size, self.hidden_size).to(device),
54                torch.zeros(self.num_layers, batch_size, self.hidden_size).to(device))
55
56# Initialize model, loss, and optimizer
57model = CharRNN(vocab_size, hidden_size, num_layers).to(device)
58criterion = nn.CrossEntropyLoss()
59optimizer = optim.Adam(model.parameters(), lr=learning_rate)
60
61# Training loop
62for epoch in range(num_epochs):
63    hidden = model.init_hidden(batch_size)
64    for i, (x, y) in enumerate(create_batches(data, batch_size, sequence_length)):
65        x, y = x.to(device), y.to(device)
66        hidden = tuple(h.detach() for h in hidden)  # Detach hidden state
67        optimizer.zero_grad()
68        output, hidden = model(x, hidden)
69        loss = criterion(output.transpose(1, 2), y)
70        loss.backward()
71        optimizer.step()
72
73        if (i+1) % 10 == 0:
74            print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}], Loss: {loss.item():.4f}')
75
76# Generate text
77def generate_text(model, start_str, length=100):
78    model.eval()
79    chars = [char_to_idx[ch] for ch in start_str]
80    hidden = model.init_hidden(1)
81    for i in range(length):
82        x = torch.tensor([chars[-1]], dtype=torch.long).unsqueeze(0).to(device)
83        with torch.no_grad():
84            output, hidden = model(x, hidden)
85        prob = torch.softmax(output.squeeze(), dim=0).cpu().numpy()
86        next_char = np.random.choice(vocab_size, p=prob)
87        chars.append(next_char)
88    return ''.join([idx_to_char[ch] for ch in chars])
89
90# Test text generation
91start_str = "Hello"
92generated_text = generate_text(model, start_str, length=200)
93print(generated_text)