Kris0404/social-story-gemma4-e2b-v5

Mystoria on-device Social Story model v5 (Gemma 4 E2B, MLX 6-bit mixed precision, 3.45 GB)

On-device text model for the Mystoria iOS app. It writes first-person Social Stories for autistic children following Carol Gray's Social Stories 10.4 criteria from a ten-field child profile (name, age, language level L1 to L4, communication notes, support level, interests, situation, target behavior, story type, people and companions).

Provenance

• Base: google/gemma-4-E2B-it (Apache 2.0), text-only extraction (vision and audio towers removed).
• Training: LoRA r16 supervised fine-tuning on 9,025 rule- and judge-filtered stories, then on-policy rejection sampling and length-balanced DPO (2,908 pairs). Adapter adapter_e2b_dpo1, sha256 152889a5c158f4a23dce964926a075e0ed5f1314eceb9c70a1fb69e7d7536d36. Trained by Mystoria co-founder YnnJ456 in the gemma4_e2b_smoke pipeline; converted and evaluated for MLX by Kris Yu.
• Quantization: MLX affine: linear layers 6-bit g64, token embedding 8-bit g64, per-layer embedding table 4-bit g32; per_layer_model_projection kept in fp16 and shared-KV attention keys present for mlx-swift-lm 3.31.x (per-layer entries in config.json["quantization"]). model.safetensors sha256 8f2e5848bad8f870860be57ad148059f2ea277bee36f462857768741cb5a13bc (3.453 GB). MLX peak memory on the real prompt plus one story: 4.27 GB. Intended for iPhones and iPads with 8 GB of RAM or more; the lite tier serves 6 GB devices.

Prompt contract

The model expects exactly the system prompt and user template of gemma4_e2b_smoke/deploy/story_service.py (prompt contract student-v5-2026-09), rendered with the bundled chat_template.jinja (system turn, user turn, add_generation_prompt=True). Output is a single JSON object {"title": ..., "pages": [{"page_number": n, "text": ...}]} with 8 to 12 pages.

Evaluation (410 held-out validation prompts, GPT judge, JSON and 11 rule checks and judge pass)

decoding	compliance
greedy	71.7% (294/410; W 113, F 2; JSON 410/410, rules 409/410)
bf16 reference 71.5%; friend's best 4-bit GGUF 63.9%

Reference: the bf16 adapter scored 71.5% under the same judge (greedy, vLLM).

Intended use and limits

Drafting aid for caregivers; every story is reviewed and edited by an adult before use. English only. The model can still produce stories that miss a rule; the app applies the same rule checker and hands such drafts to the caregiver as "needs review".

Swift verification and the config.json layout (2026-09-15)

Verified with a standalone macOS executable pinned to mlx-swift-lm 3.31.3 / mlx-swift 0.31.3 / swift-transformers 1.2.1 (the versions the Mystoria app uses): the model loads with its per-layer quantization entries, swift-jinja renders the bundled chat template to the same token ids as the Python tokenizer for all 120 probed prompts, and greedy generation through the AsyncStream chunk path scores 85/120 = 70.8% (Python mlx_lm greedy on the same 120 prompts: 72.5%) under the same JSON + rule + GPT-judge criterion as the table above; 88/120 stories are byte-identical to the Python greedy output. Peak MLX memory over the 120-prompt batch: 4.40 GB; about 40.6 tokens/s on an M3 Pro.

config.json keeps every text key (rope_parameters, sliding_window, layer_types, ...) both inside text_config and at the top level on purpose: for model_type: "gemma4_text" mlx-swift-lm 3.31.x decodes its configuration from the top level, and a layout with those keys only inside text_config makes the Swift model fall back to partial_rotary_factor 1.0 on the full-attention layers, which silently degrades output as the prompt gets longer (fenced, malformed JSON and placeholder names on the 1,000-token Mystoria prompt). If you downloaded this repo before 2026-09-15, re-download config.json.

Consumers that pin an earlier revision of this repo and unwrap text_config themselves (the Mystoria app does) need no change: their pinned config.json is the expected input for that build.