gte-naga-tool-picker
A cross-encoder that ranks the agent tools of the naga simulation engine for a request. It is gte-reranker-modernbert-base (149M parameters) fine-tuned on about 1,000 synthetic tool-choice requests for naga's 90 tools. Given a request and a tool description it returns one score; sorting the 90 tools by score gives the picker's ranking. Weights are stored in float16 (about 290 MB).
Use with naga
The command needs naga 0.14.0 or later (pip install -U nagasim).
naga mcp --tool-picker haiku --tool-picker-backup TheMattFace/gte-reranker-naga-toolpicker
Use with sentence-transformers
import torch
from sentence_transformers import CrossEncoder
# On a GPU, float16 is fastest; on CPU pass float32 (float16 is about 7x slower there).
m = CrossEncoder("TheMattFace/gte-reranker-naga-toolpicker", model_kwargs={"torch_dtype": torch.float32}) # CPU
# m = CrossEncoder("TheMattFace/gte-reranker-naga-toolpicker", model_kwargs={"torch_dtype": torch.float16}) # GPU
tools = {"add_block": "Add a block to the model ...", "connect": "Connect two blocks ..."} # naga's 90 tool descriptions
request = "Link the loader to the conveyor"
scores = m.predict([(request, d) for d in tools.values()])
print(sorted(zip(tools, scores), key=lambda t: -t[1])[:3])
Results
Top-1 and top-3 over all 90 tools (the gold tool is in the first one or three). TEST3: 150 requests written by four different writers. AGENT3: 274 search queries as an AI agent writes them. Neither set was used in training.
| Picker | TEST3 top-1 | TEST3 top-3 | AGENT3 top-1 | AGENT3 top-3 |
|---|---|---|---|---|
keyword search (naga find_tool) |
36 / 150 | 63 | 136 / 274 | 190 |
| base model, not fine-tuned | 75 / 150 | 99 | 161 / 274 | 201 |
| this model | 104 / 150 | 128 | 183 / 274 | 215 |
Timing on a 4 GB laptop GPU (RTX A500), 90 pairs per request: median about 1.2 s in float16.
float16 against float32
Every number above is float16 on the GPU. Re-running the same model in float32 (all four test sets, 549 requests) gave the same top-1 pick on every request and the same top-3 order on 546 of 549 (99.5%). On 20 TEST3 requests, loading this folder with no dtype argument (it loads as float16) matched float32 exactly. Float16 is about 3x faster on the GPU (median 1.2 s against 3.9 s) and half the size.
On a CPU it is the other way round. Measured on 10 TEST3 requests with this folder, float16 took a median 117 s per request and float32 16.5 s, with the same top-1 pick on all 10. Load it as float32 on a CPU (the naga command does this for you from 0.14.0) and as float16 on a GPU.
Limitations
- Tuned to naga's 90-tool catalogue and its tool descriptions; it will not rank other tools well, and a new naga tool needs a re-tune.
- English only. About 1,000 synthetic requests is a small training set; expect errors on unusual wording (about 30% of TEST3 top-1 picks are wrong).
- A backup, not a replacement for a hosted picker: where a hosted model is acceptable it is about as accurate.
- Timings are from a 4 GB laptop GPU and will differ on other hardware.