Create mmlu_eval_sft.py

mmlu_eval_sft.py

@@ -0,0 +1,311 @@
+#!/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()