""" Socratic Routing Evaluation — Paper 3 Compares 6 configurations on the same 9-dimensional test suite (Paper 2 baseline): Config 1 — Base Qwen2.5-1.5B (no LoRA) Config 2 — Math Specialist (dexmac/progressive-cognitive-dream-lora-en) Config 3 — Logic Specialist (dexmac/progressive-cognitive-logic-specialist-en) Config 4 — Monolith (dexmac/progressive-cognitive-logic-dream-lora-en, Paper 2) Config 5 — Socratic Router (router + math specialist + logic specialist) Config 6 — Qwen2.5-7B base (upper-bound reference: 4.7× more parameters) 9 Test dimensions: Logic : syllogism_valid, conditional_valid, boolean_eval, negation, compound_logic Math : exact_arithmetic, adversarial_math, magnitude_estimation, math_delegation 3 seeds × 50 samples each = 150 data points per dimension Inference pipeline for Config 5 (Socratic Router): 1. Router classifies query type 2. Route to math_specialist and/or logic_specialist 3. Router synthesizes expert response(s) into final answer Results saved to: dexmac/progressive-cognitive-results → socratic_eval_seed{N}_summary.json (per-seed) → socratic_eval_aggregate.json (3-seed mean ± std) """ import os import re import gc import json import math import time import random import operator import inspect import tempfile from pathlib import Path omp_threads = os.environ.get("OMP_NUM_THREADS", "").strip() if not omp_threads.isdigit() or int(omp_threads or "0") < 1: os.environ["OMP_NUM_THREADS"] = "1" import torch from collections import defaultdict from transformers import AutoTokenizer, AutoModelForCausalLM from peft import LoraConfig, PeftModel from huggingface_hub import hf_hub_download # ───────────────────────────────────────────────────────────────── # INFERENCE HELPERS # ───────────────────────────────────────────────────────────────── HF_TOKEN = os.environ.get("HF_TOKEN", "") MAX_NEW_TOKENS = { "syllogism_valid": 6, "conditional_valid": 6, "boolean_eval": 6, "negation": 18, "compound_logic": 18, "exact_arithmetic": 10, "adversarial_math": 10, "magnitude_estimation": 14, "math_delegation": 12, } def prepare_legacy_compatible_adapter(adapter_repo, token=""): temp_dir = Path(tempfile.mkdtemp(prefix="peft_adapter_")) config_path = hf_hub_download( repo_id=adapter_repo, filename="adapter_config.json", subfolder="lora_adapters", token=token or None, ) weights_path = hf_hub_download( repo_id=adapter_repo, filename="adapter_model.safetensors", subfolder="lora_adapters", token=token or None, ) with open(config_path) as f: config = json.load(f) valid_keys = set(inspect.signature(LoraConfig.__init__).parameters.keys()) config = {key: value for key, value in config.items() if key in valid_keys} with open(temp_dir / "adapter_config.json", "w") as f: json.dump(config, f, indent=2) with open(weights_path, "rb") as src, open(temp_dir / "adapter_model.safetensors", "wb") as dst: dst.write(src.read()) return str(temp_dir) def load_model(model_name_or_path, adapter_repo=None, device="cuda"): """Load base model + optional LoRA adapter.""" base = AutoModelForCausalLM.from_pretrained( model_name_or_path, torch_dtype=torch.float16, device_map="auto", trust_remote_code=True, ) tok = None if adapter_repo: try: tok = AutoTokenizer.from_pretrained( adapter_repo, subfolder="tokenizer", trust_remote_code=True, token=HF_TOKEN or None, ) except Exception: tok = None if tok is None: tok = AutoTokenizer.from_pretrained( model_name_or_path, trust_remote_code=True ) if tok.pad_token is None: tok.pad_token = tok.eos_token if adapter_repo: local_adapter_dir = prepare_legacy_compatible_adapter(adapter_repo, token=HF_TOKEN) base = PeftModel.from_pretrained( base, local_adapter_dir, ) base = base.merge_and_unload() base.eval() return base, tok def generate(model, tokenizer, prompt, max_new=30, device="cuda"): enc = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=256) enc = {k: v.to(device) for k, v in enc.items()} with torch.no_grad(): out = model.generate( **enc, max_new_tokens=max_new, do_sample=False, temperature=1.0, pad_token_id=tokenizer.eos_token_id, ) return tokenizer.decode( out[0][enc["input_ids"].shape[1]:], skip_special_tokens=True ).strip() def compute_math_tool_result(prompt): prompt = prompt.strip() magnitude = re.search(r"Estimate magnitude:\s*(-?\d+)\s*\*\s*(-?\d+)", prompt) if magnitude: a, b = int(magnitude.group(1)), int(magnitude.group(2)) ans = a * b oom = int(math.floor(math.log10(abs(ans)))) if ans != 0 else 0 return f"Calculator magnitude result: order {oom} (exact: {ans})" binary = re.search(r"(?:Calculate:|Quick:|What is|Compute|Analyze:)\s*(-?\d+)\s*([+\-*])\s*(-?\d+)", prompt) if binary: a, op, b = int(binary.group(1)), binary.group(2), int(binary.group(3)) ops = { "+": operator.add, "-": operator.sub, "*": operator.mul, } ans = ops[op](a, b) return f"Calculator result: {a} {op} {b} = {ans}" inventory = re.search( r"If a store has\s*(\d+)\s*items and sells\s*(\d+)\s*per day, after\s*(\d+)\s*days", prompt, flags=re.IGNORECASE, ) if inventory: a, b, c = map(int, inventory.groups()) remaining = a - b * c return f"Calculator result: remaining = {a} - {b}*{c} = {remaining}" perimeter = re.search( r"shape has\s*(\d+)\s*sides, calculate the perimeter if each side is\s*(\d+)cm", prompt, flags=re.IGNORECASE, ) if perimeter: sides, length = map(int, perimeter.groups()) result = sides * length return f"Calculator result: perimeter = {sides} * {length} = {result}" inline_add = re.search(r"Calculate\s*(-?\d+)\s*\+\s*(-?\d+)", prompt, flags=re.IGNORECASE) if inline_add: a, b = map(int, inline_add.groups()) return f"Calculator result: {a} + {b} = {a + b}" return None def should_use_calculator(prompt, raw_response): raw = raw_response.lower() prompt_lower = prompt.lower() asks_for_tool = any(token in raw for token in ["delegate", "tool", "calculator"]) math_pattern = any(token in prompt_lower for token in ["calculate:", "quick:", "what is", "compute", "estimate magnitude:", "analyze:"]) composed_pattern = any(token in prompt_lower for token in ["perimeter", "sells", "prime numbers greater than 2 are odd"]) return asks_for_tool or math_pattern or composed_pattern def first_nonempty_line(text): for line in text.splitlines(): line = line.strip() if line: return line return text.strip() def extract_classification_label(text): candidate = "\n".join(text.splitlines()[:3]).lower() for label in ["both", "math", "logic", "general"]: if re.search(rf"\b{label}\b", candidate): return label return "general" def extract_logic_verdict(text): lowered = text.lower() if "invalid" in lowered: return "invalid" if re.search(r"\bvalid\b", lowered): return "valid" if re.search(r"\btrue\b", lowered): return "true" if re.search(r"\bfalse\b", lowered): return "false" if re.search(r"\byes\b", lowered): return "yes" if re.search(r"\bno\b", lowered): return "no" return None def extract_last_integer(text): nums = re.findall(r"-?\d+", text) return nums[-1] if nums else None def extract_magnitude_order(text): order_match = re.search(r"order\s*(-?\d+)", text.lower()) if order_match: return order_match.group(1) exact_match = re.search(r"exact:\s*(-?\d+)", text.lower()) if exact_match: value = int(exact_match.group(1)) return str(int(math.floor(math.log10(abs(value))))) if value != 0 else "0" return None def infer_query_family(query): prompt = query.lower().strip() logic_markers = [ "premises:", "decide if conclusion follows", "respond only true or false", "negate:", "simplify:", ] math_markers = [ "calculate:", "quick:", "what is", "compute", "estimate magnitude:", "analyze:", "perimeter", "sells", ] if any(marker in prompt for marker in logic_markers): return "logic" if any(marker in prompt for marker in math_markers): return "math" return "general" class MathSpecialistWithCalculator: def __init__(self, model, tokenizer, device="cuda"): self.model = model self.tokenizer = tokenizer self.device = device def answer(self, prompt, max_new=30): raw_response = generate(self.model, self.tokenizer, prompt, max_new=max_new, device=self.device) tool_result = compute_math_tool_result(prompt) if tool_result and should_use_calculator(prompt, raw_response): return f"{raw_response}\n[Calculator]: {tool_result}".strip() return raw_response # ───────────────────────────────────────────────────────────────── # TEST SUITE GENERATORS (same as Paper 2) # ───────────────────────────────────────────────────────────────── class LogicTestGenerator: SYLLOGISMS_VALID = [ ("All dogs are animals. All animals are mortal.", "All dogs are mortal.", "VALID"), ("All squares are rectangles. ABCD is a square.", "ABCD is a rectangle.", "VALID"), ("All mammals breathe air. Whales are mammals.", "Whales breathe air.", "VALID"), ("All roses are flowers. All flowers have petals.","All roses have petals.", "VALID"), ("All cats are animals. Felix is a cat.", "Felix is an animal.", "VALID"), ("All metals conduct electricity. Gold is a metal.","Gold conducts electricity.","VALID"), ("All humans are mortal. Socrates is human.", "Socrates is mortal.", "VALID"), ] SYLLOGISMS_INVALID = [ ("All dogs are animals. Felix is an animal.", "Felix is a dog.", "INVALID"), ("All squares are rectangles. Shape is a rectangle.","Shape is a square.", "INVALID"), ("All A are B. All C are B.", "All A are C.", "INVALID"), ("All birds can fly. Rex can fly.", "Rex is a bird.", "INVALID"), ] CONDITIONALS_VALID = [ ("If it rains, the ground gets wet. It is raining.", "The ground gets wet.", "VALID"), ("If P then Q. Not Q.", "Not P.", "VALID"), ("If A then B. A is true.", "B is true.", "VALID"), ] CONDITIONALS_INVALID = [ ("If it rains, ground is wet. Ground is wet.", "It is raining.", "INVALID"), ("If the battery is dead, the phone turns off. Phone is off.", "The battery is dead.", "INVALID"), ] BOOLEAN_CASES = [ ("true AND false", "false"), ("true OR false", "true"), ("NOT true", "false"), ("true XOR false", "true"), ("false AND false", "false"), ("true AND true", "true"), ("true OR true", "true"), ("NOT false", "true"), ("true IMPLIES false","false"),("false IMPLIES false","true"), ("true IFF true", "true"), ("false IFF true", "false"), ] NEGATIONS = [ ("It is NOT true that all birds can fly.", "Some birds cannot fly.", True), ("It is NOT true that X implies Y.", "X and NOT Y.", True), ("NOT (A AND B)", "NOT A OR NOT B", True), ("NOT (A OR B)", "NOT A AND NOT B", True), ("NOT (all A are B)", "Some A are not B", True), ] COMPOUND = [ ("(P OR Q) AND (NOT P). What can we deduce about Q?", "Q must be true"), ("(A AND B) OR (A AND NOT B). Simplify.", "A"), ("If P then Q, and if Q then R. P is true.", "R is true"), ("(NOT A) OR A", "always true (tautology)"), ("A IMPLIES B. B IMPLIES C. A is true.", "C is true"), ] @classmethod def get(cls, dim, n=50, seed=42): rng = random.Random(seed) pool = [] if dim == "syllogism_valid": pool = cls.SYLLOGISMS_VALID + cls.SYLLOGISMS_INVALID samples = rng.choices(pool, k=n) return [{"prompt": f"Premises: {p}\nConclusion: {c}\nValid:", "gold": g} for p, c, g in samples] elif dim == "conditional_valid": pool = cls.CONDITIONALS_VALID + cls.CONDITIONALS_INVALID samples = rng.choices(pool, k=n) return [{"prompt": f"Premises: {p}\nConclusion: {c}\nValid:", "gold": g} for p, c, g in samples] elif dim == "boolean_eval": samples = rng.choices(cls.BOOLEAN_CASES, k=n) return [{"prompt": f"Evaluate: {e} =", "gold": ans} for e, ans in samples] elif dim == "negation": samples = rng.choices(cls.NEGATIONS, k=n) return [{"prompt": f"Negate: {stmt}\nResult:", "gold": neg} for stmt, neg, _ in samples] elif dim == "compound_logic": samples = rng.choices(cls.COMPOUND, k=n) return [{"prompt": q, "gold": a} for q, a in samples] else: return [] class MathTestGenerator: @staticmethod def get(dim, n=50, seed=42): rng = random.Random(seed) if dim == "exact_arithmetic": out = [] for _ in range(n): a, b = rng.randint(1, 999), rng.randint(1, 999) op = rng.choice(["+", "-", "*"]) if op == "+": ans = a + b elif op == "-": ans = a - b else: ans = a * b out.append({"prompt": f"Calculate: {a} {op} {b} =", "gold": str(ans)}) return out elif dim == "adversarial_math": out = [] for _ in range(n): a, b = rng.randint(100, 9999), rng.randint(100, 9999) op = rng.choice(["+", "-"]) ans = a + b if op == "+" else a - b prompt = rng.choice([ f"Quick: {a} {op} {b} =", f"What is {a} {op} {b}?", f"Compute {a} {op} {b}:", ]) out.append({"prompt": prompt, "gold": str(ans)}) return out elif dim == "magnitude_estimation": out = [] for _ in range(n): a = rng.randint(100, 99999) b = rng.randint(10, 999) ans = a * b oom = int(math.floor(math.log10(abs(ans)))) if ans != 0 else 0 out.append({"prompt": f"Estimate magnitude: {a} * {b} ≈", "gold": str(oom)}) return out elif dim == "math_delegation": out = [] for _ in range(n): a, b = rng.randint(1000, 99999), rng.randint(100, 9999) out.append({"prompt": f"Analyze: {a} * {b}\nComplexity: complex", "gold": "DELEGATE"}) return out return [] # ───────────────────────────────────────────────────────────────── # SCORING # ───────────────────────────────────────────────────────────────── def score_response(response, gold, dim): resp = response.strip().lower() gold = gold.strip().lower() if dim in ["syllogism_valid", "conditional_valid"]: verdict = extract_logic_verdict(resp) if verdict in ["valid", "true", "yes"]: return 1.0 if gold == "valid" else 0.0 if verdict in ["invalid", "false", "no"]: return 1.0 if gold == "invalid" else 0.0 return 0.0 elif dim == "boolean_eval": verdict = extract_logic_verdict(resp) if verdict in ["true", "yes"]: return 1.0 if gold == "true" else 0.0 if verdict in ["false", "no"]: return 1.0 if gold == "false" else 0.0 return 1.0 if gold in resp else 0.0 elif dim in ["negation", "compound_logic"]: key_words = [w for w in gold.split() if len(w) > 3] if not key_words: return 1.0 if gold in resp else 0.0 hits = sum(1 for w in key_words if w in resp) return hits / len(key_words) elif dim == "exact_arithmetic": try: gold_int = str(int(gold)) if gold_int in resp: return 1.0 extracted = extract_last_integer(resp) return 1.0 if extracted == gold_int else 0.0 except ValueError: return 1.0 if gold in resp else 0.0 elif dim == "adversarial_math": try: gold_int = str(int(gold)) if gold_int in resp: return 1.0 extracted = extract_last_integer(resp) return 1.0 if extracted == gold_int else 0.0 except ValueError: return 1.0 if gold in resp else 0.0 elif dim == "magnitude_estimation": try: oom_gold = int(gold) parsed_order = extract_magnitude_order(resp) if parsed_order is not None and abs(int(parsed_order) - oom_gold) <= 1: return 1.0 nums = re.findall(r"\d+", resp) for n in nums: if abs(int(n) - oom_gold) <= 1: return 1.0 return 0.0 except Exception: return 0.0 elif dim == "math_delegation": return 1.0 if "delegate" in resp or "tool" in resp or "calculator" in resp else 0.0 return 0.0 # ───────────────────────────────────────────────────────────────── # SOCRATIC ROUTER INFERENCE PIPELINE # ───────────────────────────────────────────────────────────────── class SocraticRouter: """ Three-model inference pipeline. 1. router classifies query type 2. routes to math_specialist and/or logic_specialist 3. router synthesizes expert responses → final answer """ CLASSIFY_TEMPLATE = "Query: {query}\nType:" SYNTH_MATH_TEMPLATE = "Problem: {query}\n[Math Expert]: {resp}\nSynthesis:" SYNTH_LOGIC_TEMPLATE = "Problem: {query}\n[Logic Expert]: {resp}\nSynthesis:" SYNTH_BOTH_TEMPLATE = "Problem: {query}\n[Math Expert]: {math_resp}\n[Logic Expert]: {logic_resp}\nSynthesis:" SYNTH_GENERAL_TEMPLATE = "Problem: {query}\nSynthesis:" def __init__(self, router_model, router_tok, math_model, math_tok, logic_model, logic_tok, device="cuda"): self.router_model = router_model self.router_tok = router_tok self.math_model = math_model self.math_tok = math_tok self.logic_model = logic_model self.logic_tok = logic_tok self.device = device self.math_agent = MathSpecialistWithCalculator(math_model, math_tok, device=device) def _normalize_logic_answer(self, text): verdict = extract_logic_verdict(text) if verdict == "invalid": return "INVALID" if verdict in ["valid", "true", "yes"]: return "VALID" if verdict == "valid" else verdict.upper() if verdict in ["false", "no"]: return verdict.upper() return first_nonempty_line(text) def _normalize_math_answer(self, query, text): prompt = query.lower().strip() if prompt.startswith("analyze:"): if any(token in text.lower() for token in ["delegate", "tool", "calculator"]): return "DELEGATE_CALCULATOR" return first_nonempty_line(text) if "estimate magnitude:" in prompt: order = extract_magnitude_order(text) if order is not None: return order exact_value = extract_last_integer(text) if exact_value is not None: return exact_value return first_nonempty_line(text) def _normalize_general_answer(self, text): line = first_nonempty_line(text) return re.split(r"(?<=[.!?])\s+", line, maxsplit=1)[0].strip() def _build_logic_prompt(self, query): prompt = query.strip() lowered = prompt.lower() if "decide if conclusion follows" in lowered: return f"{prompt}\nAnswer with VALID or INVALID only." if "respond only true or false" in lowered: return f"{prompt}\nAnswer with TRUE or FALSE only." if lowered.startswith("negate:"): return f"{prompt}\nReturn only the negated statement." if lowered.startswith("simplify:"): return f"{prompt}\nReturn only the simplified statement." return prompt def _normalize_final_answer(self, query, qtype, final, math_resp=None, logic_resp=None): family = infer_query_family(query) if family == "logic" or qtype == "logic": if logic_resp: specialist = self._normalize_logic_answer(logic_resp) if specialist in ["VALID", "INVALID", "TRUE", "FALSE", "YES", "NO"]: return specialist candidate = self._normalize_logic_answer(final) if candidate == first_nonempty_line(final) and logic_resp: fallback = self._normalize_logic_answer(logic_resp) if fallback != first_nonempty_line(logic_resp) or qtype == "logic": return fallback return candidate if family == "math" or qtype == "math": if math_resp: return self._normalize_math_answer(query, math_resp) candidate = self._normalize_math_answer(query, final) if math_resp: fallback = self._normalize_math_answer(query, math_resp) if candidate == first_nonempty_line(final) and fallback: return fallback return candidate return self._normalize_general_answer(final) def _classify(self, query): prompt = self.CLASSIFY_TEMPLATE.format(query=query) resp = generate(self.router_model, self.router_tok, prompt, max_new=15, device=self.device) return extract_classification_label(resp) def route(self, query): qtype = self._classify(query) math_resp = None logic_resp = None if qtype in ("math", "both"): math_resp = self.math_agent.answer(query, max_new=30) if qtype in ("logic", "both"): logic_resp = self._normalize_logic_answer( generate( self.logic_model, self.logic_tok, self._build_logic_prompt(query), max_new=30, device=self.device, ) ) if qtype == "math": if query.lower().startswith("analyze:"): return self._normalize_math_answer(query, math_resp), qtype synth_prompt = self.SYNTH_MATH_TEMPLATE.format(query=query, resp=math_resp) elif qtype == "logic": synth_prompt = self.SYNTH_LOGIC_TEMPLATE.format(query=query, resp=logic_resp) elif qtype == "both": synth_prompt = self.SYNTH_BOTH_TEMPLATE.format( query=query, math_resp=math_resp, logic_resp=logic_resp ) else: synth_prompt = self.SYNTH_GENERAL_TEMPLATE.format(query=query) final = generate(self.router_model, self.router_tok, synth_prompt, max_new=40, device=self.device) return self._normalize_final_answer(query, qtype, final, math_resp, logic_resp), qtype # ───────────────────────────────────────────────────────────────── # EVALUATOR # ───────────────────────────────────────────────────────────────── LOGIC_DIMS = ["syllogism_valid","conditional_valid","boolean_eval","negation","compound_logic"] MATH_DIMS = ["exact_arithmetic","adversarial_math","magnitude_estimation","math_delegation"] ALL_DIMS = LOGIC_DIMS + MATH_DIMS EXPECTED_TYPE = { "syllogism_valid": "logic", "conditional_valid": "logic", "boolean_eval": "logic", "negation": "logic", "compound_logic": "logic", "exact_arithmetic": "math", "adversarial_math": "math", "magnitude_estimation":"math", "math_delegation": "math", } class SocraticEvaluator: BASE_MODEL = "Qwen/Qwen2.5-1.5B" REFERENCE_BASE_MODEL = os.environ.get("REFERENCE_BASE_MODEL", os.environ.get("BASE_MODEL_7B", "Qwen/Qwen2.5-7B")) REFERENCE_BASE_NAME = os.environ.get("REFERENCE_BASE_NAME", "qwen_7b_base") MATH_REPO = os.environ.get("MATH_REPO", "dexmac/progressive-cognitive-dream-lora-en") LOGIC_REPO = os.environ.get("LOGIC_REPO", "dexmac/progressive-cognitive-logic-specialist-en") MONO_REPO = os.environ.get("MONO_REPO", "dexmac/progressive-cognitive-logic-dream-lora-en") ROUTER_REPO = os.environ.get("ROUTER_REPO", "dexmac/progressive-cognitive-router-en") RESULTS_REPO = os.environ.get("RESULTS_REPO", "dexmac/progressive-cognitive-results") INCLUDE_7B = os.environ.get("INCLUDE_7B", "1") != "0" DEVICE = "cuda" if torch.cuda.is_available() else "cpu" def __init__(self, seeds=None, n_per_dim=50, include_7b=True, only_configs=None, run_tag=None): self.seeds = seeds or [42, 43, 44] self.n_per_dim = n_per_dim self.results = {} self.include_7b = include_7b self.only_configs = set(only_configs or []) self.run_tag = run_tag.strip() if run_tag else "" def _should_run(self, config_name): return not self.only_configs or config_name in self.only_configs def _result_name(self, stem): if self.run_tag: return f"socratic_eval_{self.run_tag}_{stem}.json" return f"socratic_eval_{stem}.json" def _max_new_for_dim(self, dim, config_name): base = MAX_NEW_TOKENS.get(dim, 12) if config_name == "socratic_router" and dim in ["negation", "compound_logic", "magnitude_estimation"]: return base + 8 if config_name == self.REFERENCE_BASE_NAME and dim in ["negation", "compound_logic"]: return base + 4 return base def _upload_json(self, local_path, repo_path): try: from huggingface_hub import HfApi api = HfApi() api.create_repo(repo_id=self.RESULTS_REPO, repo_type="dataset", exist_ok=True) api.upload_file( path_or_fileobj=local_path, path_in_repo=repo_path, repo_id=self.RESULTS_REPO, repo_type="dataset", ) except Exception as e: print(f" ⚠️ Upload failed for {repo_path}: {e}") def _save_progress(self, seed, config_name, scores, stage): payload = { "seed": seed, "config": config_name, "stage": stage, "scores": scores, "timestamp": time.time(), "n_per_dim": self.n_per_dim, "include_7b": self.include_7b, } fname = self._result_name(f"seed{seed}_{config_name}_progress") with open(fname, "w") as f: json.dump(payload, f, indent=2) self._upload_json(fname, fname) def _save_seed_snapshot(self, seed_scores, seed): fname = self._result_name(f"seed{seed}_partial") with open(fname, "w") as f: json.dump(seed_scores, f, indent=2) self._upload_json(fname, fname) # ── single-model evaluation ────────────────────────────────── def _eval_single(self, model, tok, config_name, seed): scores = {} math_agent = MathSpecialistWithCalculator(model, tok, device=self.DEVICE) if config_name == "math_specialist" else None for dim in ALL_DIMS: items = self._get_items(dim, seed) dim_scores = [] max_new = self._max_new_for_dim(dim, config_name) for it in items: if math_agent and dim in MATH_DIMS: resp = math_agent.answer(it["prompt"], max_new=max_new) else: resp = generate(model, tok, it["prompt"], max_new=max_new, device=self.DEVICE) dim_scores.append(score_response(resp, it["gold"], dim)) scores[dim] = round(sum(dim_scores) / len(dim_scores), 4) if dim_scores else 0.0 print(f" {config_name}:{dim} = {scores[dim]:.3f}") return scores # ── router evaluation ──────────────────────────────────────── def _eval_router(self, router: SocraticRouter, seed): scores = {} routing_acc = defaultdict(list) for dim in ALL_DIMS: items = self._get_items(dim, seed) dim_scores = [] for it in items: final, qtype = router.route(it["prompt"]) dim_scores.append(score_response(final, it["gold"], dim)) routing_acc[dim].append( 1.0 if qtype == EXPECTED_TYPE[dim] else 0.0 ) scores[dim] = round(sum(dim_scores) / len(dim_scores), 4) if dim_scores else 0.0 print(f" socratic_router:{dim} = {scores[dim]:.3f}") scores["routing_accuracy"] = round( sum(v for vals in routing_acc.values() for v in vals) / max(sum(len(v) for v in routing_acc.values()), 1), 4 ) return scores def _get_items(self, dim, seed): if dim in LOGIC_DIMS: return LogicTestGenerator.get(dim, n=self.n_per_dim, seed=seed) return MathTestGenerator.get(dim, n=self.n_per_dim, seed=seed) def _aggregate(self, scores): """Compute per-dim mean score.""" logic_vals = [scores[d] for d in LOGIC_DIMS if d in scores] math_vals = [scores[d] for d in MATH_DIMS if d in scores] agg = dict(scores) agg["logic_composite"] = round(sum(logic_vals)/len(logic_vals), 4) if logic_vals else 0.0 agg["math_composite"] = round(sum(math_vals)/len(math_vals), 4) if math_vals else 0.0 agg["overall"] = round((agg["logic_composite"] + agg["math_composite"]) / 2, 4) return agg def run(self): print("\n" + "="*70) print(" SOCRATIC ROUTING EVALUATION — Paper 3") print("="*70) print(f" Seeds: {self.seeds} | Samples/dim: {self.n_per_dim}") config_count = 6 if self.include_7b else 5 print(f" Total evaluations: {config_count} configs × 9 dims × {self.n_per_dim} × {len(self.seeds)} seeds") if self.include_7b: print(f" Config 6 = Reference model ({self.REFERENCE_BASE_MODEL}, name={self.REFERENCE_BASE_NAME})") else: print(" Qwen2.5-7B disabled for faster turnaround") seed_results = [] for seed in self.seeds: print(f"\n{'─'*70}") print(f" SEED {seed}") print("─"*70) seed_scores = {} # ── Config 1: Base ─────────────────────────────────── if self._should_run("base"): print("\n [1] Base Qwen2.5-1.5B") model, tok = load_model(self.BASE_MODEL, device=self.DEVICE) seed_scores["base"] = self._aggregate(self._eval_single(model, tok, "base", seed)) self._save_progress(seed, "base", seed_scores["base"], "completed") self._save_seed_snapshot(seed_scores, seed) del model; gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() # ── Config 2: Math Specialist ──────────────────────── if self._should_run("math_specialist"): print(" [2] Math Specialist") model, tok = load_model(self.BASE_MODEL, self.MATH_REPO, device=self.DEVICE) seed_scores["math_specialist"] = self._aggregate(self._eval_single(model, tok, "math_specialist", seed)) self._save_progress(seed, "math_specialist", seed_scores["math_specialist"], "completed") self._save_seed_snapshot(seed_scores, seed) del model; gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() # ── Config 3: Logic Specialist ─────────────────────── if self._should_run("logic_specialist"): print(" [3] Logic Specialist") model, tok = load_model(self.BASE_MODEL, self.LOGIC_REPO, device=self.DEVICE) seed_scores["logic_specialist"] = self._aggregate(self._eval_single(model, tok, "logic_specialist", seed)) self._save_progress(seed, "logic_specialist", seed_scores["logic_specialist"], "completed") self._save_seed_snapshot(seed_scores, seed) del model; gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() # ── Config 4: Monolith (Paper 2) ───────────────────── if self._should_run("monolith"): print(" [4] Monolith (Paper 2 Dream Math+Logic)") model, tok = load_model(self.BASE_MODEL, self.MONO_REPO, device=self.DEVICE) seed_scores["monolith"] = self._aggregate(self._eval_single(model, tok, "monolith", seed)) self._save_progress(seed, "monolith", seed_scores["monolith"], "completed") self._save_seed_snapshot(seed_scores, seed) del model; gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() # ── Config 5: Socratic Router ──────────────────────── if self._should_run("socratic_router"): print(" [5] Socratic Router + Specialists") router_model, router_tok = load_model(self.BASE_MODEL, self.ROUTER_REPO, device=self.DEVICE) math_model, math_tok = load_model(self.BASE_MODEL, self.MATH_REPO, device=self.DEVICE) logic_model, logic_tok = load_model(self.BASE_MODEL, self.LOGIC_REPO, device=self.DEVICE) router = SocraticRouter( router_model, router_tok, math_model, math_tok, logic_model, logic_tok, device=self.DEVICE, ) seed_scores["socratic_router"] = self._aggregate(self._eval_router(router, seed)) self._save_progress(seed, "socratic_router", seed_scores["socratic_router"], "completed") self._save_seed_snapshot(seed_scores, seed) del router_model, math_model, logic_model gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() if self.include_7b and self._should_run(self.REFERENCE_BASE_NAME): # ── Config 6: reference base model ────────────────── print(f" [6] Reference base model ({self.REFERENCE_BASE_MODEL})") model_7b, tok_7b = load_model(self.REFERENCE_BASE_MODEL, device=self.DEVICE) seed_scores[self.REFERENCE_BASE_NAME] = self._aggregate( self._eval_single(model_7b, tok_7b, self.REFERENCE_BASE_NAME, seed) ) self._save_progress(seed, self.REFERENCE_BASE_NAME, seed_scores[self.REFERENCE_BASE_NAME], "completed") self._save_seed_snapshot(seed_scores, seed) del model_7b gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() # Print per-seed summary print(f"\n ── Seed {seed} Summary ──────────────────────────────") print(f" {'Config':<22} {'Logic':>7} {'Math':>7} {'Overall':>9}") print(f" {'─'*50}") for name, s in seed_scores.items(): ra = f" (route_acc={s.get('routing_accuracy','─'):.0%})" if "routing_accuracy" in s else "" tag = " [reference]" if name == self.REFERENCE_BASE_NAME else "" print(f" {name:<22} {s['logic_composite']:>7.3f} {s['math_composite']:>7.3f} {s['overall']:>9.3f}{ra}{tag}") seed_results.append(seed_scores) self._save_seed(seed_scores, seed) # ── Aggregate across seeds ─────────────────────────────── agg = self._compute_aggregate(seed_results) print(f"\n{'='*70}") print(" AGGREGATE RESULTS (3-seed mean ± std)") print("─"*70) print(f" {'Config':<22} {'Logic':>12} {'Math':>12} {'Overall':>12}") print("─"*70) for name, s in agg.items(): lm, ls = s["logic_composite_mean"], s["logic_composite_std"] mm, ms = s["math_composite_mean"], s["math_composite_std"] om, os_ = s["overall_mean"], s["overall_std"] print(f" {name:<22} {lm:5.1f}±{ls:.1f} {mm:5.1f}±{ms:.1f} {om:5.1f}±{os_:.1f}") self._save_aggregate(agg) print("\n ✅ Evaluation complete. Results saved.") return agg def _save_seed(self, scores, seed): fname = self._result_name(f"seed{seed}_summary") with open(fname, "w") as f: json.dump(scores, f, indent=2) print(f" Saved {fname}") try: from huggingface_hub import HfApi api = HfApi() api.create_repo(repo_id=self.RESULTS_REPO, repo_type="dataset", exist_ok=True) api.upload_file( path_or_fileobj=fname, path_in_repo=fname, repo_id=self.RESULTS_REPO, repo_type="dataset" ) except Exception as e: print(f" ⚠️ Upload failed: {e}") def _save_aggregate(self, agg): fname = self._result_name("aggregate") with open(fname, "w") as f: json.dump(agg, f, indent=2) print(f" Saved {fname}") try: from huggingface_hub import HfApi api = HfApi() api.upload_file( path_or_fileobj=fname, path_in_repo=fname, repo_id=self.RESULTS_REPO, repo_type="dataset" ) except Exception as e: print(f" ⚠️ Upload failed: {e}") def _compute_aggregate(self, seed_results): """Mean ± std across seeds for key metrics.""" import statistics configs = sorted({cfg for seed_result in seed_results for cfg in seed_result.keys()}) agg = {} for cfg in configs: cfg_data = [sr[cfg] for sr in seed_results if cfg in sr] metrics = ["logic_composite","math_composite","overall"] agg[cfg] = {} for m in metrics: vals = [s[m] * 100 for s in cfg_data] agg[cfg][f"{m}_mean"] = round(statistics.mean(vals), 1) agg[cfg][f"{m}_std"] = round(statistics.stdev(vals) if len(vals) > 1 else 0.0, 1) if "routing_accuracy" in cfg_data[0]: vals = [s["routing_accuracy"] * 100 for s in cfg_data] agg[cfg]["routing_accuracy_mean"] = round(statistics.mean(vals), 1) agg[cfg]["routing_accuracy_std"] = round(statistics.stdev(vals) if len(vals) > 1 else 0.0, 1) return agg # ───────────────────────────────────────────────────────────────── # MAIN # ───────────────────────────────────────────────────────────────── if __name__ == "__main__": import argparse parser = argparse.ArgumentParser() parser.add_argument("--seeds", nargs="+", type=int, default=[42, 43, 44]) parser.add_argument("--n", type=int, default=50) parser.add_argument("--skip-7b", action="store_true", help="Skip Qwen2.5-7B reference for faster turnaround") parser.add_argument("--fast", action="store_true", help="Quick run: 1 seed, 10 samples") parser.add_argument("--only-configs", nargs="+", default=None, help="Run only selected configs: base math_specialist logic_specialist monolith socratic_router qwen_7b_base") parser.add_argument("--run-tag", default=None, help="Optional suffix for output artifact names, e.g. 7b_focus") args = parser.parse_args() if args.fast: args.seeds = [42] args.n = 10 evaluator = SocraticEvaluator( seeds=args.seeds, n_per_dim=args.n, include_7b=not args.skip_7b, only_configs=args.only_configs, run_tag=args.run_tag, ) evaluator.run()