aaie-gft-rca-sft-8k

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