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

1import torch
2import torch.nn as nn
3from torch.utils.data import Dataset, DataLoader
4import numpy as np
5
6# Configuration
7FILE_PATH = 'dataset.txt'
8SEQ_LENGTH = 32
9BATCH_SIZE = 8
10EPOCHS = 5
11EMBEDDING_DIM = 32
12HIDDEN_DIM = 64
13LEARNING_RATE = 0.01
14
15# Read and process text
16with open(FILE_PATH, 'r', encoding='utf-8') as f:
17    text = f.read()
18
19# Vocabulary setup
20chars = sorted(set(text))
21vocab_size = len(chars)
22char_to_idx = {ch: i for i, ch in enumerate(chars)}
23idx_to_char = {i: ch for i, ch in enumerate(chars)}
24
25# Encode text
26encoded_text = np.array([char_to_idx[ch] for ch in text], dtype=np.int64)
27
28# Dataset class
29class TextDataset(Dataset):
30    def __init__(self, data, seq_length):
31        self.data = data
32        self.seq_length = seq_length
33        
34    def __len__(self):
35        return len(self.data) - self.seq_length
36    
37    def __getitem__(self, idx):
38        x = torch.tensor(self.data[idx:idx+self.seq_length], dtype=torch.long)
39        y = torch.tensor(self.data[idx+1:idx+self.seq_length+1], dtype=torch.long)
40        return x, y
41
42dataset = TextDataset(encoded_text, SEQ_LENGTH)
43dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)
44
45# Model architecture
46class CharLM(nn.Module):
47    def __init__(self):
48        super(CharLM, self).__init__()
49        self.embedding = nn.Embedding(vocab_size, EMBEDDING_DIM)
50        self.rnn = nn.GRU(EMBEDDING_DIM, HIDDEN_DIM, batch_first=True)
51        self.fc = nn.Linear(HIDDEN_DIM, vocab_size)
52        
53    def forward(self, x, hidden=None):
54        x = self.embedding(x)
55        out, hidden = self.rnn(x, hidden)
56        out = self.fc(out)
57        return out, hidden
58
59device = torch.device("cpu")
60model = CharLM().to(device)
61criterion = nn.CrossEntropyLoss()
62optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
63
64# Training loop
65for epoch in range(EPOCHS):
66    model.train()
67    total_loss = 0
68    
69    for inputs, targets in dataloader:
70        inputs, targets = inputs.to(device), targets.to(device)
71
72        optimizer.zero_grad()
73        outputs, _ = model(inputs)
74        loss = criterion(outputs.reshape(-1, vocab_size), targets.reshape(-1))
75        loss.backward()
76        optimizer.step()
77        total_loss += loss.item()
78    
79    print(f'Epoch {epoch+1}/{EPOCHS}, Loss: {total_loss / len(dataloader):.4f}')
80
81# Enhanced Text Generation Function
82def generate_text(model, start_str, length=100, temperature=0.7, top_k=0):
83    """
84    Generate text with temperature scaling and top-k sampling
85    temperature: >1.0 more random, <1.0 more conservative
86    top_k: 0=no sampling, >0 top-k tokens to consider
87    """
88    model.eval()
89    chars = [ch for ch in start_str]
90    input_seq = torch.tensor([char_to_idx[ch] for ch in chars], dtype=torch.long).unsqueeze(0).to(device)
91    hidden = None
92    
93    with torch.no_grad():
94        for _ in range(length):
95            outputs, hidden = model(input_seq, hidden)
96            logits = outputs[0, -1] / temperature
97
98            if top_k > 0:
99                top_vals, top_idx = torch.topk(logits, top_k)
100                logits[logits < top_vals[-1]] = -float('Inf')
101            
102            probs = torch.softmax(logits, dim=-1)
103            next_char = torch.multinomial(probs, num_samples=1).item()
104            chars.append(idx_to_char[next_char])
105            input_seq = torch.tensor([[next_char]], dtype=torch.long).to(device)
106    
107    return ''.join(chars)
108
109# Text generation examples
110print("\nGreedy sampling (temperature=0.5):")
111print(generate_text(model, "The ", temperature=0.5))
112
113print("\nCreative sampling (temperature=1.2):")
114print(generate_text(model, "Once ", temperature=1.2))
115
116print("\nTop-k sampling (k=5):")
117print(generate_text(model, "In ", top_k=5))
118
119print("\nCombined (temp=0.7, top_k=3):")
120print(generate_text(model, "AI ", temperature=0.7, top_k=3))
121