ferscm44/tiny-aya-agent

Tiny Aya Agent (LoRA)

A LoRA adapter for CohereLabs/tiny-aya-global that teaches multi-step agentic reasoning — chaining data between calls, refusing impossible requests, and branching on conditions — distilled from Cohere's North Mini Code through an automatic verifier.

The base model does single-turn function calling well, but it cannot chain a returned id into the next call, does not refuse when no available tool fits, and does not branch on conditionals. This adapter closes that gap through verified distillation: a stronger teacher (North Mini Code) generates candidate trajectories, a strict automatic verifier keeps only the ones that provably satisfy the task, and those train Tiny Aya with QLoRA.

Full method, benchmark, and training code: https://github.com/fscm44xyz/aya-act

Results

Scored automatically by the verifier on the benchmark (20-scenario evaluation):

Model	Total	one_shot	data_chain	rejection	conditional
Tiny Aya (base)	5 / 20	5/5	0/5	0/5	0/5
Tiny Aya Agent (SFT v2, 80 examples)	15 / 20	5/5	5/5	1/5	4/5

From base to this adapter, data_chain goes 0→5 and conditional goes 0→4: the model learns to resolve "$N.field" references and to emit both branches of a condition — capabilities entirely absent in the base model.

Key finding: refusal is sensitive to dataset proportion

Refusal is not learned as a fixed skill — it is tuned by how much of the training set consists of rejection examples. At 25% rejection, the model over-refuses, declining solvable conditional tasks with {"error": "no tool available"}. At 12.5% (this adapter), it under-refuses, building plans for impossible requests instead of declining. Both proportions overshoot in opposite directions; the optimum lies between them. For a small model, the decision to act versus decline is shaped more by dataset composition than by any single example.

How to use

This is a PEFT/LoRA adapter, loaded on top of the 4-bit base model. The tokenizer in this repo carries the chat template used during training — load it from here, not from the base.

pip install torch transformers peft bitsandbytes accelerate

import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

BASE = "CohereLabs/tiny-aya-global"
ADAPTER = "ferscm44/tiny-aya-agent"

bnb = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.float16,
)

base = AutoModelForCausalLM.from_pretrained(
    BASE, quantization_config=bnb, device_map="auto",
    attn_implementation="sdpa", dtype=torch.float16,
)
model = PeftModel.from_pretrained(base, ADAPTER).eval()

# Load the tokenizer from THIS repo — it carries the training chat template.
tokenizer = AutoTokenizer.from_pretrained(ADAPTER)

# The model expects the full planning prompt as the user turn: the available
# tools, the output protocol, and the request. See runner.build_prompt in the
# project repo for the canonical prompt builder.
prompt = '''You are an agent that plans tool calls. You respond with JSON only.

AVAILABLE TOOLS:
{
  "find_customer": {"params": {"name": "string"}, "returns": {"id": "string"}},
  "get_account": {"params": {"customer_id": "string"}, "returns": {"account_id": "string", "status": "string"}},
  "get_balance": {"params": {"account_id": "string"}, "returns": {"balance": "number"}}
}

OUTPUT PROTOCOL:
1. Default: a JSON array of calls, e.g. [{"tool": "<name>", "args": {...}}, ...]
2. Reference an earlier result with "$N.field" (1-based step N).
3. For an if/else request: {"setup": [...], "branches": [{"condition": {...}, "calls": [...]}, {"condition": "else", "calls": [...]}]}
4. If no tool fits: {"error": "no tool available"}

USER REQUEST: What is the current balance of the customer named John Smith?'''

text = tokenizer.apply_chat_template(
    [{"role": "user", "content": prompt}],
    add_generation_prompt=True, tokenize=False,
)
inputs = tokenizer(text, return_tensors="pt", add_special_tokens=False).to(model.device)

# Stop at the standard eos or Aya's end-of-turn token.
eos_ids = [tokenizer.eos_token_id]
end_turn = tokenizer.convert_tokens_to_ids("<|END_OF_TURN_TOKEN|>")
if isinstance(end_turn, int) and end_turn >= 0 and end_turn != tokenizer.unk_token_id:
    eos_ids.append(end_turn)

with torch.no_grad():
    out = model.generate(
        **inputs, max_new_tokens=200, do_sample=False,
        eos_token_id=eos_ids,
        pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
    )

print(tokenizer.decode(out[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True))
# Expected:
# [{"tool": "find_customer", "args": {"name": "John Smith"}},
#  {"tool": "get_account", "args": {"customer_id": "$1.id"}},
#  {"tool": "get_balance", "args": {"account_id": "$2.account_id"}}]

Action protocol

The model is trained to answer only in the project's JSON action protocol:

• A sequence of calls: [{"tool": "<name>", "args": {...}}, ...]
• A data reference: the string "$N.field" — the field of step N's result (1-based).
• A conditional: {"setup": [...], "branches": [{"condition": {"field": "$N.field", "op": ">", "value": 100}, "calls": [...]}, {"condition": "else", "calls": [...]}]}
• A refusal: {"error": "no tool available"} when no available tool can satisfy the request.

Training

• Base: CohereLabs/tiny-aya-global
• Method: SFT with 4-bit QLoRA (rank 32), completion-only loss, on a single 24 GB GPU.
• Data: 80 verified trajectories generated by North Mini Code and filtered through the project's automatic verifier — the model only ever sees plans that pass verification.
• Format: Cohere chat fine-tuning JSONL; the training target and the evaluation use the same JSON protocol end to end.

License

CC-BY-NC-4.0 — non-commercial use only. This adapter is a derivative of CohereLabs/tiny-aya-global and inherits its non-commercial license restriction.