Kris0404/social-story-gemma4-e2b-v5-lite

Mystoria on-device Social Story model v5 (Gemma 4 E2B, MLX 4-bit group size 32, 2.92 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 4-bit, group size 32, for every quantized layer; 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 0db0fe596118cb8946d750ff765d9f17cec11e1bc955a6b3f54d9510518c18ca (2.924 GB). MLX peak memory on the real prompt plus one story: 3.77 GB. Lite tier for 6 GB iPhones (13 Pro, 14, 15); enable only after an on-device memory test.

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)

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 84/120 = 70.0% (Python mlx_lm greedy on the same 120 prompts: 65.8%) under the same JSON + rule + GPT-judge criterion as the table above; 79/120 stories are byte-identical to the Python greedy output. Peak MLX memory over the 120-prompt batch: 3.95 GB; about 43.8 (60.8 when the machine is otherwise idle) 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.