KoGemma-E2B

<div align="center"> <img src="https://huggingface.co/waylake/KoGemma-E2B/resolve/main/banner.png" alt="KoGemma-E2B" width="100%"/>

KoGemma-E2B

Korean on-device SLM that actually calls tools — built on gemma-4-E2B-it, trained on 100% self-generated data (zero public datasets).

</div>

What you get

⚠️ Read this before using

This checkpoint was trained so that a tool-bearing system prompt means "tools are in play". Consequence, measured on held-out cases:

So: only put the tool list in the system prompt when you want tool calls. For plain chat use a plain system prompt. Root cause and the fix are documented in the GitHub repo and the corrected dataset (ko-agentic-sft ships 1,583 action vs 786 no-action trajectories for retraining).

Usage

Chat (no tools)

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

mid = "waylake/KoGemma-E2B"
tok = AutoTokenizer.from_pretrained(mid)
model = AutoModelForCausalLM.from_pretrained(mid, dtype=torch.bfloat16, device_map="auto")

msgs = [{"role": "system", "content": "너는 한국어로 정확하게 답하는 AI 비서다."},
        {"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=320, do_sample=False, repetition_penalty=1.15)
print(tok.decode(out[0][ids["input_ids"].shape[1]:], skip_special_tokens=True))

Tool calling

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": "지금 원달러 환율 얼마야?"}]
# -> {"action": "web_search", "args": {"query": "원달러 환율"}}

# feed the observation back as a user turn:
msgs += [{"role": "assistant", "content": '{"action": "web_search", "args": {"query": "원달러 환율"}}'},
         {"role": "user", "content": '[도구 결과] {"results": [{"title": "환율", "url": "https://example.kr/fx", "text": "1,387.4원"}]}'}]
# -> 현재 원/달러 환율은 1,387.4원입니다. (출처: https://example.kr/fx)

On a laptop (GGUF)

See waylake/KoGemma-E2B-GGUF — 3.4 GB Q4_K_M.

Recommended settings

Evaluation

Identical protocol for base and this model. lm-evaluation-harness 5-shot (--limit 40); CoT accuracy = generate + extract on 600 held-out KMMLU test questions; agentic = 52 hand-written held-out cases scored by rule (no LLM judge).

Honest summary: knowledge is unchanged, format and tool use improved, tool restraint regressed. Benchmark parity is expected — see the method below for why.

Method (why knowledge does not move)

1. Self-distillation (SDFT) prevents forgetting but caps quality at the base. Answers for plain chat are generated by the base model itself, so the output distribution never shifts (eval loss 1.57 → 0.44 vs a teacher-answer baseline that degraded the model).
2. Verified teacher CoT + the model's own CoT, mixed 1:1 (Mix Distillation, arXiv 2502.12143). Teacher traces are capped at 3–4 sentences because models ≤3B get worse from long chains. Teacher accuracy on those questions was 80–85%; the base scored 32% on the test split. Even so, student accuracy did not move — the bottleneck at 2.3B is capacity, not data.
3. Tool calling is structure, not knowledge, so distillation works immediately there (0% → 75–100%).

Training data

Released as ko-agentic-sft and ko-verified-cot.

Limitations

• 2.3B ceiling — not a substitute for 30B-class Korean models.
• Tool over-triggering with a tool-bearing system prompt (see warning above).
• Private-info refusal regressed (0.333); do not rely on it as a safety layer.
• Tool schema is the action-JSON format documented here, not OpenAI function calling.
• Medical / legal / financial answers are general information only.

Reproduce

Full pipeline — restart-safe generation queue, teacher distillation with gold verification, mixing, training, and all evaluators — is at github.com/waylake/kogemma-e2b, including the failed rounds and the measurement traps that produced fake gains.

License

Weights follow the Gemma Terms of Use; code MIT. Credits: Google DeepMind (gemma-4), HAERAE-HUB (KMMLU · HAE-RAE), lm-evaluation-harness. Synthetic-data teacher: ox-alpha-free.