Initial release: Gemma 4 E2B + SFT + GRPO (eval 0.940)

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,168 @@
+---
+license: apache-2.0
+language:
+- en
+base_model: google/gemma-4-E2B
+tags:
+- gemma-4
+- plugin-orchestration
+- function-calling
+- grpo
+- sft
+- on-demand
+- agentic
+library_name: transformers
+pipeline_tag: text-generation
+---
+
+# Gemma-4-E2B On-Demand Plugin Orchestrator
+
+**2B-active-parameter plugin selection and orchestration model, trained by [AIREV](https://airev.ai) for the [On-Demand](https://on-demand.io) plugin platform.**
+
+Given a user request and a candidate pool of plugins, this model picks the correct subset, orders them with proper dependencies, and hydrates every API parameter — emitting a valid JSON plan that can be executed directly.
+
+---
+
+## Eval results (100-sample held-out On-Demand set)
+
+| Model | Mean | JSON valid | Plugin-ID match | Count match | No-hallucinate | Hydrated | Deps chain |
+|---|---|---|---|---|---|---|---|
+| Gemma-4-E2B SFT-only (baseline) | 0.9180 | 95.0% | 89.0% | 93.0% | 95.0% | 95.0% | 87.0% |
+| **Gemma-4-E2B SFT + GRPO (this model)** | **0.9400** | **97.0%** | **91.0%** | **96.0%** | **97.0%** | **97.0%** | **89.0%** |
+
+**+2.2pp mean score = 26.8% relative error reduction over SFT-only.**
+
+### By category (GRPO wins on multi-step chains)
+
+| Category | SFT | GRPO | Δ |
+|---|---|---|---|
+| 1_step | 0.923 | 0.923 | — |
+| **2_step** | 0.950 | **1.000** | +5.0 |
+| **3_step** | 0.936 | **0.976** | +4.0 |
+| 4_step | 0.781 | 0.791 | +1.0 |
+| multi_turn | 1.000 | 1.000 | — |
+
+---
+
+## Usage
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch, json
+
+model_id = "airev-ai/gemma-4-e2b-ondemand"
+tok = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id, torch_dtype=torch.bfloat16
+).to("cuda")
+
+system = ("You are an AI agent orchestrator. Given a user request and available "
+          "plugins/tools, generate a precise multi-step execution plan as a valid "
+          "JSON object. Each step must use available plugins with correct parameters, "
+          "proper types, and valid JSON formatting.")
+
+candidates = [
+    {"pluginId": "plugin-1714851345", "name": "Nutrition BOT",
+     "description": "Nutrition type stuff", "identifier": "rest_api", "method": "POST"},
+    {"pluginId": "plugin-1768545918", "name": "kinetiqai-exercise-scoring",
+     "description": "Analyzes workout form using PoseTracker data",
+     "identifier": "rest_api", "method": "POST"},
+    # ... more candidates
+]
+
+user_msg = (
+    f"YOUR TASK IS TO GENERATE A JSON STRICTLY and CORRECTLY\n"
+    f"{json.dumps(candidates, indent=2)}\n\n"
+    f"User Request: I want to improve my fitness routine — analyze my workout "
+    f"form, then get nutrition guidance."
+)
+
+prompt = tok.apply_chat_template(
+    [{"role": "system", "content": system},
+     {"role": "user", "content": user_msg}],
+    tokenize=False, add_generation_prompt=True,
+)
+ids = tok(prompt, return_tensors="pt").input_ids.to("cuda")
+out = model.generate(ids, max_new_tokens=1024, temperature=0.1, do_sample=True, top_p=0.9)
+response = tok.decode(out[0][ids.shape[1]:], skip_special_tokens=True)
+
+# Response contains <think>...</think> followed by JSON:
+# {"plugins": [{"pluginId": "...", "api_request_parameters": {...},
+#               "all_parameters_hydrated": true, "dependencies": [...]}]}
+```
+
+## Output format
+
+The model emits:
+1. A `<think>...</think>` reasoning trace explaining plugin selection
+2. A JSON object: `{"plugins": [...]}` where each plugin has:
+   - `pluginId` (from the candidate list — never hallucinated)
+   - `name`, `description`, `identifier`, `method`
+   - `api_request_parameters` — fully hydrated, no placeholders
+   - `all_parameters_hydrated: true`
+   - `dependencies: []` — list of pluginIds that must run first
+
+## Training pipeline
+
+### 1. Data — 64,992 cleaned samples
+- Source: real On-Demand production traces + synthetic plans
+- Cleaning pipeline: deduplicated, JSON-validated, thinking tokens enforced,
+  parameters hydrated (no `example.com`, no empty values, no placeholders)
+- Judge: Claude Opus 4.6 via Vertex AI, 100 parallel workers
+
+### 2. SFT — 194,976 steps, 3 epochs
+- Base: `google/gemma-4-E2B` (5B total, 2B active, GDN hybrid)
+- Optimizer: **Adafactor** (AdamW causes CUDA illegal memory access on Gemma 4)
+- Single GPU only, no scheduler, no gradient clipping
+- LR = 2e-5, batch_size = 1, grad_accum = 1, max_length = 1024
+- Final loss: 0.1496 avg
+- ~20 hours on 1× H100 80GB
+- Eval score: **0.918**
+
+### 3. AutoResearch — 24 iterations hyperparameter search
+- Claude Opus 4.6 mutations with ratchet (keep best config)
+- Best finding: `num_plugins=6` candidates per prompt (down from 8)
+- Everything else stayed at defaults: `lr=1e-6`, `num_generations=4`, `top_p=0.9`
+
+### 4. GRPO — 570 steps with plugin-selection reward
+Reward (0.0–1.0) combines:
+- 0.10 valid JSON
+- 0.15 all picks in available candidate list
+- 0.25 × (correct_picked / total_correct)
+- 0.20 bonus for no wrong picks
+- −0.10 × wrong picks (capped at 3)
+- 0.10 exact count match
+- 0.15 × hydration ratio
+
+**Best checkpoint: step 500** — eval score **0.940**. Peak training-reward avg20 of 0.811 hit at step 473 before entering a noise-induced tail.
+
+## Architecture notes
+
+- **Gemma-4-E2B** = 5B total params, 2B active (MoE), 128K context, GDN-style
+- **Thinking tokens always active** — the model learned to use `<think>` for plugin reasoning
+- **Adafactor is mandatory** for training — AdamW hits illegal memory access
+- **Single GPU only** — `device_map="auto"` causes crashes
+- **No LR scheduler, no gradient clipping** — these also destabilize training
+
+## Code
+
+Full open-source training pipeline, AutoResearch harness, GRPO reward function,
+and eval scripts available at: [github.com/mk42-ai/gemma-4-e2b-ondemand](https://github.com/mk42-ai/gemma-4-e2b-ondemand)
+
+## Acknowledgements
+
+- Google DeepMind for Gemma 4
+- Berkeley RAIL for the BFCL benchmark methodology
+- HuggingFace TRL team for GRPO reference implementation
+- AIREV infrastructure team for 8× H100 cluster access
+
+## Citation
+
+```bibtex
+@misc{gemma4e2b_ondemand_2026,
+  title   = {Gemma-4-E2B On-Demand: Plugin Orchestration via SFT + GRPO},
+  author  = {Khalid, Muhammed and AIREV},
+  year    = {2026},
+  url     = {https://huggingface.co/airev-ai/gemma-4-e2b-ondemand},
+}
+```

autoresearch_best.json

@@ -0,0 +1,8 @@
+{
+  "lr": 1e-06,
+  "num_generations": 4,
+  "num_plugins": 6,
+  "temperature": 0.4,
+  "top_p": 0.9,
+  "max_samples": 5000
+}

autoresearch_history.json

@@ -0,0 +1,386 @@
+[
+  {
+    "iter": 0,
+    "reward": 0.5908203124999999,
+    "kept": true,
+    "desc": "baseline",
+    "config": {
+      "lr": 1e-06,
+      "num_generations": 4,
+      "num_plugins": 8,
+      "temperature": 0.7,
+      "top_p": 0.9,
+      "max_samples": 5000
+    },
+    "steps": 16,
+    "timestamp": "2026-04-07T19:51:48.336071"
+  },
+  {
+    "iter": 1,
+    "reward": 0.57,
+    "kept": false,
+    "desc": "Increase learning rate from 1e-6 to 5e-6 to enable stronger policy updates and faster reward signal propagation in early training",
+    "config": {
+      "lr": 5e-06,
+      "num_generations": 4,
+      "num_plugins": 8,
+      "temperature": 0.7,
+      "top_p": 0.9,
+      "max_samples": 5000
+    },
+    "steps": 15,
+    "timestamp": "2026-04-07T20:00:15.469756"
+  },
+  {
+    "iter": 2,
+    "reward": 0.4546875000000001,
+    "kept": false,
+    "desc": "Increase num_generations from 4 to 8 to improve GRPO advantage estimation with more diverse completions per prompt",
+    "config": {
+      "lr": 1e-06,
+      "num_generations": 8,
+      "num_plugins": 8,
+      "temperature": 0.7,
+      "top_p": 0.9,
+      "max_samples": 5000
+    },
+    "steps": 15,
+    "timestamp": "2026-04-07T20:08:44.085994"
+  },
+  {
+    "iter": 3,
+    "reward": 0.49166666666666664,
+    "kept": false,
+    "desc": "Decrease temperature from 0.7 to 0.5 to reduce sampling noise and produce more focused completions for structured plugin selection",
+    "config": {
+      "lr": 1e-06,
+      "num_generations": 4,
+      "num_plugins": 8,
+      "temperature": 0.5,
+      "top_p": 0.9,
+      "max_samples": 5000
+    },
+    "steps": 15,
+    "timestamp": "2026-04-07T20:17:16.055154"
+  },
+  {
+    "iter": 4,
+    "reward": 0.53625,
+    "kept": false,
+    "desc": "Increase max_samples from 5000 to 7500 to provide more training signal for learning plugin selection patterns",
+    "config": {
+      "lr": 1e-06,
+      "num_generations": 4,
+      "num_plugins": 8,
+      "temperature": 0.7,
+      "top_p": 0.9,
+      "max_samples": 7500
+    },
+    "steps": 5,
+    "timestamp": "2026-04-07T20:21:14.945985"
+  },
+  {
+    "iter": 5,
+    "reward": 0.5251953125,
+    "kept": false,
+    "desc": "Increase temperature from 0.7 to 0.8 to improve exploration diversity for better GRPO advantage contrast",
+    "config": {
+      "lr": 1e-06,
+      "num_generations": 4,
+      "num_plugins": 8,
+      "temperature": 0.8,
+      "top_p": 0.9,
+      "max_samples": 5000
+    },
+    "steps": 16,
+    "timestamp": "2026-04-07T20:30:19.973845"
+  },
+  {
+    "iter": 6,
+    "reward": 0.5225000000000001,
+    "kept": false,
+    "desc": "Decrease top_p from 0.9 to 0.85 to gently reduce sampling noise by trimming low-probability tail tokens without overly constraining diversity like temperature reduction did",
+    "config": {
+      "lr": 1e-06,
+      "num_generations": 4,
+      "num_plugins": 8,
+      "temperature": 0.7,
+      "top_p": 0.85,
+      "max_samples": 5000
+    },
+    "steps": 15,
+    "timestamp": "2026-04-07T20:38:53.846596"
+  },
+  {
+    "iter": 7,
+    "reward": 0.558203125,
+    "kept": false,
+    "desc": "Slightly increase learning rate from 1e-6 to 2e-6 for modest improvement in policy updates without instability",
+    "config": {
+      "lr": 2e-06,
+      "num_generations": 4,
+      "num_plugins": 8,
+      "temperature": 0.7,
+      "top_p": 0.9,
+      "max_samples": 5000
+    },
+    "steps": 16,
+    "timestamp": "2026-04-07T20:47:54.889003"
+  },
+  {
+    "iter": 8,
+    "reward": 0.6212500000000001,
+    "kept": true,
+    "desc": "Decrease num_plugins from 8 to 6 to simplify the selection task and reduce distractors, enabling cleaner learning signal",
+    "config": {
+      "lr": 1e-06,
+      "num_generations": 4,
+      "num_plugins": 6,
+      "temperature": 0.7,
+      "top_p": 0.9,
+      "max_samples": 5000
+    },
+    "steps": 5,
+    "timestamp": "2026-04-07T20:51:53.605514"
+  },
+  {
+    "iter": 9,
+    "reward": 0.6039583333333333,
+    "kept": false,
+    "desc": "Decrease num_plugins from 6 to 5 to further simplify plugin selection and reduce distractors",
+    "config": {
+      "lr": 1e-06,
+      "num_generations": 4,
+      "num_plugins": 5,
+      "temperature": 0.7,
+      "top_p": 0.9,
+      "max_samples": 5000
+    },
+    "steps": 15,
+    "timestamp": "2026-04-07T21:00:24.045740"
+  },
+  {
+    "iter": 10,
+    "reward": 0.5748046875,
+    "kept": false,
+    "desc": "Very slightly increase learning rate from 1e-6 to 1.5e-6 for a gentle boost in learning speed without destabilization",
+    "config": {
+      "lr": 1.5e-06,
+      "num_generations": 4,
+      "num_plugins": 6,
+      "temperature": 0.7,
+      "top_p": 0.9,
+      "max_samples": 5000
+    },
+    "steps": 16,
+    "timestamp": "2026-04-07T21:09:19.564525"
+  },
+  {
+    "iter": 11,
+    "reward": 0.458125,
+    "kept": false,
+    "desc": "Increase learning rate from 1e-6 to 3e-6 to accelerate learning on the now-simpler 6-plugin task where reward signal is cleaner",
+    "config": {
+      "lr": 3e-06,
+      "num_generations": 4,
+      "num_plugins": 6,
+      "temperature": 0.7,
+      "top_p": 0.9,
+      "max_samples": 5000
+    },
+    "steps": 15,
+    "timestamp": "2026-04-07T21:17:52.704223"
+  },
+  {
+    "iter": 12,
+    "reward": 0.626875,
+    "kept": true,
+    "desc": "Decrease temperature from 0.7 to 0.5 to produce more focused completions, which combined with the simpler 6-plugin task should improve reward signal quality for GRPO.",
+    "config": {
+      "lr": 1e-06,
+      "num_generations": 4,
+      "num_plugins": 6,
+      "temperature": 0.5,
+      "top_p": 0.9,
+      "max_samples": 5000
+    },
+    "steps": 10,
+    "timestamp": "2026-04-07T21:24:36.267719"
+  },
+  {
+    "iter": 13,
+    "reward": 0.6106250000000001,
+    "kept": false,
+    "desc": "Increase num_generations from 4 to 8 to provide richer advantage estimation per prompt for more stable GRPO updates",
+    "config": {
+      "lr": 1e-06,
+      "num_generations": 8,
+      "num_plugins": 6,
+      "temperature": 0.5,
+      "top_p": 0.9,
+      "max_samples": 5000
+    },
+    "steps": 15,
+    "timestamp": "2026-04-07T21:33:19.428449"
+  },
+  {
+    "iter": 14,
+    "reward": 0.5841796875,
+    "kept": false,
+    "desc": "Increase learning rate from 1e-6 to 2e-6, which previously showed promise (#7) and may benefit more from the current improved base config (fewer plugins + lower temperature)",
+    "config": {
+      "lr": 2e-06,
+      "num_generations": 4,
+      "num_plugins": 6,
+      "temperature": 0.5,
+      "top_p": 0.9,
+      "max_samples": 5000
+    },
+    "steps": 16,
+    "timestamp": "2026-04-07T21:42:12.799803"
+  },
+  {
+    "iter": 15,
+    "reward": 0.5007812500000001,
+    "kept": false,
+    "desc": "Decrease top_p from 0.9 to 0.8 to further sharpen sampling distribution, complementing the already-low temperature of 0.5",
+    "config": {
+      "lr": 1e-06,
+      "num_generations": 4,
+      "num_plugins": 6,
+      "temperature": 0.5,
+      "top_p": 0.8,
+      "max_samples": 5000
+    },
+    "steps": 16,
+    "timestamp": "2026-04-07T21:51:10.044767"
+  },
+  {
+    "iter": 16,
+    "reward": 0.6127672697368421,
+    "kept": false,
+    "desc": "Increase max_samples from 5000 to 7500 to provide more training data for better generalization with the current simplified and focused setup",
+    "config": {
+      "lr": 1e-06,
+      "num_generations": 4,
+      "num_plugins": 6,
+      "temperature": 0.5,
+      "top_p": 0.9,
+      "max_samples": 7500
+    },
+    "steps": 16,
+    "timestamp": "2026-04-07T22:00:01.729876"
+  },
+  {
+    "iter": 17,
+    "reward": 0.615,
+    "kept": false,
+    "desc": "Decrease temperature from 0.5 to 0.3 to further sharpen sampling for this structured selection task where deterministic outputs are beneficial",
+    "config": {
+      "lr": 1e-06,
+      "num_generations": 4,
+      "num_plugins": 6,
+      "temperature": 0.3,
+      "top_p": 0.9,
+      "max_samples": 5000
+    },
+    "steps": 15,
+    "timestamp": "2026-04-07T22:08:35.075433"
+  },
+  {
+    "iter": 18,
+    "reward": 0.5783333333333333,
+    "kept": false,
+    "desc": "Decrease learning rate from 1e-6 to 5e-7 to reduce update magnitude and improve training stability, since all higher LR attempts degraded performance",
+    "config": {
+      "lr": 5e-07,
+      "num_generations": 4,
+      "num_plugins": 6,
+      "temperature": 0.5,
+      "top_p": 0.9,
+      "max_samples": 5000
+    },
+    "steps": 15,
+    "timestamp": "2026-04-07T22:17:05.348585"
+  },
+  {
+    "iter": 19,
+    "reward": 0.6498046874999999,
+    "kept": true,
+    "desc": "Decrease temperature from 0.5 to 0.4 to explore the sweet spot between 0.3 (0.6150) and 0.5 (0.6269)",
+    "config": {
+      "lr": 1e-06,
+      "num_generations": 4,
+      "num_plugins": 6,
+      "temperature": 0.4,
+      "top_p": 0.9,
+      "max_samples": 5000
+    },
+    "steps": 16,
+    "timestamp": "2026-04-07T22:25:57.555577"
+  },
+  {
+    "iter": 20,
+    "reward": 0.5578125,
+    "kept": false,
+    "desc": "Very slightly increase learning rate from 1e-6 to 1.2e-6 to gently accelerate learning, leveraging the cleaner gradient signal from the low temperature=0.4 sampling",
+    "config": {
+      "lr": 1.2e-06,
+      "num_generations": 4,
+      "num_plugins": 6,
+      "temperature": 0.4,
+      "top_p": 0.9,
+      "max_samples": 5000
+    },
+    "steps": 16,
+    "timestamp": "2026-04-07T22:35:01.497719"
+  },
+  {
+    "iter": 21,
+    "reward": 0.6030555555555556,
+    "kept": false,
+    "desc": "Increase num_generations from 4 to 6 for better GRPO advantage estimation without the instability seen at 8",
+    "config": {
+      "lr": 1e-06,
+      "num_generations": 6,
+      "num_plugins": 6,
+      "temperature": 0.4,
+      "top_p": 0.9,
+      "max_samples": 5000
+    },
+    "steps": 15,
+    "timestamp": "2026-04-07T22:43:35.332127"
+  },
+  {
+    "iter": 22,
+    "reward": 0.5446875,
+    "kept": false,
+    "desc": "Decrease num_plugins from 6 to 5 to reduce distractor plugins, making selection easier and boosting correct pick rewards",
+    "config": {
+      "lr": 1e-06,
+      "num_generations": 4,
+      "num_plugins": 5,
+      "temperature": 0.4,
+      "top_p": 0.9,
+      "max_samples": 5000
+    },
+    "steps": 10,
+    "timestamp": "2026-04-07T22:50:41.402506"
+  },
+  {
+    "iter": 23,
+    "reward": 0.532421875,
+    "kept": false,
+    "desc": "Increase top_p from 0.9 to 0.95 to allow slightly broader nucleus sampling, complementing the already-focused temperature=0.4 for better exploration during GRPO generation.",
+    "config": {
+      "lr": 1e-06,
+      "num_generations": 4,
+      "num_plugins": 6,
+      "temperature": 0.4,
+      "top_p": 0.95,
+      "max_samples": 5000
+    },
+    "steps": 16,
+    "timestamp": "2026-04-07T22:59:41.386140"
+  }
+]

chat_template.jinja

@@ -0,0 +1,8 @@
+{% for message in messages %}{% if message.role == "system" %}<start_of_turn>system
+{{ message.content }}<end_of_turn>
+{% elif message.role == "user" %}<start_of_turn>user
+{{ message.content }}<end_of_turn>
+{% elif message.role == "assistant" %}<start_of_turn>model
+{{ message.content }}<end_of_turn>
+{% endif %}{% endfor %}{% if add_generation_prompt %}<start_of_turn>model
+{% endif %}

config.json

@@ -0,0 +1,188 @@
+{
+  "architectures": [
+    "Gemma4ForConditionalGeneration"
+  ],
+  "audio_config": {
+    "_name_or_path": "",
+    "architectures": null,
+    "attention_chunk_size": 12,
+    "attention_context_left": 13,
+    "attention_context_right": 0,
+    "attention_invalid_logits_value": -1000000000.0,
+    "attention_logit_cap": 50.0,
+    "chunk_size_feed_forward": 0,
+    "conv_kernel_size": 5,
+    "dtype": "bfloat16",
+    "gradient_clipping": 10000000000.0,
+    "hidden_act": "silu",
+    "hidden_size": 1024,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "model_type": "gemma4_audio",
+    "num_attention_heads": 8,
+    "num_hidden_layers": 12,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_proj_dims": 1536,
+    "problem_type": null,
+    "residual_weight": 0.5,
+    "return_dict": true,
+    "rms_norm_eps": 1e-06,
+    "subsampling_conv_channels": [
+      128,
+      32
+    ],
+    "use_clipped_linears": true
+  },
+  "audio_token_id": 258881,
+  "boa_token_id": 256000,
+  "boi_token_id": 255999,
+  "dtype": "bfloat16",
+  "eoa_token_id": 258883,
+  "eoa_token_index": 258883,
+  "eoi_token_id": 258882,
+  "image_token_id": 258880,
+  "initializer_range": 0.02,
+  "model_type": "gemma4",
+  "text_config": {
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "attention_k_eq_v": false,
+    "bos_token_id": 2,
+    "dtype": "bfloat16",
+    "enable_moe_block": false,
+    "eos_token_id": 1,
+    "expert_intermediate_size": null,
+    "final_logit_softcapping": 30.0,
+    "global_head_dim": 512,
+    "head_dim": 256,
+    "hidden_activation": "gelu_pytorch_tanh",
+    "hidden_size": 1536,
+    "hidden_size_per_layer_input": 256,
+    "initializer_range": 0.02,
+    "intermediate_size": 6144,
+    "layer_types": [
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention"
+    ],
+    "max_position_embeddings": 131072,
+    "model_type": "gemma4_text",
+    "moe_intermediate_size": null,
+    "num_attention_heads": 8,
+    "num_experts": null,
+    "num_global_key_value_heads": null,
+    "num_hidden_layers": 35,
+    "num_key_value_heads": 1,
+    "num_kv_shared_layers": 20,
+    "pad_token_id": 0,
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+      "full_attention": {
+        "partial_rotary_factor": 0.25,
+        "rope_theta": 1000000.0,
+        "rope_type": "proportional"
+      },
+      "sliding_attention": {
+        "rope_theta": 10000.0,
+        "rope_type": "default"
+      }
+    },
+    "sliding_window": 512,
+    "tie_word_embeddings": true,
+    "top_k_experts": null,
+    "use_bidirectional_attention": null,
+    "use_cache": true,
+    "use_double_wide_mlp": true,
+    "vocab_size": 262144,
+    "vocab_size_per_layer_input": 262144
+  },
+  "tie_word_embeddings": true,
+  "transformers_version": "5.5.0",
+  "use_cache": false,
+  "video_token_id": 258884,
+  "vision_config": {
+    "_name_or_path": "",
+    "architectures": null,
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "chunk_size_feed_forward": 0,
+    "default_output_length": 280,
+    "dtype": "bfloat16",
+    "global_head_dim": 64,
+    "head_dim": 64,
+    "hidden_activation": "gelu_pytorch_tanh",
+    "hidden_size": 768,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "intermediate_size": 3072,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "max_position_embeddings": 131072,
+    "model_type": "gemma4_vision",
+    "num_attention_heads": 12,
+    "num_hidden_layers": 16,
+    "num_key_value_heads": 12,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "patch_size": 16,
+    "pooling_kernel_size": 3,
+    "position_embedding_size": 10240,
+    "problem_type": null,
+    "return_dict": true,
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+      "rope_theta": 100.0,
+      "rope_type": "default"
+    },
+    "standardize": false,
+    "use_clipped_linears": true
+  },
+  "vision_soft_tokens_per_image": 280
+}

eval_results.json

@@ -0,0 +1,78 @@
+[
+  {
+    "model": "/root/checkpoints/gemma4_adafactor/final",
+    "n_samples": 100,
+    "elapsed_sec": 2972.9,
+    "mean_score": 0.918,
+    "pct_valid_json": 95.0,
+    "pct_pid_match": 89.0,
+    "pct_count_match": 93.0,
+    "pct_all_in_avail": 95.0,
+    "pct_hydrated": 95.0,
+    "pct_deps_match": 87.0,
+    "by_category": {
+      "3_step": 0.936,
+      "multi_turn": 1.0,
+      "4_step": 0.781,
+      "2_step": 0.95,
+      "1_step": 0.9231
+    }
+  },
+  {
+    "model": "/root/checkpoints/gemma4_grpo_v2/checkpoint-200",
+    "n_samples": 100,
+    "elapsed_sec": 2949.6,
+    "mean_score": 0.918,
+    "pct_valid_json": 96.0,
+    "pct_pid_match": 88.0,
+    "pct_count_match": 92.0,
+    "pct_all_in_avail": 96.0,
+    "pct_hydrated": 96.0,
+    "pct_deps_match": 87.0,
+    "by_category": {
+      "3_step": 0.968,
+      "multi_turn": 0.9524,
+      "4_step": 0.7429,
+      "2_step": 1.0,
+      "1_step": 0.9231
+    }
+  },
+  {
+    "model": "/root/checkpoints/gemma4_grpo_v2/checkpoint-500",
+    "n_samples": 100,
+    "elapsed_sec": 2946.2,
+    "mean_score": 0.94,
+    "pct_valid_json": 97.0,
+    "pct_pid_match": 91.0,
+    "pct_count_match": 96.0,
+    "pct_all_in_avail": 97.0,
+    "pct_hydrated": 97.0,
+    "pct_deps_match": 89.0,
+    "by_category": {
+      "3_step": 0.976,
+      "multi_turn": 1.0,
+      "4_step": 0.7905,
+      "2_step": 1.0,
+      "1_step": 0.9231
+    }
+  },
+  {
+    "model": "/root/checkpoints/gemma4_grpo_v2/final",
+    "n_samples": 100,
+    "elapsed_sec": 2977.0,
+    "mean_score": 0.934,
+    "pct_valid_json": 97.0,
+    "pct_pid_match": 90.0,
+    "pct_count_match": 94.0,
+    "pct_all_in_avail": 97.0,
+    "pct_hydrated": 97.0,
+    "pct_deps_match": 89.0,
+    "by_category": {
+      "3_step": 0.992,
+      "multi_turn": 0.9524,
+      "4_step": 0.7429,
+      "2_step": 1.0,
+      "1_step": 1.0
+    }
+  }
+]

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "bos_token_id": 2,
+  "do_sample": true,
+  "eos_token_id": 1,
+  "pad_token_id": 0,
+  "temperature": 1.0,
+  "top_k": 64,
+  "top_p": 0.95,
+  "transformers_version": "5.5.0"
+}

grpo_metrics.json

@@ -0,0 +1,578 @@
+{
+  "final_reward_avg20": 0.6996874999999999,
+  "final_reward_avg_all": 0.6659301900584759,
+  "best_reward": 0.95,
+  "total_steps": 570,
+  "reward_history": [
+    0.825,
+    0.95,
+    0.6593749999999999,
+    0.6124999999999999,
+    0.825,
+    0.5375,
+    0.6749999999999999,
+    0.12499999999999997,
+    0.825,
+    0.6000000000000001,
+    0.825,
+    0.4,
+    0.50625,
+    0.6124999999999999,
+    0.825,
+    0.825,
+    0.95,
+    0.825,
+    0.69375,
+    0.7812499999999999,
+    0.4,
+    0.50625,
+    0.4875,
+    0.4,
+    0.12499999999999997,
+    0.825,
+    0.825,
+    0.825,
+    0.825,
+    0.85625,
+    0.43125,
+    0.0,
+    0.390625,
+    0.825,
+    0.30000000000000004,
+    0.64375,
+    0.6125,
+    0.353125,
+    0.4,
+    0.4,
+    0.825,
+    0.825,
+    0.50625,
+    0.825,
+    0.95,
+    0.85625,
+    0.6187499999999999,
+    0.825,
+    0.825,
+    0.4,
+    0.825,
+    0.4,
+    0.4,
+    0.6124999999999999,
+    0.825,
+    0.825,
+    0.30625,
+    0.0,
+    0.8124999999999999,
+    0.7593749999999999,
+    0.825,
+    0.825,
+    0.95,
+    0.0,
+    0.4,
+    0.3,
+    0.825,
+    0.4,
+    0.4,
+    0.6312500000000001,
+    0.50625,
+    0.85625,
+    0.71875,
+    0.4,
+    0.5062500000000001,
+    0.50625,
+    0.4,
+    0.825,
+    0.825,
+    0.615625,
+    0.4,
+    0.71875,
+    0.825,
+    0.825,
+    0.825,
+    0.825,
+    0.5375,
+    0.91875,
+    0.4,
+    0.95,
+    0.95,
+    0.825,
+    0.0,
+    0.8875,
+    0.40625,
+    0.825,
+    0.825,
+    0.4,
+    0.725,
+    0.4062499999999999,
+    0.825,
+    0.95,
+    0.315625,
+    0.68125,
+    0.95,
+    0.30000000000000004,
+    0.0,
+    0.4,
+    0.95,
+    0.825,
+    0.5375,
+    0.09999999999999995,
+    0.95,
+    0.703125,
+    0.4,
+    0.71875,
+    0.4,
+    0.95,
+    0.825,
+    0.4,
+    0.6625,
+    0.6125,
+    0.6187499999999999,
+    0.8125,
+    0.4,
+    0.825,
+    0.825,
+    0.175,
+    0.95,
+    0.4,
+    0.95,
+    0.4,
+    0.925,
+    0.825,
+    0.4,
+    0.825,
+    0.95,
+    0.91875,
+    0.4,
+    0.825,
+    0.765625,
+    0.95,
+    0.6625,
+    0.91875,
+    0.825,
+    0.95,
+    0.50625,
+    0.40624999999999994,
+    0.725,
+    0.54375,
+    0.4,
+    0.825,
+    0.825,
+    0.5718749999999999,
+    0.5375,
+    0.4,
+    0.95,
+    0.825,
+    0.95,
+    0.825,
+    0.0,
+    0.825,
+    0.6187499999999999,
+    0.0,
+    0.6125,
+    0.71875,
+    0.4,
+    0.4375,
+    0.50625,
+    0.6124999999999999,
+    0.4,
+    0.675,
+    0.7187499999999999,
+    0.8875,
+    0.95,
+    0.328125,
+    0.50625,
+    0.71875,
+    0.95,
+    0.825,
+    0.275,
+    0.95,
+    0.36875,
+    0.825,
+    0.475,
+    0.95,
+    0.825,
+    0.4,
+    0.6125,
+    0.95,
+    0.825,
+    0.825,
+    0.825,
+    0.825,
+    0.825,
+    0.825,
+    0.825,
+    0.45937500000000003,
+    0.253125,
+    0.825,
+    0.26249999999999996,
+    0.825,
+    0.725,
+    0.825,
+    0.38125000000000003,
+    0.50625,
+    0.91875,
+    0.7999999999999999,
+    0.31249999999999994,
+    0.825,
+    0.825,
+    0.6625,
+    0.6124999999999999,
+    0.825,
+    0.7062499999999999,
+    0.69375,
+    0.95,
+    0.64375,
+    0.5093749999999999,
+    0.28125,
+    0.825,
+    0.71875,
+    0.95,
+    0.7187499999999999,
+    0.4,
+    0.95,
+    0.825,
+    0.825,
+    0.825,
+    0.95,
+    0.20625,
+    0.4,
+    0.95,
+    0.825,
+    0.85625,
+    0.4,
+    0.7718750000000001,
+    0.85625,
+    0.78125,
+    0.6124999999999999,
+    0.225,
+    0.6843750000000001,
+    0.95,
+    0.825,
+    0.58125,
+    0.4,
+    0.4,
+    0.825,
+    0.825,
+    0.825,
+    0.95,
+    0.95,
+    0.825,
+    0.58125,
+    0.825,
+    0.825,
+    0.6124999999999999,
+    0.825,
+    0.95,
+    0.825,
+    0.0,
+    0.68125,
+    0.71875,
+    0.51875,
+    0.95,
+    0.825,
+    0.95,
+    0.825,
+    0.825,
+    0.825,
+    0.5375,
+    0.4,
+    0.50625,
+    0.4,
+    0.3,
+    0.825,
+    0.95,
+    0.4,
+    0.5625,
+    0.825,
+    0.825,
+    0.825,
+    0.0,
+    0.4,
+    0.4,
+    0.725,
+    0.85625,
+    0.4,
+    0.6875,
+    0.71875,
+    0.64375,
+    0.825,
+    0.50625,
+    0.825,
+    0.825,
+    0.95,
+    0.95,
+    0.6125,
+    0.4,
+    0.4,
+    0.825,
+    0.825,
+    0.825,
+    0.8875,
+    0.825,
+    0.4,
+    0.5375000000000001,
+    0.91875,
+    0.95,
+    0.6625,
+    0.4,
+    0.6375,
+    0.7124999999999999,
+    0.4,
+    0.85625,
+    0.825,
+    0.825,
+    0.825,
+    0.6125,
+    0.40312499999999996,
+    0.825,
+    0.0,
+    0.6625,
+    0.725,
+    0.50625,
+    0.825,
+    0.6499999999999999,
+    0.825,
+    0.825,
+    0.20625,
+    0.825,
+    0.95,
+    0.825,
+    0.603125,
+    0.825,
+    0.35624999999999996,
+    0.0,
+    0.4,
+    0.95,
+    0.725,
+    0.6125,
+    0.4,
+    0.725,
+    0.825,
+    0.48124999999999996,
+    0.825,
+    0.4,
+    0.95,
+    0.6187499999999999,
+    0.95,
+    0.825,
+    0.91875,
+    0.765625,
+    0.4,
+    0.4,
+    0.83125,
+    0.7562500000000001,
+    0.71875,
+    0.95,
+    0.64375,
+    0.4,
+    0.4,
+    0.709375,
+    0.825,
+    0.825,
+    0.540625,
+    0.825,
+    0.31875,
+    0.6625,
+    0.825,
+    0.4,
+    0.71875,
+    0.49687499999999996,
+    0.09062499999999998,
+    0.35625,
+    0.95,
+    0.8875,
+    0.825,
+    0.825,
+    0.584375,
+    0.825,
+    0.825,
+    0.825,
+    0.5375000000000001,
+    0.71875,
+    0.825,
+    0.95,
+    0.825,
+    0.825,
+    0.4,
+    0.6124999999999999,
+    0.825,
+    0.825,
+    0.4,
+    0.47187499999999993,
+    0.0,
+    0.50625,
+    0.71875,
+    0.825,
+    0.4,
+    0.825,
+    0.36875,
+    0.95,
+    0.825,
+    0.4,
+    0.578125,
+    0.4,
+    0.825,
+    0.65625,
+    0.95,
+    0.5656249999999999,
+    0.89375,
+    0.95,
+    0.825,
+    0.725,
+    0.825,
+    0.95,
+    0.825,
+    0.4,
+    0.33125,
+    0.95,
+    0.825,
+    0.825,
+    0.703125,
+    0.95,
+    0.725,
+    0.825,
+    0.85625,
+    0.4,
+    0.35,
+    0.95,
+    0.95,
+    0.4,
+    0.825,
+    0.825,
+    0.4,
+    0.85625,
+    0.825,
+    0.95,
+    0.95,
+    0.0,
+    0.71875,
+    0.6125,
+    0.95,
+    0.43125,
+    0.0,
+    0.825,
+    0.71875,
+    0.91875,
+    0.725,
+    0.825,
+    0.825,
+    0.49999999999999994,
+    0.85625,
+    0.825,
+    0.4,
+    0.95,
+    0.825,
+    0.825,
+    0.95,
+    0.825,
+    0.7833333333333333,
+    0.825,
+    0.825,
+    0.6375,
+    0.825,
+    0.825,
+    0.825,
+    0.95,
+    0.825,
+    0.6625,
+    0.95,
+    0.825,
+    0.50625,
+    0.54375,
+    0.38749999999999996,
+    0.27499999999999997,
+    0.4,
+    0.0,
+    0.50625,
+    0.825,
+    0.825,
+    0.4,
+    0.4,
+    0.4,
+    0.825,
+    0.4,
+    0.825,
+    0.71875,
+    0.4,
+    0.6124999999999999,
+    0.825,
+    0.71875,
+    0.6187499999999999,
+    0.0,
+    0.825,
+    0.825,
+    0.4,
+    0.825,
+    0.4,
+    0.825,
+    0.825,
+    0.725,
+    0.6437499999999999,
+    0.95,
+    0.5875,
+    0.825,
+    0.4,
+    0.95,
+    0.0,
+    0.725,
+    0.825,
+    0.7437499999999999,
+    0.825,
+    0.825,
+    0.0,
+    0.95,
+    0.8374999999999999,
+    0.3,
+    0.825,
+    0.8125,
+    0.6937499999999999,
+    0.95,
+    0.825,
+    0.825,
+    0.725,
+    0.91875,
+    0.95,
+    0.50625,
+    0.825,
+    0.4,
+    0.0,
+    0.69375,
+    0.4,
+    0.0,
+    0.8875,
+    0.4,
+    0.825,
+    0.95,
+    0.95,
+    0.825,
+    0.95,
+    0.4,
+    0.95,
+    0.91875,
+    0.6625,
+    0.85625,
+    0.4,
+    0.4,
+    0.825,
+    0.71875,
+    0.6687500000000001,
+    0.825,
+    0.95,
+    0.825,
+    0.6625,
+    0.6125,
+    0.825,
+    0.85625,
+    0.825,
+    0.4,
+    0.825,
+    0.825,
+    0.825,
+    0.4,
+    0.049999999999999975,
+    0.4,
+    0.825,
+    0.95,
+    0.725
+  ]
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:effe6368c5be9d40523946b3c061fad56e3cac2dd5a7b4adad269702b80587b8
+size 10246621918

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:67d2db6635a88e60782a80257f796471b3601faea45c3de74958d2dff02864ec
+size 32170169

tokenizer_config.json

@@ -0,0 +1,60 @@
+{
+  "audio_token": "<|audio|>",
+  "backend": "tokenizers",
+  "boa_token": "<|audio>",
+  "boi_token": "<|image>",
+  "bos_token": "<bos>",
+  "eoa_token": "<audio|>",
+  "eoc_token": "<channel|>",
+  "eoi_token": "<image|>",
+  "eos_token": "<eos>",
+  "eot_token": "<turn|>",
+  "escape_token": "<|\"|>",
+  "etc_token": "<tool_call|>",
+  "etd_token": "<tool|>",
+  "etr_token": "<tool_response|>",
+  "extra_special_tokens": [
+    "<|video|>"
+  ],
+  "image_token": "<|image|>",
+  "is_local": true,
+  "mask_token": "<mask>",
+  "max_length": 1024,
+  "model_max_length": 1000000000000000019884624838656,
+  "model_specific_special_tokens": {
+    "audio_token": "<|audio|>",
+    "boa_token": "<|audio>",
+    "boi_token": "<|image>",
+    "eoa_token": "<audio|>",
+    "eoc_token": "<channel|>",
+    "eoi_token": "<image|>",
+    "eot_token": "<turn|>",
+    "escape_token": "<|\"|>",
+    "etc_token": "<tool_call|>",
+    "etd_token": "<tool|>",
+    "etr_token": "<tool_response|>",
+    "image_token": "<|image|>",
+    "soc_token": "<|channel>",
+    "sot_token": "<|turn>",
+    "stc_token": "<|tool_call>",
+    "std_token": "<|tool>",
+    "str_token": "<|tool_response>",
+    "think_token": "<|think|>"
+  },
+  "pad_to_multiple_of": null,
+  "pad_token": "<pad>",
+  "pad_token_type_id": 0,
+  "padding_side": "left",
+  "processor_class": "Gemma4Processor",
+  "soc_token": "<|channel>",
+  "sot_token": "<|turn>",
+  "stc_token": "<|tool_call>",
+  "std_token": "<|tool>",
+  "str_token": "<|tool_response>",
+  "stride": 0,
+  "think_token": "<|think|>",
+  "tokenizer_class": "GemmaTokenizer",
+  "truncation_side": "right",
+  "truncation_strategy": "longest_first",
+  "unk_token": "<unk>"
+}