AgentVocab

This repository contains the code and scripts for AgentVocab: Structure-Aware Vocabulary Adaptation for Efficient LLM Agents.

AgentVocab is a structure-aware vocabulary adaptation framework for efficient LLM agents. It mines real tool-calling traces, adds reusable structural and content fragments as new vocabulary entries, and uses two-stage LoRA SFT with vocabulary expansion in the second stage. The resulting tokenizer reduces tool-calling tokenization overhead while preserving competitive task performance.

Highlights

• Structure-aware token induction: mines repeated tool-calling structures such as JSON-like schemas, function signatures, delimiters, and argument templates.
• Content-aware token induction: mines reusable content fragments with a TL-BPE + VEGAD-style gradient ranking pipeline.
• Two-stage LoRA SFT: first learns tool-calling behavior with the original tokenizer, then expands the vocabulary and continues LoRA SFT.
• Reproducible evaluation utilities: includes lightweight scripts for SWIFT Native + vLLM evaluation and result aggregation.

Why AgentVocab?

LLM agents are typically trained with general-purpose tokenizers, but deployed in narrow tool-calling pipelines dominated by structured schemas, function calls, arguments, and tool observations. This creates a training-deployment mismatch: repeated tool-calling fragments are split into long token sequences, increasing context length and decoding latency.

Method

AgentVocab mines structural patterns and content fragments from real tool-calling traces, then adapts the tokenizer through second-stage vocabulary expansion and LoRA SFT for agent deployment.

Results Preview

AgentVocab reduces input tokens, output tokens, and latency while maintaining competitive aggregate task performance. Input/output tokens and latency are averaged per turn.

tau-bench

tau2-bench

Vocabulary Budget Ablation

Structural tokens are the primary efficiency driver, and small content token budgets provide useful complements for balanced adaptation.

Installation

git clone https://github.com/Starry-159/AgentVocab.git
cd AgentVocab
conda create -n agentvocab python=3.10 -y
conda activate agentvocab
pip install -r requirements.txt
pip install -e .

Notes:

• AgentVocab uses ms-swift[all]==3.12.6. Do not replace it with SWIFT 4.x, because several scripts rely on SWIFT 3.x APIs.

• Install optional vLLM support only if you want to run SWIFT Native evaluation with the vLLM backend:

  pip install "vllm>=0.5.1"

• Install optional FlashAttention support after torch is available:

  pip install flash-attn --no-build-isolation

• Install benchmark packages only if you want to reproduce tau-bench or tau2-bench evaluation:

  pip install git+https://github.com/sierra-research/tau-bench
  pip install git+https://github.com/sierra-research/tau2-bench@v0.2.0

• If SWIFT installation fails in your environment, follow the official SWIFT installation guide: https://swift.readthedocs.io/en/latest/

• If evaluation fails, check the official EvalScope documentation: https://evalscope.readthedocs.io/en/latest/

• flash-attn must be installed after torch is available. If installation fails, install the wheel or module matching your CUDA, PyTorch, and GPU environment.

• Large datasets, model checkpoints, raw predictions, and private API keys are intentionally not included in this repository.

• Long-running scripts print [AgentVocab] status messages and use tqdm progress bars for data conversion, rendering, token mining, reachability filtering, vocabulary expansion, and checkpoint export.

Repository Structure

AgentVocab/
├── assets/                 # paper figures and optional result renders
├── examples/               # training and evaluation shell templates
├── scripts/                # command-line entry points
├── src/agentvocab/         # reusable Python modules
└── requirements.txt        # SWIFT 3.12 and utility dependencies

Quick Start

If your data is already in SWIFT agent format, start from step 2.

# 1. Convert Toucan parquet files to SWIFT agent JSONL.
python scripts/convert_toucan_to_swift.py \
  --input-dir /path/to/Toucan-1.5M/SFT \
  --output outputs/data/toucan_swift_agent_format.jsonl

# 2. Render SWIFT data into the actual model inputs seen by the tokenizer.
python scripts/render_swift_actual_content.py \
  --input outputs/data/toucan_swift_agent_format.jsonl \
  --output outputs/data/toucan_swift_agent_format_actual_content.jsonl \
  --model /path/to/Qwen2.5-7B-Instruct \
  --agent-template hermes

# 3. Filter valid records and render them again for token mining.
python scripts/filter_valid_data.py \
  --input outputs/data/toucan_swift_agent_format_actual_content.jsonl \
  --output outputs/data/toucan_swift_agent_format_valid.jsonl

