Add model architecture definitions

MiniLLM.py

@@ -0,0 +1,279 @@
+import jax
+import jax.numpy as jnp
+import flax.nnx as nnx
+import grain.python as pygrain
+import optax
+import tiktoken
+from pathlib import Path
+
+tokenizer = tiktoken.get_encoding("gpt2")
+
+vocab_size = tokenizer.n_vocab
+num_transformer_blocks = 6
+maxlen = 128
+embed_dim = 192
+num_heads = 6
+feed_forward_dim = int(2/3 * 4 * embed_dim)
+batch_size = 24
+num_epochs = 3
+
+def load_stories_from_file(
+    file_path,
+    max_stories = None
+):
+    """
+    Efficiently load stories from a text file.
+    Each story ends with <|endoftext|>.
+    """
+    file_path = Path(file_path)
+    if not file_path.exists():
+        raise FileNotFoundError(f"Data file not found: {file_path}")
+
+    print(f"Loading stories from {file_path}...")
+    stories = []
+    current_story = []
+
+    with open(file_path, 'r', encoding='utf-8', errors='replace') as f:
+        for line in f:
+            if '<|endoftext|>' in line:
+                parts = line.split('<|endoftext|>')
+                for part in parts[:-1]:
+                    current_story.append(part)
+                    story_text = ''.join(current_story).strip()
+                    if story_text:
+                        stories.append(story_text + '<|endoftext|>')
+                        if max_stories and len(stories) >= max_stories:
+                            break
+                    current_story = []
+                if parts[-1].strip():
+                    current_story = [parts[-1]]
+                else:
+                    current_story = []
+                if max_stories and len(stories) >= max_stories:
+                    break
+            else:
+                current_story.append(line)
+        if current_story and (not max_stories or len(stories) < max_stories):
+            story_text = ''.join(current_story).strip()
+            if story_text:
+                stories.append(story_text + '<|endoftext|>')
+
+    print(f"Loaded {len(stories):,} stories")
+    return stories
+
+
+
+class TransformerBlock(nnx.Module):
+
+    def __init__(self, embed_dim, num_heads, ff_dim, *, rngs):
+
+        self.attention = nnx.MultiHeadAttention(
+            num_heads=num_heads,
+            in_features=embed_dim,
+            qkv_features=embed_dim,
+            out_features=embed_dim,
+            decode=False,
+            rngs=rngs
+        )
+
+    def __call__(self, x, mask=None):
+        attn_out = self.attention(x, mask=mask)
+        x = x + attn_out
+        return x
+
+class TokenAndPositionEmbedding(nnx.Module):
+    def __init__(self, maxlen, vocab_size, embed_dim, *, rngs):
+        self.token_emb = nnx.Embed(vocab_size, embed_dim, rngs=rngs)
+        self.pos_emb = nnx.Embed(maxlen, embed_dim, rngs=rngs)
+
+    def __call__(self, x):
+        seq_len = x.shape[1]
+        positions = jnp.arange(seq_len)[None, :]
+        return self.token_emb(x) + self.pos_emb(positions)
+
+class MiniGPT(nnx.Module):
+
+    def __init__(self, maxlen=maxlen, vocab_size=vocab_size, embed_dim=embed_dim, num_heads=num_heads,
+                 feed_forward_dim=feed_forward_dim, num_transformer_blocks=num_transformer_blocks, *, rngs=nnx.Rngs(0)):
+
+        self.maxlen = maxlen
+
+        self.embedding = TokenAndPositionEmbedding(maxlen, vocab_size, embed_dim, rngs=rngs)
+
+        self.transformer_blocks = [
+            TransformerBlock(embed_dim, num_heads, feed_forward_dim, rngs=rngs)
+            for _ in range(num_transformer_blocks)
+        ]
+
+        self.output_layer = nnx.Linear(embed_dim, vocab_size, use_bias=False, rngs=rngs)
+
+    def causal_attention_mask(self, seq_len):
+        return jnp.tril(jnp.ones((seq_len, seq_len)))
+
+    def __call__(self, token_ids):
+        seq_len = token_ids.shape[1]
+        mask = self.causal_attention_mask(seq_len)
+
+        x = self.embedding(token_ids)
+
+        for block in self.transformer_blocks:
+            x = block(x, mask=mask)
+
+        logits = self.output_layer(x)
+
+        return logits
+
+    # Pass the tokenizer as an argument to break the global dependency
+    def generate_text(self, prompt, tokenizer, max_new_tokens=50, temperature=1.0):
+        # 1. Encode the string prompt into integer token IDs
+        tokens = tokenizer.encode(prompt)
+
+        # Cache the end token ID so we aren't encoding it on every single loop iteration
+        end_token_id = tokenizer.encode('<|endoftext|>', allowed_special={'<|endoftext|>'})[0]
+
+        for _ in range(max_new_tokens):
+            context = tokens[-self.maxlen:]
+
+            # RIGHT-pad to match training (not left-pad!)
+            actual_len = len(context)
+            if actual_len < self.maxlen:
+                context = context + [0] * (self.maxlen - actual_len)
+
+            context_array = jnp.array(context)[None, :]
+            logits = self(context_array)
+
+            next_token_logits = logits[0, actual_len - 1, :] / temperature
+            next_token = int(jnp.argmax(next_token_logits))
+
+            if next_token == end_token_id:
+                break
+
+            tokens.append(next_token)
+
+        # 2. Decode the final list of token IDs back into a string
+        return tokenizer.decode(tokens)
+
+
+
+def generate_story(model, story_prompt, temperature, max_new_tokens):
+    start_tokens = tokenizer.encode(story_prompt)[:maxlen]
+    generated = generate_text(model, start_tokens, max_new_tokens=max_new_tokens, temperature=temperature)
+    return generated
+
+
+class StoryDataset:
+    def __init__(self, stories, maxlen, tokenizer):
+        self.stories = stories
+        self.maxlen = maxlen
+        self.tokenizer = tokenizer
+        self.end_token = tokenizer.encode('<|endoftext|>', allowed_special={'<|endoftext|>'})[0]
+
+    def __len__(self):
+        return len(self.stories)
+
+    def __getitem__(self, idx):
+        story = self.stories[idx]
+        tokens = self.tokenizer.encode(story, allowed_special={'<|endoftext|>'})
+
+        if len(tokens) > self.maxlen:
+            tokens = tokens[:self.maxlen]
+
+        tokens.extend([0] * (self.maxlen - len(tokens)))
+        return tokens
+
+
+
+def load_and_preprocess_data(
+    file_path,
+    batch_size,
+    maxlen,
+    max_stories = 100_000,
+    num_epochs = 1,
+    shuffle = False,
+    seed = 42
+):
+    """
+    Load and preprocess TinyStories data with memory-efficient chunk reading.
+
+    Args:
+        file_path: Path to the text file
+        batch_size: Batch size for training
+        maxlen: Maximum sequence length
+        max_stories: Maximum number of stories to load (for memory efficiency)
+        num_epochs: Number of training epochs
+        shuffle: Whether to shuffle the data
+        seed: Random seed for reproducibility
+
+    Returns:
+        Tuple of (Grain DataLoader, estimated_batches_per_epoch)
+    """
+
+    # Load and validate file
+    file_path = file_path
+
+    print(f"Loading data from {file_path} (max {max_stories:,} stories)")
+
+    # Read file in chunks to avoid loading entire file into memory
+    stories = []
+    current_story = []
+
+    with open(file_path, 'r', encoding='utf-8', errors='replace') as f:
+        for line in f:
+            if '<|endoftext|>' in line:
+                # Split on end token and process parts
+                parts = line.split('<|endoftext|>')
+                for i, part in enumerate(parts[:-1]):  # All but last part have end tokens
+                    current_story.append(part)
+                    story_text = ''.join(current_story).strip()
+                    if story_text:
+                        stories.append(story_text + '<|endoftext|>')
+                        if len(stories) >= max_stories:
+                            break
+                    current_story = []
+
+                # Last part becomes start of next story
+                if parts[-1].strip():
+                    current_story = [parts[-1]]
+
+                if len(stories) >= max_stories:
+                    break
+            else:
+                current_story.append(line)
+
+        # Don't forget the last story if file doesn't end with end token
+        if current_story and len(stories) < max_stories:
+            story_text = ''.join(current_story).strip()
+            if story_text:
+                stories.append(story_text + '<|endoftext|>')
+
+    print(f"Loaded {len(stories):,} stories")
+    if len(stories) == 0:
+        raise ValueError("No valid stories found in the dataset")
+
+    # Calculate estimated batches per epoch
+    estimated_batches_per_epoch = len(stories) // batch_size
+    print(f"Estimated batches per epoch: {estimated_batches_per_epoch:,}")
+
+    # Create efficient dataset
+    dataset = StoryDataset(stories, maxlen, tokenizer)
+
+    # Configure sampler with sharding support
+    sampler = pygrain.IndexSampler(
+        num_records=len(dataset),
+        shuffle=shuffle,
+        seed=seed,
+        shard_options=pygrain.NoSharding(),
+        num_epochs=num_epochs,
+    )
+
+    # Create DataLoader with efficient batching
+    dataloader = pygrain.DataLoader(
+        data_source=dataset,
+        sampler=sampler,
+        operations=[
+            pygrain.Batch(batch_size=batch_size, drop_remainder=True)
+        ]
+    )
+
+    print(f"Created DataLoader with batch_size={batch_size}, maxlen={maxlen}")
+    return dataloader, estimated_batches_per_epoch