aaie-gft-rca-sft-8k
Full-parameter SFT of AAIE_DDense_GFT_Llama
(the GFT model — distilled base + general instruction-tuning on SmolTalk) on a root-cause-analysis
(RCA) task: given a rubric, a student submission, and that submission's assessment results, produce
a structured JSON explanation of why the student scored the way they did per criterion. Trained
with YaRN 8x context scaling baked in (8,192-token context) since RCA examples run up to ~2.1K
tokens — past the base model's native 1,024-token training length.
Training code: sft_rca/ in the
source repo (train_sft.py + data/rca_dataset.py + configs/rca_sft_llama.yaml).
Request/response contract
Two-turn chat: a system message framing the task, then a user message with the rubric, the student's full submission, and that submission's per-criterion assessment results (score + excellent/good/average/etc. rating), asking for a JSON root-cause explanation:
system: You are an educational root cause analyst. Given a rubric, a student submission, and
the assessment results, identify the root causes of the student's performance as a JSON object.
user: Domain: ...
Essay Prompt: ...
Rubric Criteria:
[C1] ... (weight NN%): ...
...
Student Submission:
...
Assessment Results (overall: NN/100, <quality>):
[C1] ...: NN/NN (<rating>)
...
Quick start
from transformers import AutoModelForCausalLM, AutoTokenizer
MODEL_ID = "namquangstudy/aaie-gft-rca-sft-8k"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
messages = [
{"role": "system", "content": "You are an educational root cause analyst. Given a rubric, "
"a student submission, and the assessment results, identify "
"the root causes of the student's performance as a JSON object."},
{"role": "user", "content": "Domain: ...\n\n...\n\nAssessment Results (overall: 78/100, Good):\n ..."},
]
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(prompt, return_tensors="pt")
im_end_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
out = model.generate(
**inputs, max_new_tokens=400, do_sample=False,
eos_token_id=[im_end_id, tokenizer.eos_token_id], pad_token_id=tokenizer.eos_token_id,
)
print(tokenizer.decode(out[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True))
Explicitly passing eos_token_id=[im_end_id, tokenizer.eos_token_id] matters — this model's
chat template ends each turn with <|im_end|>, which some tooling's default eos_token_id
doesn't include, leading to ungoverned generation past the intended stop point.
Training procedure
transformers.Trainer, full-parameter finetune (every weight, no LoRA/adapters). Prompt tokens
masked out of the loss (-100), so only the target JSON contributes. 504 RCA-task examples (454
train / 50 held-out val), learning rate 1e-5 (kept small — full-parameter update starting from an
already instruction-tuned checkpoint, on a relatively small dataset, to avoid catastrophic
forgetting), 3 epochs, cosine schedule.
Result: train loss 2.98 → 0.076, eval loss → 0.086.
Validated separately by direct generation on 20 held-out examples (not the training loss
alone): 20/20 produced valid JSON and stopped cleanly at <|im_end|> with no repetition, once
eos_token_id is passed explicitly as shown above (the upstream base model's own
generation_config.json ships a stale eos_token_id left over from a generic Llama default —
already corrected in this checkpoint's generation_config.json, but worth knowing if you ever
diff it against the base model's).
Framework versions
- Transformers: 4.57.1
- Pytorch: 2.7.1+cu118
- Tokenizers: 0.22.1