"""Export instruction-tuned weights and a standalone CPU demo, without uploading.

    python export.py --checkpoint checkpoints/sft/best.pt --out checkpoints/release
"""

import argparse
import json
from pathlib import Path
import shutil
import subprocess

import torch
from safetensors.torch import save_model

from eval import load as load_checkpoint
from inference import FIELDS, TEMPLATE, digest, load


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--checkpoint", default="checkpoints/sft/best.pt")
    parser.add_argument("--data", default="data")
    parser.add_argument("--out", default="checkpoints/release")
    args = parser.parse_args()
    out = Path(args.out)
    if out.exists() and any(out.iterdir()):
        raise ValueError("Output directory is not empty. Use a new directory.")
    torch.set_num_threads(4)
    model, tok, cfg, step = load_checkpoint(args.checkpoint, args.data, "cpu")
    checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
    if "sft" not in checkpoint or checkpoint["sft"]["template"] != TEMPLATE:
        raise ValueError("Export requires an instruction-tuned checkpoint with the story template.")
    out.mkdir(parents=True, exist_ok=True)
    save_model(model, str(out / "model.safetensors"))
    shutil.copyfile(Path(args.data) / "tokenizer.model", out / "tokenizer.model")
    architecture = {key: cfg[key] for key in ("n_layer", "dim", "n_head", "ctx", "vocab")}
    architecture["pos"] = cfg.get("pos", "rope")
    root = Path(__file__).resolve().parent
    commit = subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=root, text=True).strip()
    config = dict(format_version=1, model=architecture, template=TEMPLATE, fields=FIELDS,
                  special_tokens=tok.special, dtype="float32", step=step,
                  source_checkpoint_sha256=digest(args.checkpoint), source_commit=commit,
                  training_fingerprint=checkpoint["fingerprint"],
                  instruction_revision=checkpoint["sft"]["revision"],
                  sha256={name: digest(out / name) for name in ("model.safetensors", "tokenizer.model")})
    (out / "config.json").write_text(json.dumps(config, indent=2) + "\n")
    for name in ("model.py", "tokenizer.py", "inference.py", "app.py", "requirements.txt", "LICENSE"):
        shutil.copyfile(root / name, out / name)
    params = sum(p.numel() for p in model.parameters())
    run = checkpoint.get("wandb_id")
    card = f"""---
license: mit
language:
- en
datasets:
- roneneldan/TinyStories
- roneneldan/TinyStoriesInstruct
pipeline_tag: text-generation
tags:
- pytorch
- safetensors
- custom-code
---

# Sprout story model

A decoder-only model trained from scratch on TinyStoriesV2-GPT4 and instruction tuned on TinyStoriesInstruct. It has {params:,} parameters, {cfg['n_layer']} layers, width {cfg['dim']}, {cfg['n_head']} attention heads, a {cfg['ctx']}-token context, {architecture['pos']} positions, RMSNorm, SwiGLU, tied embeddings, and a {cfg['vocab']}-token byte-level BPE vocabulary.

[Try the live browser demo](https://huggingface.co/spaces/GeorgeCodesSometimes/sprout). The static demo runs float32 ONNX weights on your device; prompts stay local.

This bundle contains float32 inference weights at instruction step {step:,}. It excludes optimizer state, training data, credentials, and training dependencies. The tied embedding/output tensor is stored once and restored with `safetensors.torch.load_model`. `config.json` records weight and tokenizer SHA-256 checksums, source checkpoint SHA-256, source code commit, and the original training fingerprint.

## Run locally

Download all files from this model repository into an empty directory, then run:

```bash
python -m pip install -r requirements.txt
python app.py
```

Use the included loading module directly:

```python
from inference import load, stream

model, tok, config = load(".")
for text in stream(model, tok, {{"Summary": "A little fox learns to be patient."}}):
    print(text, end="\\r", flush=True)
```

`load` also accepts a Hub model ID and an optional `revision` commit. The included code is the loading contract; this is not a Transformers `AutoModel` repository. Loading does not require any training corpus files. The demo uses the exact training template, keeps the full instruction in context, and supports streamed output, stopping, temperature, and top-k.

## Training and evaluation

The source checkpoint's instruction validation loss is {checkpoint['best_val']:.6f}. The [training run](https://wandb.ai/george-nasseem/sprout/runs/{run}) records configuration and losses. See the [project README](https://github.com/gnasseem/sprout#instruction-tuning) for the fixed behavioral evaluation, comparison with the pretrained model, and base-validation regression. Low validation loss does not establish reliable constraint following.

Pretraining uses the story-level seeded split in the project, with 537,686,596 training tokens and 5,402,089 validation tokens. Instruction preparation excludes malformed and duplicate stories, held-out story overlap, and conversations beyond the context limit. TinyStoriesInstruct is pinned to `{config['instruction_revision']}`. These are synthetic English stories generated by GPT models, not a general conversational corpus.

## Limitations and intended use

An educational story-generation experiment. It often misses requested words and sentences, repeats itself, and loses narrative details. Instruction tuning regresses base-story validation loss. Long instructions leave less room for the story. It has no general chat or factual-answering capability, no content filter, and no suitability evaluation for children. Outputs can contain biases or disturbing story content from its synthetic source corpus.

## License and provenance

Code and exported weights use the included MIT license. The upstream [TinyStories dataset card](https://huggingface.co/datasets/roneneldan/TinyStories) lists CDLA-Sharing-1.0; [TinyStoriesInstruct](https://huggingface.co/datasets/roneneldan/TinyStoriesInstruct) is the separate instruction corpus. No source stories are redistributed in this bundle. The dataset creators' [paper](https://arxiv.org/abs/2305.07759) describes TinyStories.
"""
    (out / "README.md").write_text(card, encoding="utf-8")
    (out / "SPACE.md").write_text("""---
title: Sprout Story Workshop
sdk: gradio
sdk_version: 6.30.0
python_version: '3.11'
app_file: app.py
pinned: false
license: mit
---

# Sprout Story Workshop

An experimental small story model. It can miss constraints, repeat itself, and lose story details.
""")
    restored, restored_tok, _ = load(out)
    prefix = tok.encode(TEMPLATE.format(instruction="Write a short story.\nSummary: A fox finds a seed."))
    x = torch.tensor([prefix])
    before = model.generate(x, 64, temperature=0)
    after = restored.generate(x, 64, temperature=0)
    assert torch.equal(before, after), "export changed greedy generation"
    assert model.head.weight is model.tok_emb.weight
    assert restored.head.weight is restored.tok_emb.weight
    assert restored_tok.decode(after[0].tolist()) == tok.decode(before[0].tolist())
    print(f"Exported {params:,} parameters; 64 greedy tokens match; tied weights and tokenizer match.")
    print(f"Bundle: {out.resolve()}; weights: {(out / 'model.safetensors').stat().st_size:,} bytes")


if __name__ == "__main__":
    main()