Upload 19 files

checkpoints/model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:231488f17b10270ae927e0c778b1887705bf890ed4d76422ae768c656ed4a44e
+size 43421418

config/model_config.py

@@ -0,0 +1,23 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class ModelConfig:
+    # Model architecture
+    n_embeds: int = 384
+    n_heads: int = 6
+    n_layers: int = 6
+    dropout: float = 0.3
+
+    # Training
+    batch_size: int = 64
+    block_size: int = 128
+    max_iters: int = 20000
+    eval_interval: int = 250
+    eval_iters: int = 200
+    learning_rate: float = 3e-4
+    weight_decay: float = 0.1
+
+    # Paths
+    checkpoint_path: str = "checkpoints/model.pth"
+    data_path: str = "/data/nikhil_workspace/assn11/input.txt"

inference.py

@@ -0,0 +1,55 @@
+import torch
+from config.model_config import ModelConfig
+from src.data.tokenizer import CharacterTokenizer
+from src.utils.helpers import generate, setup_logging
+
+
+def main():
+    # Setup logging
+    logger = setup_logging()
+
+    # Load config
+    config = ModelConfig()
+
+    # Setup device
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    logger.info(f"Using device: {device}")
+
+    # Load tokenizer
+    with open(config.data_path) as f:
+        text = f.read()
+    tokenizer = CharacterTokenizer(text)
+
+    # Load trained model
+    try:
+        model = torch.load(config.checkpoint_path, map_location=device)
+        model.eval()
+    except Exception as e:
+        logger.error(f"Error loading model: {e}")
+        return
+
+    # Generate text from prompts
+    while True:
+        try:
+            prompt = input("\nEnter a prompt (or 'quit' to exit): ")
+            if prompt.lower() == "quit":
+                break
+
+            max_tokens = 200
+
+            logger.info("\nGenerating...")
+            result = generate(model, tokenizer, prompt, max_tokens, device)
+            logger.info("\nGenerated text:")
+            logger.info("=" * 50)
+            logger.info(prompt + result)
+            logger.info("=" * 50)
+        except KeyboardInterrupt:
+            logger.info("\nExiting...")
+            break
+        except Exception as e:
+            logger.error(f"Error during generation: {e}")
+            continue
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,2 @@
+torch>=2.0.0
+gradio>=3.50.0 

src/data/dataset.py

@@ -0,0 +1,39 @@
+import torch
+from torch.utils.data import Dataset, DataLoader
+
+class TextDataset(Dataset):
+    def __init__(self, data, block_size):
+        self.data = data
+        self.block_size = block_size
+    
+    def __len__(self):
+        return len(self.data) - self.block_size
+    
+    def __getitem__(self, idx):
+        x = self.data[idx:idx + self.block_size]
+        y = self.data[idx + 1:idx + self.block_size + 1]
+        return x, y
+
+def create_dataloaders(text, tokenizer, config, device):
+    data = torch.tensor(tokenizer.encode(text), dtype=torch.long)
+    n = int(0.9 * len(data))
+    train_data = data[:n]
+    val_data = data[n:]
+    
+    train_dataset = TextDataset(train_data, config.block_size)
+    val_dataset = TextDataset(val_data, config.block_size)
+    
+    train_loader = DataLoader(
+        train_dataset, 
+        batch_size=config.batch_size, 
+        shuffle=True,
+        pin_memory=True
+    )
+    val_loader = DataLoader(
+        val_dataset, 
+        batch_size=config.batch_size, 
+        shuffle=False,
+        pin_memory=True
+    )
+    
+    return train_loader, val_loader 

src/data/tokenizer.py

@@ -0,0 +1,12 @@
+class CharacterTokenizer:
+    def __init__(self, text):
+        chars = sorted(list(set(text)))
+        self.vocab_size = len(chars)
+        self.stoi = {ch: i for i, ch in enumerate(chars)}
+        self.itos = {i: ch for i, ch in enumerate(chars)}
+    
+    def encode(self, s):
+        return [self.stoi[c] for c in s]
+    
+    def decode(self, l):
+        return "".join([self.itos[i] for i in l]) 

src/model/gpt.py

