Add Story GPT Python Space

README.md

@@ -1,12 +1,28 @@
 ---
-title: Story Gpt Python
-emoji: 🐢
-colorFrom: pink
+title: Story GPT Python
+colorFrom: yellow
 colorTo: red
 sdk: gradio
-sdk_version: 6.10.0
 app_file: app.py
 pinned: false
+license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Story GPT Python
+
+This is a tiny story-writing GPT-style language model project written in Python from scratch.
+
+## What it includes
+
+- Word-level tokenizer
+- Causal transformer decoder with self-attention
+- Story-focused local training corpus
+- Local CPU training loop
+- Checkpoint save and load
+- Gradio user interface
+
+## Important
+
+- No external pretrained LLM is used
+- This is a small educational GPT-like story model
+- The first generate or train call will initialize and train the model locally

app.py

@@ -0,0 +1,118 @@
+import gradio as gr
+
+from story_gpt.config import StoryGPTConfig
+from story_gpt.service import StoryGPTService
+
+
+config = StoryGPTConfig()
+service = StoryGPTService(config=config)
+
+
+def generate_story(title, genre, tone, idea, opening_line, max_new_tokens, temperature, top_k):
+    return service.generate_story(
+        title=title,
+        genre=genre,
+        tone=tone,
+        idea=idea,
+        opening_line=opening_line,
+        max_new_tokens=int(max_new_tokens),
+        temperature=float(temperature),
+        top_k=int(top_k),
+    )
+
+
+def train_story_model(extra_story_text, steps):
+    return service.train(extra_story_text=extra_story_text, steps=int(steps))
+
+
+def reset_story_model():
+    return service.reset()
+
+
+with gr.Blocks(
+    title="Story GPT Python",
+    theme=gr.themes.Soft(primary_hue="amber", secondary_hue="orange"),
+) as demo:
+    gr.Markdown(
+        """
+        # Story GPT Python
+        A tiny story-writing GPT-style model written in Python from scratch.
+
+        - Causal transformer decoder
+        - Word-level tokenizer
+        - Story-focused local training corpus
+        - No external pretrained LLM
+        """
+    )
+
+    with gr.Tab("Write Story"):
+        with gr.Row():
+            title_input = gr.Textbox(label="Title", value="The Lantern in the Rain")
+            genre_input = gr.Dropdown(
+                label="Genre",
+                choices=["Fantasy", "Adventure", "Mystery", "Sci-Fi", "Friendship", "Folktale"],
+                value="Fantasy",
+            )
+            tone_input = gr.Dropdown(
+                label="Tone",
+                choices=["Warm", "Wonder", "Suspense", "Playful", "Calm", "Heroic"],
+                value="Wonder",
+            )
+
+        idea_input = gr.Textbox(
+            label="Story Idea",
+            value="A child finds a glowing lantern that reveals hidden paths after a storm.",
+            lines=5,
+        )
+        opening_line_input = gr.Textbox(
+            label="Opening Line",
+            value="When the rain stopped, the alley behind Mira's house began to shine.",
+            lines=2,
+        )
+
+        with gr.Row():
+            max_tokens_input = gr.Slider(30, 220, value=110, step=5, label="Story Length")
+            temperature_input = gr.Slider(0.2, 1.4, value=0.85, step=0.05, label="Temperature")
+            top_k_input = gr.Slider(1, 24, value=10, step=1, label="Top-K")
+
+        generate_button = gr.Button("Generate Story", variant="primary")
+        output_text = gr.Textbox(label="Story Output", lines=14)
+        output_status = gr.Textbox(label="Status", lines=4)
+
+    with gr.Tab("Train"):
+        extra_story_text_input = gr.Textbox(
+            label="Extra Story Examples",
+            placeholder="Add more short stories, story prompts, or endings to continue training the model.",
+            lines=12,
+        )
+        steps_input = gr.Slider(10, 500, value=140, step=10, label="Training Steps")
+        train_button = gr.Button("Train Story Model", variant="primary")
+        reset_button = gr.Button("Reset Model")
+        train_status = gr.Textbox(label="Training Status", lines=6)
+
+    generate_button.click(
+        fn=generate_story,
+        inputs=[
+            title_input,
+            genre_input,
+            tone_input,
+            idea_input,
+            opening_line_input,
+            max_tokens_input,
+            temperature_input,
+            top_k_input,
+        ],
+        outputs=[output_text, output_status],
+    )
+
+    train_button.click(
+        fn=train_story_model,
+        inputs=[extra_story_text_input, steps_input],
+        outputs=[train_status],
+    )
+
+    reset_button.click(fn=reset_story_model, outputs=[train_status])
+
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,2 @@
+gradio>=5.23.0
+torch>=2.3.0

