Upload StoryLLM.py

StoryLLM.py

@@ -0,0 +1,346 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+import tiktoken
+from datasets import load_dataset
+import matplotlib.pyplot as plt
+import numpy as np
+from datetime import datetime
+import os
+
+# Define hyperparameters
+vocab_size = 50257
+n_heads = 8
+n_layers = 6
+head_size = 64
+n_embd = 512
+block_size = 128
+dropout = 0.1
+learning_rate = 3e-4
+weight_decay = 0.1
+
+# Set Hugging Face cache directories on the external disk
+os.environ['HF_HOME'] = '/media/adrian/FamilyBackup/adrian_ai_workspace/hf_cache'
+os.environ['HF_DATASETS_CACHE'] = '/media/adrian/FamilyBackup/adrian_ai_workspace/datasets_cache'
+
+# Load the BookCorpus dataset and ensure it's cached on the external disk
+dataset = load_dataset("bookcorpus", cache_dir='/media/adrian/FamilyBackup/adrian_ai_workspace/')
+
+# Keep only 10% of the dataset
+total_samples = len(dataset["train"])
+one_percent_samples = int(total_samples * 0.001)
+dataset_subset = dataset["train"].select(range(one_percent_samples))  # Select only the first 1%
+
+# Split the subset into train (90%) and test (10%)
+split_dataset = dataset_subset.train_test_split(test_size=0.1)  # 10% for testing
+train_dataset = split_dataset["train"]
+test_dataset = split_dataset["test"]
+
+# Print the size of the train and the test sets
+print(f"Train size: {len(train_dataset)}")
+print(f"Test size: {len(test_dataset)}")
+
+# Initialize the tiktoken encoder
+enc = tiktoken.get_encoding("gpt2")
+
+# Define the tokenization function
+def tokenize_function(examples):
+    return {
+        "input_ids": [enc.encode(text) for text in examples["text"]],
+        "attention_mask": [[1] * len(enc.encode(text)) for text in examples["text"]]
+    }
+
+# Function to pad or truncate sequences
+def pad_or_truncate(batch):
+    max_length = 512
+    for key in ['input_ids', 'attention_mask']:
+        batch[key] = [
+            seq[:max_length] + [0] * (max_length - len(seq)) if len(seq) < max_length else seq[:max_length]
+            for seq in batch[key]
+        ]
+    return batch
+
+# Tokenize and process the datasets
+def process_dataset(dataset, split_name):
+    # Tokenize
+    tokenized_dataset = dataset.map(
+        tokenize_function,
+        batched=True,
+        num_proc=20,
+        remove_columns=dataset.column_names
+    )
+
+    # Pad or truncate
+    processed_dataset = tokenized_dataset.map(
+        pad_or_truncate,
+        batched=True,
+        num_proc=20,
+    )
+
+    # Set format to PyTorch tensors
+    processed_dataset.set_format(type="torch", columns=["input_ids", "attention_mask"])
+
+    return processed_dataset
+
+# Process both train and test datasets
+train_dataset = process_dataset(train_dataset, "train")
+test_dataset = process_dataset(test_dataset, "test")
+
+# Print some examples
+print(f"Example train data: {train_dataset[0]}")
+print(f"Example test data: {test_dataset[0]}")
+
+# Create DataLoaders
+train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=8, shuffle=True)
+test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=8, shuffle=False)
+
+# Print an example batch
+for batch in train_loader:
+    print(f"Batch input ids shape: {batch['input_ids'].shape}")
+    print(f"Batch attention mask shape: {batch['attention_mask'].shape}")
+    break
+
+# Print an example batch
+for batch in train_loader:
+    print(f"Batch input ids shape: {batch['input_ids'].shape}")
+    print(f"Batch attention mask shape: {batch['attention_mask'].shape}")
+    break
+
+# Define model
+class Head(nn.Module):
+    """ One head of self-attention """
+    def __init__(self, head_size, n_embd, block_size, dropout):
+        super().__init__()
+        self.key = nn.Linear(n_embd, head_size, bias=False)
+        self.query = nn.Linear(n_embd, head_size, bias=False)
+        self.value = nn.Linear(n_embd, head_size, bias=False)
+        self.register_buffer("tril", torch.tril(torch.ones(block_size, block_size)))
+
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        B, T, C = x.shape
+        k = self.key(x)
+        q = self.query(x)
+        v = self.value(x)
+
+        assert C == self.key.in_features, f"Input size {C} doesn't match expected size {self.key.in_features}"
+
+        wei = q @ k.transpose(-2, -1) * k.shape[-1]**-0.5
+        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
+        wei = F.softmax(wei, dim=-1)
+        wei = self.dropout(wei)
+
+        out = wei @ v
+        return out
+
+class MultiHeadAttention(nn.Module):
+    """ Multiple heads of self-attention in parallel """
+
+    def __init__(self, n_heads, head_size, n_embd, dropout):
+        super().__init__()
+        self.heads = nn.ModuleList([Head(head_size, n_embd, block_size, dropout) for _ in range(n_heads)])
+        self.proj = nn.Linear(n_heads *  head_size, n_embd)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        # Collects the outputs from each head
+        head_outputs = [head(x) for head in self.heads]
+        # Concatenate the outputs
+        concatenated = torch.cat(head_outputs, dim=-1)
+        # Apply linear transformation and dropout
+        out = self.proj(concatenated)
+        out = self.dropout(out)
+        return out
+
+
+class FeedForward(nn.Module):
+    """ A simple linear layer followed by non-linearity """
+
+    def __init__(self, n_embd, dropout=0.1, expansion_factor=4):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(n_embd, expansion_factor * n_embd),
+            nn.ReLU(),
+            nn.Linear(expansion_factor * n_embd, n_embd),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+class Block(nn.Module):
+    """ Transformer block: communication followed by computation """
+
+    def __init__(self, n_embd, n_head, dropout=0.1):
+        # n_embed: embedding dimension, n_head: the number of heads we'd like
+        super().__init__()
+        head_size = n_embd // n_head
+        self.sa = MultiHeadAttention(n_head, head_size, n_embd, dropout)
+        self.ffwd = FeedForward(n_embd, dropout)
+        self.ln1 = nn.LayerNorm(n_embd)
+        self.ln2 = nn.LayerNorm(n_embd)
+
+    def forward(self, x):
+        x = x + self.sa(self.ln1(x))
+        x = x + self.ffwd(self.ln2(x))
+        return x
+
+class GPTLanguageModel(nn.Module):
+    def __init__(self, vocab_size, n_embd, block_size, n_layer, n_head, device="cpu"):
+        super().__init__()
+        self.device = device
+        self.block_size = block_size
+        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
+        self.position_embedding_table = nn.Embedding(block_size, n_embd)
+        self.blocks = nn.Sequential(*[Block(n_embd, n_head) for _ in range(n_layer)])
+        self.ln_f = nn.LayerNorm(n_embd)
+        self.lm_head = nn.Linear(n_embd, vocab_size)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.1, std=0.02)
+
+    def forward(self, idx, targets=None):
+        B, T = idx.shape
+
+        # Truncate sequence length to block_size
+        T = min(T, self.block_size)
+        idx = idx[:, :T]
+
+        # Get token embeddings for input indices
+        tok_emb = self.token_embedding_table(idx)  # (B, T, C)
+
+        # Get position embeddings (truncate to match input length)
+        pos_emb = self.position_embedding_table(torch.arange(T, device=idx.device))  # (T, C)
+
+        # Combine token and position embeddings
+        x = tok_emb + pos_emb.unsqueeze(0)  # (B, T, C)
+
+        # Apply transformer blocks
+        x = self.blocks(x)  # (B, T, C)
+
+        # Final layer normalization
+        x = self.ln_f(x)  # (B, T, C)
+
+        # Get logits for vocabulary prediction
+        logits = self.lm_head(x)  # (B, T, vocab_size)
+
+        # Optionally calculate loss if targets are provided
+        loss = None
+        if targets is not None:
+            # Ensure targets are the same size as logits
+            targets = targets[:, :T]
+            B, T, C = logits.shape
+            logits = logits.reshape(B*T, C)
+            targets = targets.reshape(B*T)
+            loss = F.cross_entropy(logits, targets)
+
+        return logits, loss
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens):
+        for _ in range(max_new_tokens):
+            idx_cond = idx[:, -self.block_size:]  # Crop to the last block_size tokens
+            logits, _ = self(idx_cond)  # Get Predictions
+            logits = logits[:, -1, :]  # Focus on the last time step
+            probs = F.softmax(logits, dim=-1)  # Get probabilities
+            idx_next = torch.multinomial(probs, num_samples=1)  # Samples from the distribution
+            idx = torch.cat((idx, idx_next), dim=1)  # Append sampled index
+        return idx
+
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+print (f"Using device: {device}")
+
+# Instantiate the model
+model = GPTLanguageModel(vocab_size, n_embd, block_size, n_layers, n_heads, device=device)
+
+# Move the model to the GPU (if available)
+model = model.to(device)
+
+# Define criterion and optimizer
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.Adam(model.parameters(), lr=learning_rate)
+
+# Training loop with progress reporting
+def batch_gh(model, criterion, optimizer, train_loader, test_loader, epochs):
+    train_losses = np.zeros(epochs)
+    test_losses = np.zeros(epochs)
+
+    for it in range(epochs):
+        model.train()  # Set model to training mode
+        t0 = datetime.now()
+        train_loss = []
+
+        for i, batch in enumerate(train_loader):
+            inputs = batch["input_ids"].to(device)
+            attention_mask = batch["attention_mask"].to(device)
+
+            # Create targets by shifting inputs by one position
+            targets = inputs[:, 1:].contiguous()
+            inputs = inputs[:, :-1].contiguous()
+
+            # Zero parameter gradients
+            optimizer.zero_grad()
+
+            # Forward pass
+            outputs, loss = model(inputs, targets)
+
+            # Backward and optimize
+            loss.backward()
+            optimizer.step()
+
+            train_loss.append(loss.item())
+
+            # Print progress every 100 batches
+            if (i + 1) % 100 == 0:
+                print(f'Epoch {it + 1}/{epochs}, Batch {i + 1}/{len(train_loader)}, Loss: {loss.item():.4f}')
+
+        # Get average train_loss
+        train_loss = np.mean(train_loss)
+
+        model.eval()  # Set model to evaluation mode
+        test_loss = []
+        with torch.no_grad():
+            for batch in test_loader:
+                inputs = batch["input_ids"].to(device)
+                attention_mask = batch["attention_mask"].to(device)
+
+                # Create targets by shifting inputs by one position
+                targets = inputs[:, 1:].contiguous()
+                inputs = inputs[:, :-1].contiguous()
+
+                outputs, loss = model(inputs, targets)
+                test_loss.append(loss.item())
+
+            test_loss = np.mean(test_loss)
+
+        # Save losses
+        train_losses[it] = train_loss
+        test_losses[it] = test_loss
+
+        dt = datetime.now() - t0
+        print(f'Epoch {it + 1}/{epochs}, Train Loss: {train_loss:.4f}, '
+              f'Test Loss: {test_loss:.4f}, Duration: {dt}')
+
+    return train_losses, test_losses
+
+# Run the training
+train_losses, test_losses = batch_gh(model, criterion, optimizer, train_loader, test_loader, epochs=2)
+
+# Plot loss
+plt.plot(train_losses, label="train_loss")
+plt.plot(test_losses, label="test_loss")
+plt.legend()
+plt.show()
+
+# Save model weights
+model_save_path = "/home/adrian/Documents/StoryCrafterLLM/model_weights.pth"
+torch.save(model.state_dict(), model_save_path)
+print(f"Model saved to {model_save_path}")