@@ -0,0 +1,114 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+class GPTModel(nn.Module):
+    def __init__(self, config, vocab_size):
+        super().__init__()
+        self.config = config
+        
+        self.token_embedding = nn.Embedding(vocab_size, config.n_embeds)
+        self.position_embedding = nn.Embedding(config.block_size, config.n_embeds)
+        self.blocks = nn.ModuleList([
+            TransformerBlock(config) for _ in range(config.n_layers)
+        ])
+        self.ln_f = nn.LayerNorm(config.n_embeds)
+        self.lm_head = nn.Linear(config.n_embeds, vocab_size)
+
+    def forward(self, idx, targets=None):
+        B, T = idx.shape
+        
+        tok_emb = self.token_embedding(idx)
+        pos_emb = self.position_embedding(torch.arange(T, device=idx.device))
+        x = tok_emb + pos_emb
+        
+        for block in self.blocks:
+            x = block(x)
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+        
+        if targets is None:
+            loss = None
+        else:
+            B, T, C = logits.shape
+            logits = logits.view(B*T, C)
+            targets = targets.view(B*T)
+            loss = F.cross_entropy(logits, targets)
+            
+        return logits, loss
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(config.n_embeds)
+        self.ln2 = nn.LayerNorm(config.n_embeds)
+        self.attn = MultiHeadAttention(config)
+        self.mlp = FeedForward(config)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x):
+        # Self-attention with residual connection
+        x = x + self.dropout(self.attn(self.ln1(x)))
+        # FFN with residual connection
+        x = x + self.dropout(self.mlp(self.ln2(x)))
+        return x
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.n_heads = config.n_heads
+        self.head_size = config.n_embeds // config.n_heads
+        self.n_embeds = config.n_embeds
+        
+        # Single linear layer for Q, K, V projections
+        self.c_attn = nn.Linear(config.n_embeds, 3 * config.n_embeds)
+        self.c_proj = nn.Linear(config.n_embeds, config.n_embeds)
+        self.dropout = nn.Dropout(config.dropout)
+        
+        # Causal mask to prevent attending to future tokens
+        self.register_buffer(
+            "mask", 
+            torch.tril(torch.ones(config.block_size, config.block_size))
+            .view(1, 1, config.block_size, config.block_size)
+        )
+
+    def forward(self, x):
+        B, T, C = x.shape
+        
+        # Calculate Q, K, V with a single linear projection
+        q, k, v = self.c_attn(x).split(self.n_embeds, dim=2)
+        
+        # Reshape to (B, nh, T, hs)
+        q = q.view(B, T, self.n_heads, self.head_size).transpose(1, 2)
+        k = k.view(B, T, self.n_heads, self.head_size).transpose(1, 2)
+        v = v.view(B, T, self.n_heads, self.head_size).transpose(1, 2)
+
+        # Compute attention scores
+        att = (q @ k.transpose(-2, -1)) * (1.0 / torch.sqrt(torch.tensor(self.head_size)))
+        
+        # Apply causal mask
+        att = att.masked_fill(self.mask[:,:,:T,:T] == 0, float('-inf'))
+        att = F.softmax(att, dim=-1)
+        att = self.dropout(att)
+        
+        # Apply attention to values
+        y = att @ v  # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
+        
+        # Reshape and project back
+        y = y.transpose(1, 2).contiguous().view(B, T, C)  # (B, T, C)
+        y = self.c_proj(y)
+        
+        return y
+
+class FeedForward(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(config.n_embeds, 4 * config.n_embeds),
+            nn.GELU(),
+            nn.Linear(4 * config.n_embeds, config.n_embeds),
+            nn.Dropout(config.dropout),
+        )
+
+    def forward(self, x):
+        return self.net(x) 

src/training/trainer.py

@@ -0,0 +1,61 @@
+import torch
+import logging
+from src.utils.helpers import get_batch
+
+@torch.no_grad()
+def estimate_loss(model, eval_iters, block_size, batch_size, device):
+    out = {}
+    model.eval()
+    for split in ['train', 'val']:
+        losses = torch.zeros(eval_iters)
+        for k in range(eval_iters):
+            xb, yb = get_batch(split, block_size, batch_size)
+            xb, yb = xb.to(device), yb.to(device)
+            logits, loss = model(xb, yb)
+            losses[k] = loss.item()
+        out[split] = losses.mean().item()
+    model.train()
+    return out
+
+def train(
+    model,
+    optimizer,
+    max_iters,
+    eval_interval,
+    eval_iters,
+    block_size,
+    batch_size,
+    device,
+    checkpoint_path="checkpoints/model.pth"
+):
+    logger = logging.getLogger(__name__)
+    best_val_loss = float('inf')
+    
+    for iter in range(max_iters):
+        # Evaluation
+        if iter % eval_interval == 0:
+            losses = estimate_loss(model, eval_iters, block_size, batch_size, device)
+            logger.info(
+                f"Step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}"
+            )
+            
+            # Save best model
+            if losses['val'] < best_val_loss:
+                best_val_loss = losses['val']
+                logger.info(f"Saving model with val loss: {best_val_loss:.4f}")
+                torch.save(model, checkpoint_path)
+
+        # Training step
+        xb, yb = get_batch('train', block_size, batch_size)
+        xb, yb = xb.to(device), yb.to(device)
+
+        # Forward pass
+        logits, loss = model(xb, yb)
+        
+        # Backward pass
+        optimizer.zero_grad(set_to_none=True)
+        loss.backward()
+        optimizer.step()
+
+    # Save final model
+    torch.save(model, checkpoint_path) 