story_gpt/__init__.py

@@ -0,0 +1,4 @@
+from .config import StoryGPTConfig
+from .service import StoryGPTService
+
+__all__ = ["StoryGPTConfig", "StoryGPTService"]

story_gpt/config.py

@@ -0,0 +1,24 @@
+from dataclasses import dataclass
+from pathlib import Path
+
+
+@dataclass
+class StoryGPTConfig:
+    block_size: int = 72
+    batch_size: int = 16
+    d_model: int = 112
+    n_heads: int = 4
+    n_layers: int = 4
+    dropout: float = 0.1
+    learning_rate: float = 2.0e-3
+    bootstrap_steps: int = 120
+    cpu_threads: int = 4
+    seed: int = 42
+
+    @property
+    def root_dir(self) -> Path:
+        return Path(__file__).resolve().parents[1]
+
+    @property
+    def checkpoint_path(self) -> Path:
+        return self.root_dir / "artifacts" / "story_gpt_checkpoint.pt"

story_gpt/data.py

@@ -0,0 +1,81 @@
+BASE_CORPUS = """
+Instruction: Write a warm fantasy story.
+Title: The Lantern in the Rain
+Genre: Fantasy
+Tone: Warm
+Idea: A child finds a glowing lantern after a storm.
+Opening: When the rain stopped, the alley behind Mira's house began to shine.
+Story: When the rain stopped, the alley behind Mira's house began to shine. A brass lantern waited beside the last puddle, glowing with a quiet gold light. Mira lifted it, and the beam painted a path through walls the way moonlight slips through leaves. The hidden path led to a garden growing on rooftops, where dripping vines carried silver pears and tiny birds sang in voices soft as bells. An old gardener bowed and told Mira that the lantern appears only to people who still notice wonder after thunder. She picked one silver pear, carried the lantern home, and hung it by her window. Every storm after that left her less afraid, because she knew the dark was only the place where secret gardens begin.
+
+Instruction: Write a playful friendship story.
+Title: The Bicycle with Two Bells
+Genre: Friendship
+Tone: Playful
+Idea: Two friends fix a broken bicycle to win a town parade.
+Opening: Ravi found the bicycle under a blanket of marigold petals.
+Story: Ravi found the bicycle under a blanket of marigold petals in his grandfather's shed. The front wheel leaned sideways, one pedal was missing, and two bright brass bells dangled from the handlebars. He called for Sana, who arrived with a box of screws, purple paint, and opinions louder than thunder. They spent the afternoon tightening bolts, polishing spokes, and arguing about whether a bicycle should sparkle. At sunset they rode into the parade with both bells ringing at once, startling pigeons and making children laugh. They did not win the biggest trophy, but the mayor gave them the prize for happiest invention. Ravi said the bicycle rode better because of the bells. Sana said it rode better because they fixed it together.
+
+Instruction: Write a suspenseful mystery story.
+Title: The Library Key
+Genre: Mystery
+Tone: Suspense
+Idea: A student discovers a key hidden inside an old atlas.
+Opening: The atlas fell open by itself on the coldest evening of winter.
+Story: The atlas fell open by itself on the coldest evening of winter, dropping a brass key onto Noor's desk. She recognized the number tied to it with blue thread: nineteen, the same number painted on the locked door behind the school library staircase. After everyone left, Noor followed the silent corridor and slipped the key into the iron lock. Inside the room she found shelves of letters never mailed, each one written by students who had once been too afraid to speak. At the center stood a journal kept by the retired librarian, who had hidden the room so secrets could become courage when the time was right. Noor spent all night reading promises, apologies, and impossible dreams. In the morning she placed one fresh letter on the desk outside. It read, The door is open now.
+
+Instruction: Write a gentle folktale.
+Title: The River That Remembered
+Genre: Folktale
+Tone: Calm
+Idea: A village forgets its songs until a quiet child listens to the river.
+Opening: Every morning the river moved past the village like folded blue silk.
+Story: Every morning the river moved past the village like folded blue silk, but no one noticed its music anymore. The bakers had forgotten the kneading song, the weavers had forgotten the shuttle rhyme, and even weddings had grown quiet. Only little Tara sat beside the bank and listened with her feet in the water. She heard a melody hidden in the current, a tune that sounded like grandmothers laughing inside clay walls. Tara carried the melody home and hummed it while sweeping the courtyard. Her mother remembered the next line. The baker remembered the drum beat. By evening the entire village was singing to one another across their rooftops. From then on, they said the river remembers for those who are too busy to hear.
+
+Instruction: Write a bright science fiction story.
+Title: The Pocket Star Map
+Genre: Sci-Fi
+Tone: Wonder
+Idea: A mechanic builds a map that points toward lost satellites.
+Opening: Arin's workshop smelled like copper, dust, and hot tea.
+Story: Arin's workshop smelled like copper, dust, and hot tea when the tiny projector finally woke in his palm. It scattered blue light over the ceiling, shaping a star map that shifted each time he whispered a satellite name. One marker blinked above the desert, where an old weather station had vanished years ago. Arin packed his tools, followed the projector through the dunes, and found the missing machine half-buried beneath moonlit sand. When he repaired its cracked panel, the station answered with stored messages from children who had once used it to learn constellations. Arin transmitted the lessons back to the city that same night. By morning, every rooftop telescope was turned upward. The mayor called it a scientific recovery. Arin called it returning a lost piece of the sky.
+
+Instruction: Write a heroic adventure story.
+Title: The Bridge of Kites
+Genre: Adventure
+Tone: Heroic
+Idea: A village must cross a broken bridge during a wind festival.
+Opening: On the morning of the wind festival, the bridge ropes snapped.
+Story: On the morning of the wind festival, the bridge ropes snapped, leaving the mountain village stranded above a roaring gorge. Mehul looked at the kites lined up for the celebration and saw not decorations, but sails. He tied bamboo poles to the strongest frames, stitched the kite cloth into long panels, and asked every household for spare string. Under a hard blue sky the villagers worked as fast as prayer. By noon a ribbon bridge of red, gold, and indigo stretched from one cliff to the other, trembling but strong. Children crossed first with medicine for the far side, then farmers, then drums, then dancers. When the evening wind rose, the bridge sang above the gorge like a hundred bright birds. The festival carried on, and every kite overhead looked like a flag for courage.
+
+Instruction: Write a bedtime story.
+Title: The Moon Rabbit's Blanket
+Genre: Fantasy
+Tone: Calm
+Idea: A moon rabbit sews dreams into a blanket for a sleepless child.
+Opening: Above the sleeping city, one window was still lit.
+Story: Above the sleeping city, one window was still lit, and from the moon a rabbit noticed at once. The little rabbit gathered threads from cloud edges, silver dust from quiet stars, and a single feather from the night wind. It hopped down a beam of moonlight and sat on the child's windowsill, sewing without a sound. Into one square it stitched a forest with glowing fireflies. Into another it stitched a lake so still it could hold a whole sky. When the blanket touched the child, the room filled with the smell of rain and jasmine, and the tired eyes finally closed. By dawn the rabbit was back on the moon, but one silver thread remained on the pillow. The child kept it for years and slept as if the sky itself had tucked them in.
+""".strip()
+
+
+def build_training_text(extra_text: str = "") -> str:
+    extra = (extra_text or "").strip()
+    if not extra:
+        return BASE_CORPUS
+    return BASE_CORPUS + "\n\n" + extra
+
+
+def build_story_prompt(title: str, genre: str, tone: str, idea: str, opening_line: str) -> str:
+    clean_title = (title or "A New Story").strip()
+    clean_genre = (genre or "Fantasy").strip()
+    clean_tone = (tone or "Wonder").strip()
+    clean_idea = (idea or "Write a short imaginative story.").strip()
+    clean_opening = (opening_line or "Begin with a vivid first line.").strip()
+    return (
+        f"Instruction: Write a {clean_tone.lower()} {clean_genre.lower()} story.\n"
+        f"Title: {clean_title}\n"
+        f"Genre: {clean_genre}\n"
+        f"Tone: {clean_tone}\n"
+        f"Idea: {clean_idea}\n"
+        f"Opening: {clean_opening}\n"
+        f"Story: {clean_opening}"
+    )

