File size: 6,461 Bytes
7170091
 
bfe9979
7170091
 
bfe9979
7170091
bfe9979
 
 
 
 
 
7170091
 
bfe9979
 
 
 
 
 
7170091
 
bfe9979
7170091
bfe9979
7170091
bfe9979
7170091
bfe9979
7170091
bfe9979
7170091
bfe9979
7170091
bfe9979
 
 
 
7170091
bfe9979
7170091
bfe9979
7170091
bfe9979
7170091
bfe9979
7170091
bfe9979
7170091
bfe9979
 
 
7170091
bfe9979
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7170091
bfe9979
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7170091
bfe9979
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
---
base_model: CohereLabs/tiny-aya-global
license: cc-by-nc-4.0
library_name: peft
tags:
- peft
- lora
- qlora
- tool-calling
- agentic
- function-calling
- cohere
- tiny-aya
- sft
- trl
language:
- en
- es
- pt
- hi
pipeline_tag: text-generation
---

# Tiny Aya Agent (LoRA)

**A LoRA adapter for [`CohereLabs/tiny-aya-global`](https://huggingface.co/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.

```bash
pip install torch transformers peft bitsandbytes accelerate
```

```python
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.