src/utils/helpers.py

@@ -0,0 +1,71 @@
+import torch
+import logging
+import os
+from datetime import datetime
+
+# Global variables for data
+train_data = None
+val_data = None
+
+
+def setup_logging(log_dir="logs"):
+    # Create logs directory if it doesn't exist
+    os.makedirs(log_dir, exist_ok=True)
+
+    # Create a timestamp for the log file
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    log_file = os.path.join(log_dir, f"training_{timestamp}.log")
+
+    # Configure logging
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s",
+        handlers=[
+            logging.FileHandler(log_file),
+            logging.StreamHandler(),  # Also print to console
+        ],
+    )
+
+    logging.info(f"Logging to {log_file}")
+    return logging.getLogger(__name__)
+
+
+def count_parameters(model):
+    return sum(p.numel() for p in model.parameters())
+
+
+def get_batch(split, block_size, batch_size):
+    data = train_data if split == "train" else val_data
+    ix = torch.randint(len(data) - block_size, (batch_size,))
+    x = torch.stack([data[i : i + block_size] for i in ix])
+    y = torch.stack([data[i + 1 : i + block_size + 1] for i in ix])
+    return x, y
+
+
+def prepare_data(text, tokenizer):
+    """Prepare train and validation data"""
+    global train_data, val_data
+
+    # Encode the text
+    data = torch.tensor(tokenizer.encode(text), dtype=torch.long)
+
+    # Split into train and validation sets
+    n = int(0.9 * len(data))
+    train_data = data[:n]
+    val_data = data[n:]
+
+
+def generate(model, tokenizer, prompt, max_tokens, device):
+    model.eval()
+    tokens = torch.tensor(tokenizer.encode(prompt), dtype=torch.long)[None].to(device)
+    block_size = model.config.block_size
+
+    for _ in range(max_tokens):
+        with torch.no_grad():
+            logits, _ = model(tokens[:, -block_size:])
+            logits = logits[:, -1, :]  # / temperature
+            probs = torch.softmax(logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+            tokens = torch.cat([tokens, next_token], dim=1)
+
+    return tokenizer.decode(tokens[0].tolist())[len(prompt) :]

train.py

@@ -0,0 +1,60 @@
+import torch
+from config.model_config import ModelConfig
+from src.data.tokenizer import CharacterTokenizer
+from src.model.gpt import GPTModel
+from src.training.trainer import train
+from src.utils.helpers import generate, setup_logging, prepare_data
+
+
+def main():
+    # Setup logging
+    logger = setup_logging()
+
+    # Load config
+    config = ModelConfig()
+
+    # Setup device
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    logger.info(f"Using device: {device}")
+
+    # Load data
+    with open(config.data_path) as f:
+        text = f.read()
+    tokenizer = CharacterTokenizer(text)
+
+    # Prepare data
+    prepare_data(text, tokenizer)
+
+    # Create model
+    model = GPTModel(config, tokenizer.vocab_size)
+    model = model.to(device)
+
+    # Setup optimizer
+    optimizer = torch.optim.AdamW(
+        model.parameters(), lr=config.learning_rate, weight_decay=config.weight_decay
+    )
+
+    # Train
+    train(
+        model=model,
+        optimizer=optimizer,
+        max_iters=config.max_iters,
+        eval_interval=config.eval_interval,
+        eval_iters=config.eval_iters,
+        block_size=config.block_size,
+        batch_size=config.batch_size,
+        device=device,
+        checkpoint_path=config.checkpoint_path,
+    )
+
+    # Generate samples
+    model = torch.load(config.checkpoint_path, map_location=device)
+    for prompt in ["hello", "my name is", "america is"]:
+        result = generate(model, tokenizer, prompt, max_tokens=200, device=device)
+        logger.info(f"\nPrompt: {prompt}")
+        logger.info(f"Generated: {result}")
+        logger.info("=" * 40)
+
+
+if __name__ == "__main__":
+    main()