story_gpt/model.py

@@ -0,0 +1,111 @@
+import math
+
+import torch
+from torch import nn
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, d_model, n_heads, block_size, dropout):
+        super().__init__()
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.qkv = nn.Linear(d_model, 3 * d_model)
+        self.out_proj = nn.Linear(d_model, d_model)
+        self.dropout = nn.Dropout(dropout)
+        mask = torch.tril(torch.ones(block_size, block_size)).view(1, 1, block_size, block_size)
+        self.register_buffer("mask", mask)
+
+    def forward(self, x):
+        batch, seq_len, channels = x.shape
+        qkv = self.qkv(x)
+        q, k, v = qkv.chunk(3, dim=-1)
+
+        q = q.view(batch, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+        k = k.view(batch, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+        v = v.view(batch, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+
+        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        att = att.masked_fill(self.mask[:, :, :seq_len, :seq_len] == 0, float("-inf"))
+        att = torch.softmax(att, dim=-1)
+        att = self.dropout(att)
+
+        out = att @ v
+        out = out.transpose(1, 2).contiguous().view(batch, seq_len, channels)
+        return self.out_proj(out)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, d_model, dropout):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(d_model, 4 * d_model),
+            nn.GELU(),
+            nn.Linear(4 * d_model, d_model),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class GPTBlock(nn.Module):
+    def __init__(self, d_model, n_heads, block_size, dropout):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(d_model)
+        self.attn = CausalSelfAttention(d_model, n_heads, block_size, dropout)
+        self.ln2 = nn.LayerNorm(d_model)
+        self.ff = FeedForward(d_model, dropout)
+
+    def forward(self, x):
+        x = x + self.attn(self.ln1(x))
+        x = x + self.ff(self.ln2(x))
+        return x
+
+
+class StoryGPTModel(nn.Module):
+    def __init__(self, vocab_size, block_size, d_model, n_heads, n_layers, dropout):
+        super().__init__()
+        self.block_size = block_size
+        self.token_emb = nn.Embedding(vocab_size, d_model)
+        self.pos_emb = nn.Embedding(block_size, d_model)
+        self.dropout = nn.Dropout(dropout)
+        self.blocks = nn.Sequential(
+            *[GPTBlock(d_model, n_heads, block_size, dropout) for _ in range(n_layers)]
+        )
+        self.ln_f = nn.LayerNorm(d_model)
+        self.head = nn.Linear(d_model, vocab_size, bias=False)
+        self.head.weight = self.token_emb.weight
+
+    def forward(self, idx, targets=None):
+        batch, seq_len = idx.shape
+        positions = torch.arange(seq_len, device=idx.device)
+        x = self.token_emb(idx) + self.pos_emb(positions)[None, :, :]
+        x = self.dropout(x)
+        x = self.blocks(x)
+        x = self.ln_f(x)
+        logits = self.head(x)
+
+        loss = None
+        if targets is not None:
+            loss = nn.functional.cross_entropy(
+                logits.reshape(-1, logits.size(-1)),
+                targets.reshape(-1),
+            )
+        return logits, loss
+
+    def generate(self, idx, max_new_tokens, eos_id, temperature=1.0, top_k=8):
+        for _ in range(max_new_tokens):
+            idx_cond = idx[:, -self.block_size :]
+            logits, _ = self(idx_cond)
+            logits = logits[:, -1, :] / max(temperature, 1e-4)
+
+            if top_k is not None and top_k > 0:
+                values, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < values[:, [-1]]] = float("-inf")
+
+            probs = torch.softmax(logits, dim=-1)
+            next_id = torch.multinomial(probs, num_samples=1)
+            idx = torch.cat([idx, next_id], dim=1)
+            if int(next_id.item()) == eos_id:
+                break
+        return idx

story_gpt/service.py

@@ -0,0 +1,139 @@
+import shutil
+
+import torch
+
+from .config import StoryGPTConfig
+from .data import build_story_prompt
+from .model import StoryGPTModel
+from .tokenizer import WordTokenizer
+from .trainer import create_model_and_tokenizer, set_seed, train_model
+
+
+class StoryGPTService:
+    def __init__(self, config: StoryGPTConfig):
+        self.config = config
+        torch.set_num_threads(max(1, self.config.cpu_threads))
+        self.model = None
+        self.tokenizer = None
+
+    def generate_story(
+        self,
+        title: str,
+        genre: str,
+        tone: str,
+        idea: str,
+        opening_line: str,
+        max_new_tokens: int,
+        temperature: float,
+        top_k: int,
+    ):
+        clean_prompt = build_story_prompt(
+            title=title,
+            genre=genre,
+            tone=tone,
+            idea=idea,
+            opening_line=opening_line,
+        )
+        self._ensure_ready()
+        encoded = self.tokenizer.encode(clean_prompt, add_bos=True)
+        idx = torch.tensor(encoded, dtype=torch.long).unsqueeze(0)
+        self.model.eval()
+
+        with torch.inference_mode():
+            output = self.model.generate(
+                idx=idx,
+                max_new_tokens=max_new_tokens,
+                eos_id=self.tokenizer.eos_id,
+                temperature=temperature,
+                top_k=top_k,
+            )
+
+        text = self.tokenizer.decode(output[0].tolist())
+        if "Story:" in text:
+            text = text.split("Story:", 1)[1].strip()
+        status = (
+            f"Generated with Story GPT Python. "
+            f"Architecture=causal transformer, Vocab={self.tokenizer.vocab_size}, Layers={self.config.n_layers}."
+        )
+        return text, status
+
+    def train(self, extra_story_text: str, steps: int):
+        steps = max(1, steps)
+        checkpoint_exists = self.config.checkpoint_path.exists()
+        training_text = extra_story_text or ""
+
+        if checkpoint_exists:
+            self._load_or_initialize(extra_text="")
+
+        model, tokenizer, encoded = create_model_and_tokenizer(self.config, training_text)
+        if checkpoint_exists and self.model is not None and self.tokenizer is not None:
+            if tokenizer.stoi == self.tokenizer.stoi:
+                model.load_state_dict(self.model.state_dict())
+
+        losses = train_model(model, encoded, self.config, steps)
+        self.model = model
+        self.tokenizer = tokenizer
+        self._save_checkpoint(extra_text=training_text)
+
+        return (
+            f"Story GPT training finished.\n"
+            f"Steps: {steps}\n"
+            f"Start Loss: {losses[0]:.4f}\n"
+            f"End Loss: {losses[-1]:.4f}\n"
+            f"Checkpoint: {self.config.checkpoint_path}"
+        )
+
+    def reset(self):
+        checkpoint_dir = self.config.checkpoint_path.parent
+        if checkpoint_dir.exists():
+            shutil.rmtree(checkpoint_dir)
+        self.model = None
+        self.tokenizer = None
+        return "Story GPT reset complete. Next train or generate call will rebuild from scratch."
+
+    def _ensure_ready(self):
+        if self.model is not None and self.tokenizer is not None:
+            return
+        self._load_or_initialize(extra_text="")
+
+    def _load_or_initialize(self, extra_text: str):
+        checkpoint = self.config.checkpoint_path
+        if checkpoint.exists():
+            state = torch.load(checkpoint, map_location="cpu")
+            self.tokenizer = WordTokenizer.from_state_dict(state["tokenizer"])
+            self.model = StoryGPTModel(
+                vocab_size=state["config"]["vocab_size"],
+                block_size=state["config"]["block_size"],
+                d_model=state["config"]["d_model"],
+                n_heads=state["config"]["n_heads"],
+                n_layers=state["config"]["n_layers"],
+                dropout=state["config"]["dropout"],
+            )
+            self.model.load_state_dict(state["model"])
+            self.model.eval()
+            return
+
+        set_seed(self.config.seed)
+        self.model, self.tokenizer, encoded = create_model_and_tokenizer(self.config, extra_text)
+        train_model(self.model, encoded, self.config, self.config.bootstrap_steps)
+        self._save_checkpoint(extra_text=extra_text)
+
+    def _save_checkpoint(self, extra_text: str):
+        checkpoint = self.config.checkpoint_path
+        checkpoint.parent.mkdir(parents=True, exist_ok=True)
+        torch.save(
+            {
+                "model": self.model.state_dict(),
+                "tokenizer": self.tokenizer.state_dict(),
+                "config": {
+                    "vocab_size": self.tokenizer.vocab_size,
+                    "block_size": self.config.block_size,
+                    "d_model": self.config.d_model,
+                    "n_heads": self.config.n_heads,
+                    "n_layers": self.config.n_layers,
+                    "dropout": self.config.dropout,
+                    "extra_text": extra_text,
+                },
+            },
+            checkpoint,
+        )

story_gpt/tokenizer.py

@@ -0,0 +1,69 @@
+import re
+
+
+TOKEN_PATTERN = re.compile(r"\n|[A-Za-z0-9_']+|[^\w\s]")
+
+
+class WordTokenizer:
+    def __init__(self):
+        self.special_tokens = ["<pad>", "<unk>", "<bos>", "<eos>"]
+        self.stoi = {}
+        self.itos = {}
+
+    @property
+    def vocab_size(self):
+        return len(self.stoi)
+
+    @property
+    def bos_id(self):
+        return self.stoi["<bos>"]
+
+    @property
+    def eos_id(self):
+        return self.stoi["<eos>"]
+
+    def tokenize(self, text: str):
+        return TOKEN_PATTERN.findall(text)
+
+    def fit(self, text: str):
+        vocab = self.special_tokens + sorted(set(self.tokenize(text)))
+        self.stoi = {token: idx for idx, token in enumerate(vocab)}
+        self.itos = {idx: token for token, idx in self.stoi.items()}
+        return self
+
+    def encode(self, text: str, add_bos: bool = False, add_eos: bool = False):
+        tokens = self.tokenize(text)
+        ids = [self.stoi.get(token, self.stoi["<unk>"]) for token in tokens]
+        if add_bos:
+            ids = [self.bos_id] + ids
+        if add_eos:
+            ids = ids + [self.eos_id]
+        return ids
+
+    def decode(self, ids):
+        tokens = []
+        for idx in ids:
+            token = self.itos.get(int(idx), "<unk>")
+            if token in self.special_tokens:
+                continue
+            tokens.append(token)
+
+        text = ""
+        for token in tokens:
+            if token == "\n":
+                text = text.rstrip() + "\n"
+            elif token in {".", ",", "!", "?", ":", ";"}:
+                text = text.rstrip() + token + " "
+            else:
+                text += token + " "
+        return text.strip()
+
+    def state_dict(self):
+        return {"stoi": self.stoi}
+
+    @classmethod
+    def from_state_dict(cls, state):
+        tok = cls()
+        tok.stoi = dict(state["stoi"])
+        tok.itos = {idx: token for token, idx in tok.stoi.items()}
+        return tok

story_gpt/trainer.py

@@ -0,0 +1,52 @@
+import random
+
+import torch
+
+from .data import build_training_text
+from .model import StoryGPTModel
+from .tokenizer import WordTokenizer
+
+
+def set_seed(seed: int):
+    random.seed(seed)
+    torch.manual_seed(seed)
+
+
+def create_model_and_tokenizer(config, extra_text=""):
+    text = build_training_text(extra_text)
+    tokenizer = WordTokenizer().fit(text)
+    encoded = tokenizer.encode(text, add_bos=True, add_eos=True)
+    encoded = torch.tensor(encoded, dtype=torch.long)
+    model = StoryGPTModel(
+        vocab_size=tokenizer.vocab_size,
+        block_size=config.block_size,
+        d_model=config.d_model,
+        n_heads=config.n_heads,
+        n_layers=config.n_layers,
+        dropout=config.dropout,
+    )
+    return model, tokenizer, encoded
+
+
+def build_batch(encoded, block_size, batch_size):
+    max_start = max(1, len(encoded) - block_size - 1)
+    starts = torch.randint(0, max_start, (batch_size,))
+    x = torch.stack([encoded[start : start + block_size] for start in starts])
+    y = torch.stack([encoded[start + 1 : start + block_size + 1] for start in starts])
+    return x, y
+
+
+def train_model(model, encoded, config, steps):
+    optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate)
+    model.train()
+    losses = []
+
+    for _ in range(steps):
+        xb, yb = build_batch(encoded, config.block_size, config.batch_size)
+        _, loss = model(xb, targets=yb)
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        losses.append(float(loss.item()))
+
+    return losses