python scripts/render_swift_actual_content.py \
  --input outputs/data/toucan_swift_agent_format_valid.jsonl \
  --output outputs/data/toucan_swift_agent_format_valid_actual_content.jsonl \
  --model /path/to/Qwen2.5-7B-Instruct \
  --agent-template hermes

# 4. Split rendered inputs into structural and content streams.
python scripts/split_structural_content.py \
  --input outputs/data/toucan_swift_agent_format_valid_actual_content.jsonl \
  --structural-output outputs/data/structural_spans.txt \
  --content-output outputs/data/content_text.txt \
  --sample-size 10000 \
  --random-seed 42

Omit --sample-size to process the full rendered corpus.

Token Mining

Structural Tokens

python scripts/mine_structural_tokens.py \
  --input outputs/data/structural_spans.txt \
  --tokenizer /path/to/Qwen2.5-7B-Instruct \
  --output outputs/tokens/structural_scored.json \
  --max-new-tokens 10000 \
  --min-frequency 10

Content Tokens

Content tokens are mined with a content-aware pipeline: TL-BPE candidate mining + VEGAD-style gradient ranking. The base tokenizer is used to generate token-level BPE candidates, while the scoring model is used to rank candidates with gradient signals. They can point to the same base model path, or the scoring model can point to a first-stage checkpoint.

python scripts/mine_content_tokens.py \
  --input outputs/data/content_text.txt \
  --base-tokenizer /path/to/Qwen2.5-7B-Instruct \
  --corpus outputs/data/toucan_swift_agent_format_valid_actual_content.jsonl \
  --scoring-model /path/to/stage1/checkpoint-or-base-model \
  --output outputs/tokens/content_scored.json \
  --max-new-tokens 10000 \
  --min-frequency 10 \
  --max-subwords 4 \
  --sample-size 1000 \
  --max-seq-len 4096

Reachability Filtering

Structural and content tokens share the same reachability filtering step. This simulates tokenizer insertion and keeps tokens that are actually selected on rendered training inputs.

python scripts/select_reachable_tokens.py \
  --scored-tokens outputs/tokens/structural_scored.json \
  --corpus outputs/data/toucan_swift_agent_format_valid_actual_content.jsonl \
  --tokenizer /path/to/Qwen2.5-7B-Instruct \
  --output-dir outputs/tokens \
  --token-type structural \
  --targets 200 500 800 1000 \
  --sample-size 10000 \
  --random-seed 42

python scripts/select_reachable_tokens.py \
  --scored-tokens outputs/tokens/content_scored.json \
  --corpus outputs/data/toucan_swift_agent_format_valid_actual_content.jsonl \
  --tokenizer /path/to/Qwen2.5-7B-Instruct \
  --output-dir outputs/tokens \
  --token-type content \
  --targets 100 200 400 800 \
  --sample-size 10000 \
  --random-seed 42

Omit --sample-size to run reachability filtering on the full corpus.

Mix Structural and Content Tokens

mix_tokens.py only merges token lists with order-preserving deduplication. Since mixed tokens can still affect each other's tokenizer segmentation, run select_reachable_tokens.py again on the mixed list to obtain the final reachable mixed vocabulary.

python scripts/mix_tokens.py \
  --inputs outputs/tokens/top_800_reachable_structural_tokens.json \
           outputs/tokens/top_200_reachable_content_tokens.json \
  --output outputs/tokens/mixed_800_structural_and_200_content_tokens.json

python scripts/select_reachable_tokens.py \
  --scored-tokens outputs/tokens/mixed_800_structural_and_200_content_tokens.json \
  --corpus outputs/data/toucan_swift_agent_format_valid_actual_content.jsonl \
  --tokenizer /path/to/Qwen2.5-7B-Instruct \
  --output-dir outputs/tokens \
  --token-type mixed_800_structural_and_200_content \
  --targets -1 \
  --sample-size 10000 \
  --random-seed 42

Vocabulary Expansion

New token embeddings and LM-head rows are initialized by averaging the original subtoken vectors.

python scripts/expand_tokenizer.py \
  --base-model /path/to/stage1/checkpoint \
  --tokens outputs/tokens/top_1000_reachable_mixed_800_structural_and_200_content_tokens.json \
  --output outputs/models/expanded_agentvocab_step0 \
  --device-map cpu

If fewer than the requested number of tokens are reachable, use the actual filename generated by select_reachable_tokens.py, e.g. top_953_reachable_..._tokens.json.

Training

