KoGemma-E2B-tools-v6 (experimental)

LoRA adapter that teaches gemma-4-E2B-it to emit tool actions reliably — and the honest record of what it broke.

What this is

A 47 MB LoRA adapter (r=16) trained on 6,895 self-generated records: verified teacher CoT + base self-CoT (Mix Distillation) + tool trajectories + self-distilled replay.

Ship this if you want tool calling. Do not ship it if you need tool restraint.

Measured trade-off (held-out 52 cases, rule-based scoring, no LLM judge)

Capability	base	v6 adapter
web_search action correct	0.533	1.000
calculator action correct	0.000	0.750
fetch_page action correct	0.000	1.000
Searchable-future questions handled	0.750	1.000
Tool restraint (no action when unnecessary)	1.000	0.000
Private-info refusal	1.000	0.333
Overall	0.596	0.615

Reasoning / format (KMMLU test, 600 questions)	base	v6
CoT accuracy	0.3200	0.3250
CoT format failures	48 (8.0%)	27 (4.5%)
KMMLU 5-shot (lm-eval)	0.3000	0.3089
HAE-RAE 5-shot (lm-eval)	0.4600	0.4350

Why restraint collapsed (root cause, not speculation)

Every training example that carried a tool-bearing system prompt also contained a tool action. The data assembler routed the direct records (tool prompt + answer without an action) into a different pool, so the model only ever saw "tools visible ⇒ call a tool". It learned that rule perfectly.

Rule of thumb: when teaching a capability, include negative examples under the same prompt conditions. Target ratio ≈ 1 action : 1 no-action. The corrected dataset (ko-agentic-sft) ships 1,040 action and 543 no-action trajectories for exactly this reason.

Usage

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

BASE = "google/gemma-4-E2B-it"
tok = AutoTokenizer.from_pretrained(BASE)
model = AutoModelForCausalLM.from_pretrained(BASE, dtype=torch.bfloat16, device_map="auto")
model = PeftModel.from_pretrained(model, "waylake/KoGemma-E2B-tools-v6")

SYS = ('너는 한국어로 정확하게 답하고, 필요하면 도구를 쓰는 AI 비서다. '
       '도구가 필요하면 한 줄 JSON으로만 답한다: '
       '{"action": "web_search", "args": {"query": "..."}} '
       '사용 가능 도구: web_search{query}, fetch_page{url}, calculator{expr}, now{}')

msgs = [{"role": "system", "content": SYS},
        {"role": "user", "content": "지금 원달러 환율 얼마야?"}]
text = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True,
                               enable_thinking=False)
ids = tok(text, return_tensors="pt", add_special_tokens=False).to(model.device)
out = model.generate(**ids, max_new_tokens=120, do_sample=False)
print(tok.decode(out[0][ids["input_ids"].shape[1]:], skip_special_tokens=True))
# -> {"action": "web_search", "args": {"query": "원달러 환율"}}

Observation turns are fed back as a user message prefixed with [도구 결과] followed by JSON:

msgs += [{"role": "assistant", "content": '{"action": "web_search", "args": {"query": "원달러 환율"}}'},
         {"role": "user", "content": '[도구 결과] {"results": [{"title": "환율", "url": "https://example.kr/fx", "text": "1,387.4원"}]}'}]

Notes

• enable_thinking=False is recommended, not mandatory. Measured behaviour with thinking on: plain questions get an English reasoning block that eats the token budget, but tool questions still produce a correct action JSON — and the reasoning block comes out in Korean (thought 사용자는 현재 원/달러 환율을 묻고 있다. ... {"action": "web_search", ...}). Keep it off for latency; if you keep it on, raise max_new_tokens to ≳1000.
• Greedy decoding (temp 0) with repetition_penalty≈1.15 is the tested configuration.
• Adapter targets q/k/v/o and ffw_layer_1/2 after unwrapping Gemma4ClippableLinear; see the GitHub repo for the unwrap helper.

License

Follows the Gemma Terms of Use. Code MIT.