README.md

---
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)