AgentVocab uses two-stage LoRA SFT. Stage 1 adapts the base model to the agent-format training data with the original tokenizer. After mining and adding new vocabulary entries, Stage 2 continues LoRA SFT from the expanded Step0 model and saves embed_tokens / lm_head so that the new token embeddings are preserved.

The example scripts are templates. You can either edit the variables at the top of each script or pass them as environment variables.

Stage 1: Agent Adaptation Without Vocabulary Expansion

Inputs:

• MODEL_PATH: original base model, e.g. Qwen2.5-7B-Instruct.
• DATASET_PATH: SWIFT-format training JSONL.
• OUTPUT_DIR: directory for stage-1 LoRA checkpoints.

Default training settings in examples/train_stage1_lora.sh:

• LoRA target: all-linear
• LoRA rank / alpha: 64 / 128
• max length: 32768
• learning rate: 1e-4
• epochs: 4.0
• batch / accumulation: 1 / 32
• save/eval interval: 500
• template / agent template: qwen2_5 / hermes

MODEL_PATH=/path/to/Qwen2.5-7B-Instruct \
DATASET_PATH=outputs/data/toucan_swift_agent_format_valid.jsonl \
OUTPUT_DIR=outputs/stage1_lora \
CUDA_VISIBLE_DEVICES=0,1 \
NPROC_PER_NODE=2 \
MASTER_PORT=29501 \
bash examples/train_stage1_lora.sh

Stage 2: Continue SFT After Vocabulary Expansion

Before Stage 2, expand the tokenizer and initialize new embedding / LM-head rows:

python scripts/expand_tokenizer.py \
  --base-model /path/to/stage1/checkpoint \
  --tokens outputs/tokens/top_1000_reachable_mixed_800_structural_and_200_content_tokens.json \
  --output outputs/models/expanded_agentvocab_step0 \
  --device-map cpu

Inputs:

• BASE_MODEL: expanded Step0 model from scripts/expand_tokenizer.py.
• DATASET_PATH: SWIFT-format training JSONL.
• OUTPUT_DIR: directory for stage-2 LoRA checkpoints.

Default training settings in examples/train_stage2_lora.sh:

• LoRA target: all-linear
• modules to save: embed_tokens lm_head
• LoRA rank / alpha: 64 / 128
• max length: 8192
• learning rate: 5e-5
• scheduler / warmup: cosine / 0.05
• epochs: 3.0
• batch / accumulation: 1 / 16
• save/eval interval: 500
• template / agent template: qwen2_5 / hermes

BASE_MODEL=outputs/models/expanded_agentvocab_step0 \
DATASET_PATH=outputs/data/toucan_swift_agent_format_valid.jsonl \
OUTPUT_DIR=outputs/stage2_lora \
CUDA_VISIBLE_DEVICES=0,1,2,3 \
NPROC_PER_NODE=4 \
MASTER_PORT=29502 \
bash examples/train_stage2_lora.sh

If flash-attn is unavailable in your environment, remove or change --attn_impl flash_attn in the shell scripts.

Export LoRA Checkpoints

For vLLM evaluation, merge LoRA checkpoints into full model directories:

python scripts/merge_lora_checkpoints.py \
  --series-dir outputs/stage2_lora_series \
  --base-model outputs/models/expanded_agentvocab_step0 \
  --output-dir outputs/models \
  --output-prefix AgentVocab-Step

The exporter copies tokenizer files directly from the expanded Step0 model to preserve token-id mappings.

Evaluation

For a simple single-GPU SWIFT Native + vLLM evaluation:

TAU_BENCH_API_KEY=<your-key> \
MODEL_PATH=outputs/models/AgentVocab-StepXXXX \
MODEL_NAME=AgentVocab \
EVAL_DATASET=tau2_bench \
GPU_ID=0 \
PORT=10000 \
bash examples/evaluate_swift_native.sh

The example script is intentionally minimal and does not include multi-GPU scheduling. Do not commit API keys or private benchmark logs.

Aggregate tidy result files:

python scripts/aggregate_results.py \
  --input outputs/eval/tidy_results.xlsx \
  --output outputs/eval/overall_results.xlsx

Citation

If you use this code, please cite:

@inproceedings{bian2026agentvocab,
  title = {AgentVocab: Structure-Aware Vocabulary Adaptation for Efficient LLM Agents},
  author = {Kai Bian and Haosi Mo and Xuebo Liu and Shuangyong Song and Jing Li and Yongxiang Li and Min Zhang and Xuelong Li},
  booktitle = {Proceedings of the 43rd International Conference on Machine Learning},
  year = {2026}
}

License

This project is released under the MIT License. See LICENSE for details.