Hugging Face's logo

CMSManhattan
/
JiRackUltra_32b

Text Generation
Safetensors
GGUF
qwen2
ternary
bitnet
1.58bit
cpu
deepseek
efficient
low-memory
jirack
web-ui
routing
tool-call
robotics
conversational
Model card Files
xet
 Community
JiRackUltra_32b / mmlu_eval_sft.py
kgrabko's picture
kgrabko
Create mmlu_eval_sft.py
09e4b57 verified
Raw History Blame Contribute Delete
13.8 kB
#!/usr/bin/env python3
# mmlu_eval_sft.py — eval loop for JiRack models + automatic SFT set from the failures.
#
# 1. cais/mmlu ("all", split "test": 14042 questions, 57 subjects)
# 2. your model behind an OpenAI-compatible endpoint (Ollama: http://localhost:11434/v1)
# 3. the answer letter is checked against the MMLU key HERE (no judge)
# 4. accuracy overall and per subject -> summary.json
# 5. only for wrong / unparsed answers: DeepSeek writes the step-by-step reasoning
# for the KNOWN correct answer -> sft.jsonl ("question -> reasoning -> answer")
#
# Both endpoints use the same openai SDK client; only base_url / api_key / model differ.
# Keys come from the environment only:
# MODEL_API_KEY (Ollama ignores it; default "ollama")
# DEEPSEEK_API_KEY (required for the sft stage)
#
# Install: pip install openai datasets
# Examples:
# # quick smoke test, 50 questions, eval only
# python3 mmlu_eval_sft.py --model jirack-llama2b --limit 50 --stage eval
# # full run: eval, then SFT from the failures
# export DEEPSEEK_API_KEY=... # in the terminal, never in chat
# python3 mmlu_eval_sft.py --model jirack-llama2b --stage all --workers 4
# # later the ternary 16B: only the config changes
# python3 mmlu_eval_sft.py --model jirack-16b-ternary --out runs/16b
#
# Restart-safe: results.jsonl and sft.jsonl are appended per question; a rerun with the
# same --out skips everything already done.
#
# License note: sft.jsonl rows carry "generator" = the DeepSeek model. Output of the
# DeepSeek API falls under DeepSeek's terms (distillation allowed, AI-generated content
# must be disclosed) — keep that field when the data goes into training.
import argparse
import json
import os
import random
import re
import sys
import threading
import time
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
LETTERS = "ABCD"
ANSWER_SYSTEM = ("You are taking a multiple-choice test. Reply with the letter of the "
"correct option (A, B, C or D) and nothing else.")
REASON_SYSTEM = ("You are an expert tutor. You are given a multiple-choice question and its "
"correct answer. Explain step by step why that answer is correct and, briefly, "
"why the other options are wrong. Use facts, not guesses. Be concise "
"(at most about 200 words). The last line must be exactly: Answer: <letter>")
# ------------------------------------------------------------------ helpers
def format_question(q, choices):
lines = [q.strip(), ""]
for letter, text in zip(LETTERS, choices):
lines.append(f"{letter}. {str(text).strip()}")
return "\n".join(lines)
def parse_letter(text, choices):
"""Letter A-D from a model reply; None if it cannot be read unambiguously."""
if not text:
return None
t = text.strip()
# reasoning models: drop a <think>...</think> block
t = re.sub(r"<think>.*?</think>", "", t, flags=re.S).strip()
m = re.search(r"(?:answer|ответ)\s*(?:is|:)?\s*\(?([ABCD])\b", t, flags=re.I)
if m:
return m.group(1).upper()
m = re.match(r"^\(?([ABCD])(?:[\).:,\s]|$)", t)
if m:
return m.group(1).upper()
found = set(re.findall(r"\b([ABCD])\b", t))
if len(found) == 1:
return found.pop()
# model replied with the option text instead of the letter
low = t.lower().rstrip(".")
hits = [LETTERS[i] for i, c in enumerate(choices) if str(c).strip().lower() == low]
return hits[0] if len(hits) == 1 else None
def final_letter(text):
m = re.findall(r"Answer:\s*\(?([ABCD])\b", text or "", flags=re.I)
return m[-1].upper() if m else None
def load_done(path, key="idx"):
done = {}
if os.path.exists(path):
with open(path, encoding="utf-8") as f:
for line in f:
line = line.strip()
if line:
row = json.loads(line)
done[row[key]] = row
return done
class JsonlWriter:
def __init__(self, path):
self.f = open(path, "a", encoding="utf-8")
self.lock = threading.Lock()
def write(self, row):
with self.lock:
self.f.write(json.dumps(row, ensure_ascii=False) + "\n")
self.f.flush()
def close(self):
self.f.close()
def chat(client, model, messages, max_tokens, temperature, retries, timeout):
"""chat.completions with retry on rate limits, timeouts, connection and 5xx errors."""
import openai
delay = 2.0
for attempt in range(1, retries + 1):
try:
r = client.chat.completions.create(model=model, messages=messages,
max_tokens=max_tokens, temperature=temperature,
timeout=timeout)
return r.choices[0].message.content or ""
except (openai.RateLimitError, openai.APITimeoutError, openai.APIConnectionError,
openai.InternalServerError) as e:
err = e
except openai.APIStatusError as e:
if e.status_code < 500 and e.status_code != 429:
raise # 400/401/404: retrying will not help
err = e
if attempt == retries:
raise err
time.sleep(delay + random.random())
delay = min(delay * 2, 60)
# ------------------------------------------------------------------ data
def load_mmlu(args):
from datasets import load_dataset
ds = load_dataset("cais/mmlu", args.config, split=args.split)
rows = [{"idx": i, "subject": r["subject"], "question": r["question"],
"choices": list(r["choices"]), "answer": LETTERS[int(r["answer"])]}
for i, r in enumerate(ds)]
if args.subjects:
keep = set(args.subjects.split(","))
rows = [r for r in rows if r["subject"] in keep]
if args.shuffle:
random.Random(args.seed).shuffle(rows)
if args.limit:
rows = rows[:args.limit]
return rows
# ------------------------------------------------------------------ stage 1: eval
def run_eval(args, rows):
from openai import OpenAI
client = OpenAI(base_url=args.model_base_url,
api_key=os.environ.get("MODEL_API_KEY", "ollama"), max_retries=0)
path = os.path.join(args.out, "results.jsonl")
done = load_done(path)
todo = [r for r in rows if r["idx"] not in done]
print(f"[eval] {args.model} @ {args.model_base_url}: {len(rows)} questions, "
f"{len(rows) - len(todo)} already done, {len(todo)} to run")
w = JsonlWriter(path)
def one(r):
prompt = format_question(r["question"], r["choices"])
try:
reply = chat(client, args.model,
[{"role": "system", "content": ANSWER_SYSTEM},
{"role": "user", "content": prompt}],
args.answer_max_tokens, args.temperature, args.retries, args.timeout)
error = None
except Exception as e: # keep going; the row is marked
reply, error = "", f"{type(e).__name__}: {e}"
pred = parse_letter(reply, r["choices"])
row = {**r, "model": args.model, "reply": reply, "pred": pred,
"correct": pred == r["answer"], "error": error}
w.write(row)
return row
t0, n = time.time(), 0
with ThreadPoolExecutor(max_workers=args.workers) as ex:
for fut in as_completed([ex.submit(one, r) for r in todo]):
fut.result()
n += 1
if n % 50 == 0 or n == len(todo):
print(f"[eval] {n}/{len(todo)} | {(time.time() - t0) / n:.2f} s/question")
w.close()
def summarize(args, rows):
wanted = {r["idx"] for r in rows}
res = [r for i, r in load_done(os.path.join(args.out, "results.jsonl")).items() if i in wanted]
by = defaultdict(lambda: [0, 0])
for r in res:
by[r["subject"]][0] += r["correct"]
by[r["subject"]][1] += 1
n = len(res)
correct = sum(r["correct"] for r in res)
summary = {
"model": args.model, "dataset": f"cais/mmlu/{args.config}/{args.split}",
"questions": n, "correct": correct,
"accuracy": round(correct / n, 4) if n else None,
"unparsed": sum(r["pred"] is None and not r["error"] for r in res),
"api_errors": sum(bool(r["error"]) for r in res),
"per_subject": {s: {"accuracy": round(c / t, 4), "n": t}
for s, (c, t) in sorted(by.items(), key=lambda kv: kv[1][0] / kv[1][1])},
}
with open(os.path.join(args.out, "summary.json"), "w", encoding="utf-8") as f:
json.dump(summary, f, ensure_ascii=False, indent=2)
print(f"[summary] accuracy {summary['accuracy']} ({correct}/{n}) | unparsed "
f"{summary['unparsed']} | api errors {summary['api_errors']} (25% = random guessing)")
for s, v in list(summary["per_subject"].items())[:5]:
print(f" weakest: {s:40s} {v['accuracy']:.3f} (n={v['n']})")
return res
# ------------------------------------------------------------------ stage 2: SFT from failures
def run_sft(args, res):
key = os.environ.get("DEEPSEEK_API_KEY")
if not key:
sys.exit("DEEPSEEK_API_KEY is not set (export it in the terminal)")
from openai import OpenAI
client = OpenAI(base_url=args.deepseek_base_url, api_key=key, max_retries=0)
# API errors are excluded: the model never answered, that is not a knowledge failure
fails = [r for r in res if not r["correct"] and not r["error"]]
path = os.path.join(args.out, "sft.jsonl")
done = load_done(path, key="idx")
rej_path = os.path.join(args.out, "sft_rejected.jsonl")
rejected = load_done(rej_path, key="idx")
todo = [r for r in fails if r["idx"] not in done and r["idx"] not in rejected]
print(f"[sft] {len(fails)} failures, {len(fails) - len(todo)} already handled, "
f"{len(todo)} to send to {args.deepseek_model}")
w, wr = JsonlWriter(path), JsonlWriter(rej_path)
def one(r):
prompt = format_question(r["question"], r["choices"])
ans = r["answer"]
ask = (f"{prompt}\n\nCorrect answer: {ans}. "
f"{r['choices'][LETTERS.index(ans)]}\n\nExplain why.")
try:
text = chat(client, args.deepseek_model,
[{"role": "system", "content": REASON_SYSTEM},
{"role": "user", "content": ask}],
args.reason_max_tokens, 0.3, args.retries, args.timeout).strip()
except Exception as e:
wr.write({"idx": r["idx"], "reason": f"api: {type(e).__name__}: {e}"})
return False
if final_letter(text) != ans: # explanation argues for another option -> drop it
wr.write({"idx": r["idx"], "reason": "final letter mismatch", "text": text})
return False
body = re.sub(r"\n*\s*Answer:\s*\(?[ABCD]\)?\.?\s*$", "", text, flags=re.I).strip()
assistant = f"{body}\n\nAnswer: {ans}. {r['choices'][LETTERS.index(ans)]}"
w.write({
"idx": r["idx"], "subject": r["subject"],
"messages": [{"role": "user", "content": prompt},
{"role": "assistant", "content": assistant}],
"model_pred": r["pred"], "answer": ans,
"source": f"cais/mmlu/{args.config}/{args.split}",
"generator": f"deepseek:{args.deepseek_model}",
})
return True
ok = 0
with ThreadPoolExecutor(max_workers=args.workers) as ex:
for i, fut in enumerate(as_completed([ex.submit(one, r) for r in todo]), 1):
ok += fut.result()
if i % 25 == 0 or i == len(todo):
print(f"[sft] {i}/{len(todo)} | written {ok}")
w.close()
wr.close()
print(f"[sft] {path}: {len(load_done(path))} rows | rejected: {rej_path}")
# ------------------------------------------------------------------ main
def main():
ap = argparse.ArgumentParser(description="MMLU eval + SFT-from-failures via OpenAI-compatible APIs")
ap.add_argument("--stage", choices=["eval", "sft", "all"], default="all")
ap.add_argument("--out", default="runs/llama2b", help="output folder")
ap.add_argument("--model", required=True, help="model name on your server (ollama list)")
ap.add_argument("--model-base-url", default=os.environ.get("MODEL_BASE_URL", "http://localhost:11434/v1"))
ap.add_argument("--deepseek-base-url", default=os.environ.get("DEEPSEEK_BASE_URL", "https://api.deepseek.com"))
ap.add_argument("--deepseek-model", default=os.environ.get("DEEPSEEK_MODEL", "deepseek-chat"))
ap.add_argument("--config", default="all")
ap.add_argument("--split", default="test")
ap.add_argument("--subjects", default="", help="comma list, e.g. high_school_geography,world_religions")
ap.add_argument("--limit", type=int, default=0, help="0 = all questions")
ap.add_argument("--shuffle", action="store_true", help="random subset with --limit")
ap.add_argument("--seed", type=int, default=42)
ap.add_argument("--workers", type=int, default=1, help="parallel requests")
ap.add_argument("--temperature", type=float, default=0.0)
ap.add_argument("--answer-max-tokens", type=int, default=16)
ap.add_argument("--reason-max-tokens", type=int, default=600)
ap.add_argument("--retries", type=int, default=5)
ap.add_argument("--timeout", type=float, default=120.0)
args = ap.parse_args()
os.makedirs(args.out, exist_ok=True)
rows = load_mmlu(args)
if args.stage in ("eval", "all"):
run_eval(args, rows)
res = summarize(args, rows)
if args.stage in ("sft", "all"):
run_sft(args, res)
if __name__ == "__main__":
main()