Upload README.md with huggingface_hub

README.md

@@ -0,0 +1,271 @@
+---
+license: mit
+language:
+  - en
+library_name: pytorch
+tags:
+  - pytorch
+  - nanogpt
+  - language-model
+  - from-scratch
+  - small-language-model
+  - tinystories
+  - story-generation
+  - childrens-stories
+  - text-generation
+  - rlhf
+  - rlvr
+  - reinforcement-learning
+  - policy-gradient
+  - sft
+  - sentiment
+datasets:
+  - roneneldan/TinyStories
+pipeline_tag: text-generation
+widget:
+  - text: "Once upon a time there was a little rabbit"
+---
+
+# nanoGPT SLM -- Cheerful TinyStories (3-Stage Pipeline: Pretrain -> SFT -> RLVR)
+
+A **124M-parameter nanoGPT (GPT-2 small)** language model trained **entirely from scratch** on
+the **TinyStories** dataset, then aligned to write **consistently cheerful, positive
+children's stories** through a 3-stage RLHF-style pipeline:
+
+**Pretraining -> Supervised Fine-Tuning (SFT) -> Reinforcement Learning with Verifiable Rewards (RLVR).**
+
+This repository ships **all three checkpoints** so you can load and compare every stage of
+the pipeline yourself.
+
+## What This Model Does
+
+The headline model (**RLVR**) generates short, age-appropriate children's stories (ages 3-5)
+that are reliably **warm, upbeat, and resolve happily**. Give it a story opening and it
+continues in simple, cheerful language:
+
+```
+Input:  "The little girl was sad until"
+Output: "The little girl was sad until she found a tiny puppy in the garden.
+         The puppy wagged its tail and licked her hand. She laughed and hugged
+         it close. They played together all afternoon and became best friends."
+```
+
+## The 3-Stage Pipeline
+
+| Stage | Checkpoint | What it does | How |
+|:--|:--|:--|:--|
+| **1. Pretraining** | `nanogpt_slm_tinystories_best.pth` | Learns general next-token competence on TinyStories | 70k iterations, AdamW, cosine LR |
+| **2. SFT** | `nanogpt_slm_sft_best.pth` | Shifts the *prior* toward positive stories | Next-token training on a VADER-filtered positive subset (low LR) |
+| **3. RLVR** | `nanogpt_slm_rlvr_final.pth` | *Optimizes* positivity directly | Vanilla policy gradient against a VADER sentiment reward, with a KL penalty to a frozen SFT reference |
+
+## Headline Result -- Positivity Climbs at Every Stage
+
+Mean VADER `compound` sentiment over generated stories (higher = more cheerful, range `-1..+1`):
+
+| Stage | Mean Sentiment | Std |
+|:--|:--:|:--:|
+| Pretrained | `+0.8428` | 0.3907 |
+| SFT (positive) | `+0.8703` | 0.2853 |
+| **RLVR** | **`+0.9001`** | 0.3371 |
+
+RLVR raises mean positivity *and* the SFT stage tightens the spread -- the pipeline makes
+the model both **happier** and **more consistent**. Individual RLVR stories routinely score
+`+0.98` and above.
+
+## Quick Start -- Gradio Space (no install)
+
+Try the model in your browser, including a side-by-side **3-model comparison** view:
+
+[**nanoGPT SLM -- Cheerful Story Generator + Illustration**](https://huggingface.co/spaces/nishantup/nanogpt-rlvr-slm-tinystories)
+
+## Programmatic Use -- the RLVR Model
+
+### Option 1: Run the inference script directly
+```bash
+# downloads weights from the Hub and runs sample generations
+pip install torch tiktoken huggingface_hub nltk
+python nanogpt_slm_rlvr_inference_tinystories.py
+```
+
+### Option 2: Import and generate
+```python
+# pip install torch tiktoken huggingface_hub nltk
+# place nanogpt_slm_rlvr_inference_tinystories.py in your working directory
+from nanogpt_slm_rlvr_inference_tinystories import tell_story, ask, generate_text
+
+print(tell_story("Once upon a time there was a little kitten"))
+print(ask("The friendly dragon lived in"))
+print(generate_text("A girl named Lily went to the park",
+                     max_tokens=300, temperature=0.8, top_k=40))
+```
+
+### Option 3: Load the RLVR weights manually
+```python
+from huggingface_hub import hf_hub_download
+import torch
+from nanogpt_slm_rlvr_inference_tinystories import GPTKV, GPTConfig
+
+path  = hf_hub_download("nishantup/nanogpt-rlvr-slm-tinystories-124m",
+                        "nanogpt_slm_rlvr_final.pth")
+model = GPTKV(GPTConfig())          # KV cache enabled for fast generation
+model.load_state_dict(torch.load(path, map_location="cpu"))
+model.eval()
+```
+
+## Comparing the Three Models
+
+The snippet below downloads all three checkpoints and prints each stage's story for the
+same prompt, scored with the same VADER metric the RLVR stage was trained against:
+
+```python
+# pip install torch tiktoken huggingface_hub nltk
+import torch, tiktoken, nltk
+from huggingface_hub import hf_hub_download
+from nanogpt_slm_rlvr_inference_tinystories import GPTKV, GPTConfig
+
+nltk.download("vader_lexicon", quiet=True)
+from nltk.sentiment.vader import SentimentIntensityAnalyzer
+
+REPO = "nishantup/nanogpt-rlvr-slm-tinystories-124m"
+enc  = tiktoken.get_encoding("gpt2")
+cfg  = GPTConfig()
+sia  = SentimentIntensityAnalyzer()
+
+def load(fname):
+    m = GPTKV(cfg)
+    m.load_state_dict(torch.load(hf_hub_download(REPO, fname), map_location="cpu"))
+    m.eval()
+    return m
+
+models = {
+    "Pretrained": load("nanogpt_slm_tinystories_best.pth"),
+    "SFT":        load("nanogpt_slm_sft_best.pth"),
+    "RLVR":       load("nanogpt_slm_rlvr_final.pth"),
+}
+
+def story(m, prompt, max_tokens=250, seed=1234):
+    torch.manual_seed(seed)                       # same RNG start for a fair compare
+    idx  = torch.tensor(enc.encode_ordinary(prompt)).unsqueeze(0)
+    out  = m.generate(idx, max_new_tokens=max_tokens, temperature=0.8, top_k=40)
+    toks = out.squeeze(0).tolist()
+    if 50256 in toks:
+        toks = toks[:toks.index(50256)]
+    return enc.decode(toks)
+
+prompt = "The little girl was sad until"
+for name, m in models.items():
+    s     = story(m, prompt)
+    score = sia.polarity_scores(s)["compound"]
+    print(f"\n=== {name}  (sentiment {score:+.3f}) ===\n{s}")
+```
+
+Typical result: the **Pretrained** model may take the story in any emotional direction, the
+**SFT** model leans positive, and the **RLVR** model produces the most reliably cheerful,
+high-sentiment continuation.
+
+## Model Architecture
+
+All three checkpoints share the same architecture:
+
+| Attribute | Value |
+|:---|:---|
+| Architecture | nanoGPT (GPT-2 small: 12 layers, 12 heads, 768 dim) |
+| Parameters | 124M (85.4M unique, with weight tying) |
+| Context length | 512 tokens |
+| Tokenizer | tiktoken GPT-2 BPE (50,257 tokens) |
+| Attention | Flash Attention when available, causal mask |
+| Normalization | Pre-norm (LayerNorm before attention/MLP) |
+| KV Cache | `GPTKV` variant included for O(1) per-token decode |
+| EOS token | `<|endoftext|>` (50256) -- learned story boundary |
+
+## Training Details
+
+### Stage 1 -- Pretraining
+| Attribute | Value |
+|:---|:---|
+| Data | TinyStories (~2.1M stories, ~470M tokens) |
+| Iterations | 70,000 |
+| Optimizer | AdamW (lr `6e-4` -> `1e-5` cosine, betas `(0.9, 0.95)`, wd `0.1`) |
+| Batch | 32 x 512 tokens, grad-accum 4 |
+| Precision | bfloat16 (A100) |
+
+### Stage 2 -- Supervised Fine-Tuning (SFT)
+| Attribute | Value |
+|:---|:---|
+| Data | Positive-sentiment subset of TinyStories (VADER compound > `+0.05`) -- 1.91M stories (90.2%), ~424M tokens |
+| Iterations | 12,952 (~2 epochs) |
+| Optimizer | AdamW, peak lr `5e-5` -> `5e-6` cosine (about 12x below pretraining) |
+| Batch | 32 x 512 tokens, grad-accum 4 |
+| Best val loss | 1.2037 |
+
+### Stage 3 -- Reinforcement Learning with Verifiable Rewards (RLVR)
+| Attribute | Value |
+|:---|:---|
+| Algorithm | Vanilla policy gradient |
+| Reward | VADER `compound` sentiment of the completed story (verifiable, deterministic) |
+| Reward broadcasting | Sequence-level reward applied to every token in the trajectory |
+| KL penalty | Against a **frozen SFT reference** (`beta = 0.1`) -- prevents reward hacking |
+| Generation batch | 16 trajectories, 200 tokens each |
+| Iterations | 200 |
+| Optimizer | AdamW, lr `5e-6` |
+| Mean reward | `+0.6485` -> `+0.8652` (KL stays bounded, ~`0.022`) |
+
+**How RLVR works (one paragraph):** each iteration the policy samples a batch of stories;
+a VADER sentiment analyzer scores each completed story (one scalar reward); that scalar is
+broadcast to every generated token; a KL penalty against the frozen SFT model is subtracted
+so the policy cannot drift into degenerate text that merely games the scorer; and the
+vanilla policy-gradient loss `-(log_probs * final_rewards).mean()` is back-propagated.
+
+## Files
+
+| File | Description |
+|:---|:---|
+| `nanogpt_slm_tinystories_best.pth` | Stage 1 -- pretrained weights |
+| `nanogpt_slm_sft_best.pth` | Stage 2 -- SFT (positive-sentiment) weights |
+| `nanogpt_slm_rlvr_final.pth` | Stage 3 -- RLVR weights (**primary model**) |
+| `nanogpt_slm_rlvr_inference_tinystories.py` | Standalone inference script (RLVR + 3-model compare) |
+| `config.json` | Architecture, pipeline, and training metadata |
+
+## API Reference (`nanogpt_slm_rlvr_inference_tinystories.py`)
+
+| Function | Description |
+|:---|:---|
+| `tell_story(beginning, max_tokens=500, temperature=0.8, top_k=40)` | Generate a cheerful story from an opening line (RLVR model) |
+| `ask(prompt, ...)` | General text completion (alias of `generate_text`, RLVR model) |
+| `generate_text(prompt, ...)` | Low-level generation with full parameter control (RLVR model) |
+| `compare_models(prompt, ...)` | Generate the same prompt from all 3 stages and return stories + VADER scores |
+
+## Example Outputs (RLVR Model)
+
+**Prompt:** "Once upon a time"  *(sentiment +0.99)*
+> Once upon a time, there was a little girl named Lily. She loved to play with her toys
+> and her friends. One day, Lily's mommy gave her a present... She hugged the doll and
+> said, "I love you, doll!"
+
+**Prompt:** "On a bright morning"  *(sentiment +0.99)*
+> On a bright morning, Molly was very excited for her first day ever. She put on her dress
+> and ran outside to the garden... The rabbit smiled and said, "Thank you for coming down
+> to play with us!"
+
+## Limitations
+
+- Short stories only (512-token context window)
+- Simple vocabulary and narrative structures (by design -- TinyStories style)
+- No instruction-following ability
+- Strongly biased toward positive sentiment (that is the goal of the pipeline)
+- English only; may occasionally repeat or produce minor inconsistencies
+
+## Citation
+
+```
+Eldan, R., & Li, Y. (2023). TinyStories: How Small Can Language Models Be
+and Still Speak Coherent English? arXiv preprint arXiv:2305.07759.
+```
+
+## Notes
+
+- Big shout-out to **Dr. Raj Dandekar** (vizuara.ai) -- the RL/RLHF workshop this pipeline follows.
+- Trained completely from scratch (no pretrained initialization).
+- Architecture follows Karpathy's nanoGPT; weight tying between token embeddings and LM head.
+- RLVR uses a *verifiable* reward (VADER) -- deterministic, CPU-only, no reward model to train.
+- All three checkpoints are provided so the full Pretrain -> SFT -> RLVR progression is reproducible.