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  # Context window size
9BATCH_SIZE = 8
10EPOCHS = 1
11EMBEDDING_DIM = 64
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(list(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])
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 - 1
36    
37    def __getitem__(self, idx):
38        x = self.data[idx:idx+self.seq_length]
39        y = self.data[idx+1:idx+self.seq_length+1]
40        return torch.from_numpy(x).long(), torch.from_numpy(y).long()
41
42dataset = TextDataset(encoded_text, SEQ_LENGTH)
43dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
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
59model = CharLM()
60criterion = nn.CrossEntropyLoss()
61optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
62
63# Training loop
64for epoch in range(EPOCHS):
65    model.train()
66    total_loss = 0
67    
68    for inputs, targets in dataloader:
69        optimizer.zero_grad()
70        outputs, _ = model(inputs)
71        loss = criterion(outputs.reshape(-1, vocab_size), targets.reshape(-1))
72        loss.backward()
73        optimizer.step()
74        total_loss += loss.item()
75    
76    print(f'Epoch {epoch+1}/{EPOCHS}, Loss: {total_loss/len(dataloader):.4f}')
77
78# Enhanced Text Generation Functions
79def generate_text(model, start_str, length=100, temperature=0.7, top_k=0):
80    """
81    Generate text with temperature scaling and top-k sampling
82    temperature: >1.0 more random, <1.0 more conservative
83    top_k: 0=no sampling, >0 top-k tokens to consider
84    """
85    model.eval()
86    chars = [ch for ch in start_str]
87    input_seq = torch.tensor([char_to_idx[ch] for ch in chars]).unsqueeze(0)
88    hidden = None
89    
90    with torch.no_grad():
91        for _ in range(length):
92            outputs, hidden = model(input_seq, hidden)
93            logits = outputs[0, -1] / temperature
94            
95            if top_k > 0:
96                top_vals, top_idx = torch.topk(logits, top_k)
97                logits[logits < top_vals[-1]] = -float('Inf')
98            
99            probs = torch.softmax(logits, dim=-1)
100            next_char = torch.multinomial(probs, num_samples=1).item()
101            chars.append(idx_to_char[next_char])
102            input_seq = torch.tensor([[next_char]])
103    
104    return ''.join(chars)
105
106# Generation examples
107print("\nGreedy sampling (temperature=0.5):")
108print(generate_text(model, "The ", temperature=0.5))
109
110print("\nCreative sampling (temperature=1.2):")
111print(generate_text(model, "Once ", temperature=1.2))
112
113print("\nTop-k sampling (k=5):")
114print(generate_text(model, "In ", top_k=5))
115
116print("\nCombined (temp=0.7, top_k=3):")
117print(generate_text(model, "AI ", temperature=0.7, top_k=3))
118
119