A 26M-parameter fine-tune of Cactus-Compute/needle specialized for transit tool-calling. It translates a natural-language transit query into one of two tool calls — or refuses when no tool applies.

The model is an extractor: it lifts the relevant slots from the query. Canonical resolution (station disambiguation, line lookup, routing) is the backend's job.


Parameters	26M
Base model	Cactus-Compute/needle
Architecture	Encoder-decoder Simple Attention Network (d_model=512, enc×12, dec×8, GQA 8H/4KV)
Vocab	8192 (SentencePiece BPE)
Format	JAX/Flax checkpoint (`.pkl`)
Languages	French, English, FR/EN code-switching
License	MIT

Supported tools

json

[
  {
    "name": "search_itinerary",
    "description": "Plan a route between two points.",
    "parameters": {
      "origin":      {"type": "string", "description": "Start location.", "required": true},
      "destination": {"type": "string", "description": "End location.", "required": true},
      "time_human":  {"type": "string", "description": "Human time expression, e.g. 'a 14h'.", "required": false},
      "time_mode":   {"type": "string", "description": "'depart_at' or 'arrive_by'.", "required": false}
    }
  },
  {
    "name": "get_next_arrivals",
    "description": "Next departures/arrivals at a stop.",
    "parameters": {
      "station": {"type": "string", "description": "Stop or station name.", "required": true},
      "line":    {"type": "string", "description": "Line identifier, e.g. '1', 'RER A'.", "required": false}
    }
  }
]

Capabilities

Two-tool routing — picks search_itinerary vs get_next_arrivals from intent.
Refusal — emits no tool call for off-topic / under-specified / ambiguous queries.
Robustness — handles typos, SMS-style compression, colloquial French, and FR/EN code-switching.
Time handling — depart_at vs arrive_by disambiguation from natural phrasing.

Examples

Example query → tool call (verified on this checkpoint):

Query	Output
`itinéraire pour Montparnasse`	`[{"name":"search_itinerary","arguments":{"destination":"Montparnasse"}}]`
`de Bastille à Nation vers 17h`	`[{"name":"search_itinerary","arguments":{"origin":"Bastille","destination":"Nation","time_human":"vers 17h","time_mode":"depart_at"}}]`
`Je voudrais arriver à Gare de Lyon pour 9h`	`[{"name":"search_itinerary","arguments":{"destination":"Gare de Lyon","time_human":"pour 9h","time_mode":"arrive_by"}}]`
`je suis à Bastille, direction La Défense`	`[{"name":"search_itinerary","arguments":{"origin":"Bastille","destination":"La Défense"}}]`
`nation bastille 8h`	`[{"name":"search_itinerary","arguments":{"origin":"nation","destination":"bastille","time_human":"8h","time_mode":"depart_at"}}]`
`I need to be at Nation by 9am from Bastille`	`[{"name":"search_itinerary","arguments":{"origin":"Bastille","destination":"Nation","time_human":"by 9am","time_mode":"arrive_by"}}]`
`Prochain RER B à Gare du Nord`	`[{"name":"get_next_arrivals","arguments":{"station":"Gare du Nord","line":"RER B"}}]`
`raconte une blague sur la SNCF`	`[]`
`combien coûte un ticket de métro ?`	`[]`

The model is an extractor — it returns slot values verbatim from the query (note the lowercased nation/bastille above); canonical resolution is the backend's job.

Results

Dataset and evaluation: coming soon. The training dataset and a held-out real-OOD evaluation suite will be released shortly.

Usage

Install the official needle package (Cactus Compute), then:

python

from needle import SimpleAttentionNetwork, load_checkpoint, generate, get_tokenizer

params, config = load_checkpoint("needle-transit.pkl")
model = SimpleAttentionNetwork(config)
tokenizer = get_tokenizer()

result = generate(
    model, params, tokenizer,
    query="Prochain RER B à Gare du Nord",
    tools='[{"name":"get_next_arrivals","description":"Next departures at a stop.","parameters":{"station":{"type":"string","description":"Stop name.","required":true},"line":{"type":"string","description":"Line id.","required":false}}}]',
    stream=False,
)
print(result)
# [{"name":"get_next_arrivals","arguments":{"station":"Gare du Nord","line":"RER B"}}]

Finetuning

Adapt the model to your own tools with the customized finetuning scripts in github.com/Rasaboun/needle-transit — tunable LR/Muon-LR, per-field loss weighting (--w-name/--w-value/--w-key), and metrics logging:

needle finetune data.jsonl --lr 3e-5 --w-value 4.0

Limitations

Single-call only — emits at most one tool call per query.
Domain-specific — tuned for transit tool-calling; off-domain tools are out of scope.
Extractor, not a resolver — returns surface slots; the backend resolves canonical names.
Small model — can be finicky on adversarial phrasing; finetune on your own data to adapt.

License & attribution

Citation

@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}
}

rasaboun/needle-transit

Needle Transit