File size: 5,569 Bytes
564152c d962e86 564152c d962e86 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 | ---
license: mit
library_name: jax
tags:
- function-calling
- tool-use
- milogs
- action-engine
- personal-ai
- on-device
- journaling
- finance
- planning
- encoder-decoder
- edge
- jax
- flax
datasets:
- milogs-action-engine
base_model: Cactus-Compute/needle
---
# Milogs Action Engine
A fine-tuned version of [Needle](https://github.com/cactus-compute/needle) — a 26M-parameter Simple Attention Network — specialized for routing natural language voice transcriptions to structured journal, finance, planning, and search actions in the [Milogs](https://milogs.app) personal knowledge workspace.
The whole process was orchestrated using DeepSeek V4 Flash on Opencode, and this report too was crated using it.
**Base model:** [Cactus-Compute/needle](https://huggingface.co/Cactus-Compute/needle)
**Paper:** [Simple Attention Networks](https://github.com/cactus-compute/needle/blob/main/docs/simple_attention_networks.md)
**Original authors:** Henry Ndubuaku, Jakub Mroz, Karen Mosoyan, Roman Shemet, Parkirat Sandhu, Satyajit Kumar, Noah Cylich, Justin H. Lee
## What it does
Milogs Action Engine converts transcribed speech into structured tool calls. A user says:
> "Spent 3500 naira on lunch at that new place"
And the model outputs:
```json
[{"name":"log_transaction","arguments":{"type":"expense","amount":3500,"currency":"NGN","description":"lunch at new place","category":"food"}}]
```
It handles 7 tools:
| Tool | Purpose |
|------|---------|
| `log_activity` | Time-bounded activity journal entry |
| `log_transaction` | Financial transaction (expense/income/transfer) |
| `log_gratitude` | Gratitude journal entry |
| `create_plan` | Task or plan creation |
| `create_journal_entry` | Free-form journal note |
| `search_content` | Full-text search across content |
| `add_to_plan` | Add notes to existing plans |
## Architecture
| Property | Value |
|----------|-------|
| Parameters | 26M |
| Architecture | Encoder-decoder, pure attention (no FFN) |
| Encoder | 12 layers, GQA (8H/4KV), RoPE, gated residuals |
| Decoder | 8 layers, self-attn + cross-attn, gated residuals |
| d_model | 512 |
| Vocab | 8192 (SentencePiece BPE) |
| Norm | ZCRMSNorm (zero-centred, init=0) |
| Precision | bfloat16 |
| File size | ~50 MB |
## Training
### Base Model
The base Needle model was:
- Pretrained on 200B tokens using 16x TPU v6e (27 hours)
- Post-trained on 2B tokens of single-shot function call data (45 minutes)
### Fine-tuning (this model)
| Detail | Value |
|--------|-------|
| **Dataset** | 3,210 synthetic examples (450 per tool + 60 multi-tool) |
| **Generated by** | Large language model (via structured prompt) |
| **Train / Val / Test** | 3,070 / 70 / 70 (per-tool stratified split) |
| **Epochs** | 2 |
| **Batch size** | 64 |
| **Optimiser** | AdamW (non-kernel) + Muon (Dense kernels) |
| **Learning rate** | Adam: 3e-5, Muon: 0.02 |
| **Schedule** | WSD (Warmup-Stable-Decay): 4 warmup / 86 stable / 4 decay steps |
| **Loss weighting** | Base: 1.0, Name: 2.0, Value: 4.0, Key: 1.5 |
| **Hardware** | Apple M1 Pro (16 GB), CPU (JAX bfloat16) |
| **Duration** | ~3.5 hours |
| **Date** | 2026-07-21 |
### Results
| Metric | Base model | Fine-tuned |
|--------|-----------|------------|
| Name F1 | 69.2% | **100%** |
| Call F1 | 20.0% | **98.6%** |
| Exact match | 18.6% | **98.6%** |
| Parse rate | 100% | **100%** |
| Args accuracy | 28.9% | **98.6%** |
Per-tool (fine-tuned, test set):
- `add_to_plan`: 10/10
- `create_journal_entry`: 10/10
- `create_plan`: 10/10
- `log_activity`: 10/10
- `log_gratitude`: 10/10
- `log_transaction`: 10/11
- `search_content`: 10/10
## Usage
```python
from needle import SimpleAttentionNetwork, load_checkpoint, generate, get_tokenizer
params, config = load_checkpoint("milogs_action_engine.pkl")
model = SimpleAttentionNetwork(config)
tokenizer = get_tokenizer()
tools = '[{"name":"log_transaction","description":"Record a financial transaction.","parameters":{"type":{"type":"string","required":true},"amount":{"type":"number","required":true}}},{"name":"log_activity","description":"Log an activity.","parameters":{"text":{"type":"string","required":true}}},{"name":"log_gratitude","description":"Log gratitude.","parameters":{"text":{"type":"string","required":true}}},{"name":"create_plan","description":"Create a task.","parameters":{"title":{"type":"string","required":true}}},{"name":"create_journal_entry","description":"Create a note.","parameters":{"text":{"type":"string","required":true}}},{"name":"search_content","description":"Search content.","parameters":{"query":{"type":"string","required":true}}},{"name":"add_to_plan","description":"Add notes to a plan.","parameters":{"plan_title":{"type":"string","required":true},"note":{"type":"string","required":true}}}]'
result = generate(model, params, tokenizer, query="Spent 3500 on lunch", tools=tools, stream=False)
print(result)
# [{"name":"log_transaction","arguments":{"type":"expense","amount":3500,"description":"lunch"}}]
```
## Attribution
This model is a fine-tuned derivative of **Needle** by Cactus Compute.
```bibtex
@misc{ndubuaku2026needle,
title={Needle},
author={Henry Ndubuaku and Jakub Mroz and Karen Mosoyan and Roman Shemet and Parkirat Sandhu and Satyajit Kumar and Noah Cylich and Justin H. Lee},
year={2026},
url={https://github.com/cactus-compute/needle}
}
```
```bibtex
@misc{milogs2026actionengine,
title={Milogs Action Engine},
author={Milogs Team},
year={2026},
url={https://milogs.app}
}
```
## License
MIT (same as base Needle model)
|