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

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4from torch.utils.data import Dataset, DataLoader, random_split
5from torch.cuda.amp import autocast, GradScaler
6from torch.utils.tensorboard import SummaryWriter
7import numpy as np
8from tqdm import tqdm
9import json
10import argparse
11from datetime import datetime
12
13# Configuration with Transformer-specific parameters
14CONFIG = {
15    "FILE_PATH": 'dataset.txt',
16    "SEQ_LENGTH": 32,
17    "BATCH_SIZE": 8,
18    "EPOCHS": 1,
19    "EMBEDDING_DIM": 64,
20    "N_HEADS": 1,
21    "FFN_DIM": 64,
22    "NUM_LAYERS": 3,
23    "DROPOUT": 0.1,
24    "LEARNING_RATE": 0.0005,
25    "WEIGHT_DECAY": 0.01,
26    "CLIP_GRAD": 1.0,
27    "LABEL_SMOOTHING": 0.1,
28    "GRAD_ACCUM_STEPS": 2,
29    "VAL_SPLIT": 0.1,
30    "EARLY_STOP_PATIENCE": 3,
31    "MODEL_SAVE_PATH": "transformer_lm_model.pth",
32    "TEMPERATURE": 0.7,
33    "TOP_K": 50,
34    "TOP_P": 0.9,
35    "LOG_DIR": "runs"
36}
37
38# Device configuration
39device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
40scaler = GradScaler(enabled=device.type == 'cuda')
41
42# Handle command-line arguments
43parser = argparse.ArgumentParser()
44parser.add_argument('--config', type=str, help='Path to config JSON file')
45args = parser.parse_args()
46
47if args.config:
48    with open(args.config) as f:
49        CONFIG.update(json.load(f))
50
51# Initialize TensorBoard
52writer = SummaryWriter(f"{CONFIG['LOG_DIR']}/{datetime.now().strftime('%Y%m%d-%H%M%S')}")
53
54# Read and process text
55with open(CONFIG["FILE_PATH"], 'r', encoding='utf-8') as f:
56    text = f.read()
57
58# Vocabulary setup
59chars = sorted(list(set(text)))
60vocab_size = len(chars)
61char_to_idx = {ch: i for i, ch in enumerate(chars)}
62idx_to_char = {i: ch for i, ch in enumerate(chars)}
63
64# Encode text
65encoded_text = np.array([char_to_idx[ch] for ch in text])
66
67# Dataset class with memory mapping
68class TextDataset(Dataset):
69    def __init__(self, data, seq_length):
70        self.data = torch.from_numpy(data).long()
71        self.seq_length = seq_length
72        
73    def __len__(self):
74        return len(self.data) - self.seq_length - 1
75    
76    def __getitem__(self, idx):
77        x = self.data[idx:idx+self.seq_length]
78        y = self.data[idx+1:idx+self.seq_length+1]
79        return x, y
80
81dataset = TextDataset(encoded_text, CONFIG["SEQ_LENGTH"])
82val_size = int(len(dataset) * CONFIG["VAL_SPLIT"])
83train_size = len(dataset) - val_size
84train_dataset, val_dataset = random_split(dataset, [train_size, val_size])
85
86train_loader = DataLoader(train_dataset, batch_size=CONFIG["BATCH_SIZE"], 
87                         shuffle=True, pin_memory=True, num_workers=4)
88val_loader = DataLoader(val_dataset, batch_size=CONFIG["BATCH_SIZE"], 
89                       pin_memory=True, num_workers=4)
90
91# Transformer-based Language Model
92class TransformerLM(nn.Module):
93    def __init__(self):
94        super().__init__()
95        self.embedding = nn.Embedding(vocab_size, CONFIG["EMBEDDING_DIM"])
96        self.pos_embed = nn.Embedding(CONFIG["SEQ_LENGTH"], CONFIG["EMBEDDING_DIM"])
97        
98        self.transformer = nn.TransformerEncoder(
99            nn.TransformerEncoderLayer(
100                d_model=CONFIG["EMBEDDING_DIM"],
101                nhead=CONFIG["N_HEADS"],
102                dim_feedforward=CONFIG["FFN_DIM"],
103                dropout=CONFIG["DROPOUT"],
104                activation='gelu',
105                batch_first=True
106            ),
107            num_layers=CONFIG["NUM_LAYERS"]
108        )
109        
110        self.ln = nn.LayerNorm(CONFIG["EMBEDDING_DIM"])
111        self.fc = nn.Linear(CONFIG["EMBEDDING_DIM"], vocab_size)
112        
113        self.apply(self._init_weights)
114        
115    def _init_weights(self, module):
116        if isinstance(module, nn.Linear):
117            nn.init.xavier_uniform_(module.weight)
118            if module.bias is not None:
119                nn.init.constant_(module.bias, 0)
120        elif isinstance(module, nn.Embedding):
121            nn.init.xavier_uniform_(module.weight)
122            
123    def forward(self, x, mask=None):
124        batch_size, seq_len = x.size()
125        positions = torch.arange(seq_len, device=device).expand(batch_size, seq_len)
126        x = self.embedding(x) + self.pos_embed(positions)
127        
128        if mask is None:
129            mask = nn.Transformer.generate_square_subsequent_mask(seq_len).to(device)
130            
131        x = self.transformer(x, mask)
132        x = self.ln(x)
133        return self.fc(x), None  # Return None for compatibility with generation code
134
135model = TransformerLM().to(device)
136optimizer = torch.optim.AdamW(model.parameters(), 
137                             lr=CONFIG["LEARNING_RATE"], 
138                             weight_decay=CONFIG["WEIGHT_DECAY"])
139scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
140    optimizer, mode='min', factor=0.5, patience=2)
141
142# Training loop with advanced features
143best_val_loss = float('inf')
144patience_counter = 0
145
146for epoch in range(CONFIG["EPOCHS"]):
147    model.train()
148    train_loss = 0
149    progress_bar = tqdm(train_loader, desc=f'Epoch {epoch+1}/{CONFIG["EPOCHS"]}')
150    
151    for i, (inputs, targets) in enumerate(progress_bar):
152        inputs, targets = inputs.to(device), targets.to(device)
153        
154        with autocast(enabled=device.type == 'cuda'):
155            outputs, _ = model(inputs)
156            logits = outputs.view(-1, vocab_size)
157            targets = targets.view(-1)
158            
159            if CONFIG["LABEL_SMOOTHING"]:
160                loss = F.cross_entropy(logits, targets, 
161                                      label_smoothing=CONFIG["LABEL_SMOOTHING"])
162            else:
163                loss = F.cross_entropy(logits, targets)
164            
165            loss = loss / CONFIG["GRAD_ACCUM_STEPS"]
166        
167        scaler.scale(loss).backward()
168        
169        if (i + 1) % CONFIG["GRAD_ACCUM_STEPS"] == 0:
170            scaler.unscale_(optimizer)
171            nn.utils.clip_grad_norm_(model.parameters(), CONFIG["CLIP_GRAD"])
172            scaler.step(optimizer)
173            scaler.update()
174            optimizer.zero_grad()
175            
176            lr = optimizer.param_groups[0]['lr']
177            progress_bar.set_postfix({'loss': loss.item() * CONFIG["GRAD_ACCUM_STEPS"], 'lr': lr})
178        
179        train_loss += loss.item() * CONFIG["GRAD_ACCUM_STEPS"]
180    
181    # Validation phase
182    model.eval()
183    val_loss = 0
184    with torch.no_grad():
185        for inputs, targets in val_loader:
186            inputs, targets = inputs.to(device), targets.to(device)
187            outputs, _ = model(inputs)
188            loss = F.cross_entropy(outputs.view(-1, vocab_size), targets.view(-1))
189            val_loss += loss.item()
190    
191    avg_train_loss = train_loss / len(train_loader)
192    avg_val_loss = val_loss / len(val_loader)
193    scheduler.step(avg_val_loss)
194    
195    # Log metrics
196    writer.add_scalar('Loss/train', avg_train_loss, epoch)
197    writer.add_scalar('Loss/val', avg_val_loss, epoch)
198    writer.add_scalar('Learning Rate', optimizer.param_groups[0]['lr'], epoch)
199    writer.add_scalar('Perplexity/train', np.exp(avg_train_loss), epoch)
200    writer.add_scalar('Perplexity/val', np.exp(avg_val_loss), epoch)
201    
202    print(f'Epoch {epoch+1} | Train Loss: {avg_train_loss:.4f} | Val Loss: {avg_val_loss:.4f}')
203    
204    # Early stopping and checkpointing
205    if avg_val_loss < best_val_loss:
206        best_val_loss = avg_val_loss
207        torch.save({
208            'epoch': epoch,
209            'model_state_dict': model.state_dict(),
210            'optimizer_state_dict': optimizer.state_dict(),
211            'scheduler_state_dict': scheduler.state_dict(),
212            'config': CONFIG
213        }, CONFIG["MODEL_SAVE_PATH"])
214        patience_counter = 0
215    else:
216        patience_counter += 1
217        if patience_counter >= CONFIG["EARLY_STOP_PATIENCE"]:
218            print("Early stopping triggered")
219            break
220
221writer.close()
222print(f'Best model saved to {CONFIG["MODEL_SAVE_PATH"]} with validation loss: {best_val_loss:.4f}')
223
224# Advanced generation with multiple sampling strategies
225def generate_text(model, start_str, length=200, temperature=CONFIG["TEMPERATURE"],
226                 top_k=CONFIG["TOP_K"], top_p=CONFIG["TOP_P"]):
227    model.eval()
228    chars = list(start_str)
229    input_seq = torch.tensor([char_to_idx[ch] for ch in chars], device=device).unsqueeze(0)
230    
231    with torch.no_grad():
232        for _ in tqdm(range(length), desc="Generating text"):
233            mask = nn.Transformer.generate_square_subsequent_mask(input_seq.size(1)).to(device)
234            outputs, _ = model(input_seq[:, -CONFIG["SEQ_LENGTH"]:], mask)
235            logits = outputs[:, -1] / temperature
236            
237            # Apply nucleus sampling first
238            if top_p > 0:
239                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
240                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
241                sorted_indices_to_remove = cumulative_probs > top_p
242                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
243                sorted_indices_to_remove[..., 0] = 0
244                indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
245                logits = logits.masked_fill(indices_to_remove, float('-inf'))
246            
247            # Then apply top-k filtering
248            if top_k > 0:
249                top_k = min(top_k, logits.size(-1))
250                indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
251                logits = logits.masked_fill(indices_to_remove, float('-inf'))
252            
253            probs = F.softmax(logits, dim=-1)
254            next_char = torch.multinomial(probs, num_samples=1)
255            chars.append(idx_to_char[next_char.item()])
256            input_seq = torch.cat([input_seq, next_char], dim=1)
257    
258    return ''.join(chars)
259
260# Generate examples with different parameters
261print("\nConservative sampling:")
262print(generate_text(model, "The ", temperature=0.5, top_p=0))
263
264print("\nCreative sampling:")
265print(generate_text(model, "Once ", temperature=1.2, top_p=0.9))
266
267print("\nTop-k sampling:")
268print(generate_text(model, "In ", top_k=50))
269
270print("\nCombined sampling:")
271print(generate_text(model, "Artificial intelligence ", temperature=0.8, top_k=50, top_p=0.9))