abhishekai/slm-125m-legal-rm

slm-125m-legal-rm

A 125.8M-parameter scalar reward model for grounded legal/financial Q&A, trained on AI-generated preferences (RLAIF). This is the reward signal used to train slm-125m-legal-ppo.

This is not a chat model. It emits a single scalar per sequence and cannot generate text. Load it with AutoModelForSequenceClassification, not AutoModelForCausalLM.

What it does

Given a rendered prompt concatenated with a candidate answer, it returns one number. Higher means "closer to what Gemini 2.5 Flash preferred." Only differences between scores are meaningful — the Bradley-Terry objective it was trained with fixes score gaps, not absolute magnitude, so the raw value is arbitrary and not comparable across reward models.

Note on the absolute-score matrix (2026-08-15)

Every generating checkpoint in this project was re-scored by a Claude Sonnet judge on four axes out of 10 over a shared set of 300 held-out prompts. A reward model has no cell in that table: it emits a scalar preference score rather than an answer, so there is nothing for an answer-quality rubric to grade. Its pairwise accuracy above remains the relevant measure. The full table is in MODEL_INDEX.md in the project repository.

Usage

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

name = "abhishekai/slm-125m-legal-rm"
tok = AutoTokenizer.from_pretrained(name)
rm = AutoModelForSequenceClassification.from_pretrained(name, num_labels=1).eval()
rm.config.pad_token_id = tok.convert_tokens_to_ids("<|pad|>")

def score(question, context, answer):
    msgs = [
        {"role": "system", "content": "You are a precise legal and financial assistant. Answer only from the provided context."},
        {"role": "user", "content": f"Context: {context}\n\nQuestion: {question}"},
    ]
    prompt = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
    enc = tok(prompt + answer.strip() + "<|eos|>", add_special_tokens=False,
              return_tensors="pt", truncation=True, max_length=1024)
    # input_ids/attention_mask only — the tokenizer also emits token_type_ids,
    # which LlamaForSequenceClassification.forward() rejects.
    with torch.no_grad():
        return rm(input_ids=enc["input_ids"],
                  attention_mask=enc["attention_mask"]).logits.squeeze().item()

The prompt must be rendered exactly as above. The reward model was trained on this template; scoring a differently-formatted string puts it out of distribution and the number becomes meaningless without any error being raised.

Training data

888 preference triplets built on-policy from the frozen SFT model:

1. 1.Sample 4 candidates per prompt (temp 0.9, top-p 0.95)
2. 2.Gemini 2.5 Flash scores each 1–10 on correctness and grounding
3. 3.Keep the best/worst pair only if the gap is ≥ 2
4. 4.An independent pairwise judge re-checks the ordering, A/B randomized
5. 5.Embedding dedup on prompts, then split

Limitations

• —It is a proxy, not a quality oracle. 0.837 pairwise accuracy means it disagrees with the AI judge on roughly one pair in six — and the judge itself is not ground truth.
• —It inherits Gemini 2.5 Flash's preferences, including any stylistic bias toward fluent prose over terse correctness.
• —In-domain only. Trained on grounded legal/financial QA against a supplied passage; scores on other tasks are not meaningful.
• —Optimizing hard against it will find its failure modes rather than improve answers — the usual reward-hacking caveat applies, which is why the PPO run it fed used an explicit KL anchor to the SFT policy.