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
5from tqdm import tqdm
6import os
7
8# Configuration
9class Config:
10    FILE_PATH = 'dataset.txt'
11    SEQ_LENGTH = 8  # Context window size
12    BATCH_SIZE = 8
13    EPOCHS = 1
14    EMBEDDING_DIM = 16
15    HIDDEN_DIM = 32
16    LEARNING_RATE = 0.01
17    DROPOUT_RATE = 0.2
18    MODEL_SAVE_PATH = "char_lm_model_f4.pth"
19    GRAD_CLIP = 1.0
20    TOP_K = 5  # For generation
21    NUM_LAYERS = 4  # GRU layers
22
23# Check for GPU availability
24device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
25print(f"Using device: {device}")
26
27# Read and process text
28with open(Config.FILE_PATH, 'r', encoding='utf-8') as f:
29    text = f.read()
30
31# Vocabulary setup
32chars = sorted(list(set(text)))
33vocab_size = len(chars)
34char_to_idx = {ch: i for i, ch in enumerate(chars)}
35idx_to_char = {i: ch for i, ch in enumerate(chars)}
36
37# Encode text
38encoded_text = np.array([char_to_idx[ch] for ch in text])
39
40# Dataset class
41class TextDataset(Dataset):
42    def __init__(self, data, seq_length):
43        self.data = torch.tensor(data, dtype=torch.long)
44        self.seq_length = seq_length
45        
46    def __len__(self):
47        return len(self.data) - self.seq_length - 1
48    
49    def __getitem__(self, idx):
50        x = self.data[idx:idx+self.seq_length]
51        y = self.data[idx+1:idx+self.seq_length+1]
52        return x, y
53
54dataset = TextDataset(encoded_text, Config.SEQ_LENGTH)
55dataloader = DataLoader(dataset, batch_size=Config.BATCH_SIZE, shuffle=True, num_workers=4)
56
57# Model architecture
58class CharLM(nn.Module):
59    def __init__(self, vocab_size, config):
60        super(CharLM, self).__init__()
61        self.embedding = nn.Embedding(vocab_size, config.EMBEDDING_DIM)
62        self.gru = nn.GRU(config.EMBEDDING_DIM, config.HIDDEN_DIM,
63                         num_layers=config.NUM_LAYERS,
64                         batch_first=True,
65                         dropout=config.DROPOUT_RATE if config.NUM_LAYERS > 1 else 0)
66        self.dropout = nn.Dropout(config.DROPOUT_RATE)
67        self.fc = nn.Linear(config.HIDDEN_DIM, vocab_size)
68        
69        self.init_weights()
70        
71    def init_weights(self):
72        for name, param in self.named_parameters():
73            if 'weight' in name:
74                nn.init.xavier_normal_(param)
75            elif 'bias' in name:
76                nn.init.constant_(param, 0.0)
77                
78    def forward(self, x, hidden=None):
79        x = self.embedding(x)
80        out, hidden = self.gru(x, hidden)
81        out = self.dropout(out)
82        out = self.fc(out)
83        return out, hidden
84
85model = CharLM(vocab_size, Config).to(device)
86criterion = nn.CrossEntropyLoss()
87optimizer = torch.optim.Adam(model.parameters(), lr=Config.LEARNING_RATE)
88scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=2)
89
90# Training loop
91best_loss = float('inf')
92for epoch in range(Config.EPOCHS):
93    model.train()
94    epoch_loss = 0
95    progress_bar = tqdm(dataloader, desc=f'Epoch {epoch+1}/{Config.EPOCHS}')
96    
97    for inputs, targets in progress_bar:
98        inputs, targets = inputs.to(device), targets.to(device)
99        
100        optimizer.zero_grad()
101        outputs, _ = model(inputs)
102        loss = criterion(outputs.view(-1, vocab_size), targets.view(-1))
103        loss.backward()
104        
105        # Gradient clipping
106        nn.utils.clip_grad_norm_(model.parameters(), Config.GRAD_CLIP)
107        
108        optimizer.step()
109        epoch_loss += loss.item()
110        
111        # Update progress bar
112        progress_bar.set_postfix({'loss': f'{loss.item():.4f}'})
113    
114    avg_loss = epoch_loss / len(dataloader)
115    scheduler.step(avg_loss)
116    
117    # Save best model
118    if avg_loss < best_loss:
119        best_loss = avg_loss
120        torch.save({
121            'model_state_dict': model.state_dict(),
122            'char_to_idx': char_to_idx,
123            'idx_to_char': idx_to_char,
124            'config': Config
125        }, Config.MODEL_SAVE_PATH)
126    
127    print(f'Epoch {epoch+1} complete. Avg loss: {avg_loss:.4f}')
128
129print(f'Model saved to {Config.MODEL_SAVE_PATH}')
130
131# Improved Text Generation Function
132def generate_text(model, start_str, length=100, temperature=1.0, top_k=None):
133    """
134    Generate text with temperature scaling and top-k sampling
135    - Maintains proper context window size
136    - Handles start strings of any length
137    - Returns original start_str + generated text
138    """
139    model.eval()
140    initial_chars = list(start_str)
141    generated = initial_chars.copy()
142    
143    # Initialize sequence with proper length
144    if len(initial_chars) < Config.SEQ_LENGTH:
145        # Pad with repeated characters if needed
146        padded = (initial_chars * Config.SEQ_LENGTH)[:Config.SEQ_LENGTH]
147    else:
148        # Take last SEQ_LENGTH characters
149        padded = initial_chars[-Config.SEQ_LENGTH:]
150    
151    current_seq = torch.tensor([char_to_idx[c] for c in padded], 
152                              dtype=torch.long, device=device).unsqueeze(0)
153    
154    with torch.no_grad():
155        for _ in range(length):
156            outputs, _ = model(current_seq)
157            logits = outputs[:, -1, :] / temperature
158            
159            if top_k is not None and top_k > 0:
160                top_values, top_indices = torch.topk(logits, top_k)
161                logits[logits < top_values[:, -1:]] = -float('Inf')
162                
163            probs = torch.softmax(logits, dim=-1)
164            next_idx = torch.multinomial(probs, num_samples=1)
165            next_char = idx_to_char[next_idx.item()]
166            
167            generated.append(next_char)
168            # Update sequence: remove first character, add new
169            current_seq = torch.cat([current_seq[:, 1:], next_idx.unsqueeze(1)], dim=1)
170    
171    # Return original start string plus generated text
172    return start_str + ''.join(generated[len(initial_chars):])
173
174# Load best model for generation
175checkpoint = torch.load(Config.MODEL_SAVE_PATH, map_location=device)
176model.load_state_dict(checkpoint['model_state_dict'])
177char_to_idx = checkpoint['char_to_idx']
178idx_to_char = checkpoint['idx_to_char']
179
180# Generation examples
181print("\n--- Generation Examples ---")
182for prompt in ["The ", "Once ", "In ", "AI "]:
183    generated = generate_text(
184        model, 
185        prompt, 
186        length=100,
187        temperature=0.4,
188        top_k=Config.TOP_K
189    )
190    print(f"\nPrompt: '{prompt}'\n{generated}\n{'-'*50}")