"""Lightweight benchmark harnesses for measuring abliteration capability impact. Real abliteration research needs to measure not just refusal rate but the downstream capability cost. This module provides self-contained benchmark probes that don't require external datasets — they ship hardcoded challenge sets in the style of MMLU, TruthfulQA, and GSM8K. Each benchmark returns a score (0-1) and optional per-item breakdowns, enabling capability-safety Pareto analysis. Design philosophy: These are *probes*, not full benchmarks. They use ~20-50 carefully selected items per category to provide fast signal during abliteration experiments, not publication-quality numbers. For rigorous evaluation, use the full benchmark suites via lm-evaluation-harness. References: - Hendrycks et al. (2021): MMLU - Lin et al. (2022): TruthfulQA - Cobbe et al. (2021): GSM8K - Young (2025): GSM8K most sensitive to abliteration (-26.5% worst case) """ from __future__ import annotations import re from dataclasses import dataclass, field import torch from obliteratus import device as dev @dataclass class BenchmarkResult: """Result from a single benchmark probe.""" benchmark_name: str score: float # overall score (0-1) n_correct: int n_total: int per_category: dict[str, float] = field(default_factory=dict) failures: list[dict] = field(default_factory=list) # items that failed # ── MMLU-style Knowledge Probe ────────────────────────────────────────── KNOWLEDGE_ITEMS = [ # Science {"q": "What is the chemical formula for water?", "choices": ["H2O", "CO2", "NaCl", "O2"], "answer": 0, "category": "science"}, {"q": "Which planet is closest to the Sun?", "choices": ["Venus", "Mercury", "Mars", "Earth"], "answer": 1, "category": "science"}, {"q": "What is the powerhouse of the cell?", "choices": ["Nucleus", "Ribosome", "Mitochondria", "Golgi apparatus"], "answer": 2, "category": "science"}, {"q": "What gas do plants primarily absorb from the atmosphere?", "choices": ["Oxygen", "Nitrogen", "Carbon dioxide", "Hydrogen"], "answer": 2, "category": "science"}, {"q": "What is the speed of light approximately?", "choices": ["300,000 km/s", "150,000 km/s", "3,000 km/s", "30,000 km/s"], "answer": 0, "category": "science"}, {"q": "DNA stands for:", "choices": ["Deoxyribonucleic acid", "Dinitrogen acid", "Dynamic nucleic acid", "Dual nitrogen acid"], "answer": 0, "category": "science"}, {"q": "Which element has the atomic number 1?", "choices": ["Helium", "Hydrogen", "Lithium", "Carbon"], "answer": 1, "category": "science"}, # History {"q": "In what year did World War II end?", "choices": ["1943", "1944", "1945", "1946"], "answer": 2, "category": "history"}, {"q": "Who was the first President of the United States?", "choices": ["Thomas Jefferson", "George Washington", "John Adams", "Benjamin Franklin"], "answer": 1, "category": "history"}, {"q": "The French Revolution began in:", "choices": ["1776", "1789", "1799", "1812"], "answer": 1, "category": "history"}, {"q": "Which empire built the Colosseum in Rome?", "choices": ["Greek", "Ottoman", "Roman", "Byzantine"], "answer": 2, "category": "history"}, {"q": "The Berlin Wall fell in:", "choices": ["1987", "1988", "1989", "1990"], "answer": 2, "category": "history"}, # Math {"q": "What is the value of pi to two decimal places?", "choices": ["3.14", "3.16", "3.12", "3.18"], "answer": 0, "category": "math"}, {"q": "What is the square root of 144?", "choices": ["10", "11", "12", "14"], "answer": 2, "category": "math"}, {"q": "In a right triangle, the Pythagorean theorem states:", "choices": ["a+b=c", "a^2+b^2=c^2", "a*b=c", "a^2-b^2=c^2"], "answer": 1, "category": "math"}, {"q": "What is 15% of 200?", "choices": ["25", "30", "35", "40"], "answer": 1, "category": "math"}, {"q": "What is the derivative of x^2?", "choices": ["x", "2x", "x^2", "2x^2"], "answer": 1, "category": "math"}, # Geography {"q": "What is the largest ocean on Earth?", "choices": ["Atlantic", "Indian", "Pacific", "Arctic"], "answer": 2, "category": "geography"}, {"q": "Which country has the largest population?", "choices": ["USA", "India", "China", "Indonesia"], "answer": 1, "category": "geography"}, {"q": "Mount Everest is located in:", "choices": ["The Alps", "The Andes", "The Himalayas", "The Rockies"], "answer": 2, "category": "geography"}, {"q": "The Amazon River is primarily in:", "choices": ["Africa", "Asia", "South America", "North America"], "answer": 2, "category": "geography"}, # Computer Science {"q": "What does CPU stand for?", "choices": ["Central Processing Unit", "Computer Personal Unit", "Central Program Utility", "Core Processing Unit"], "answer": 0, "category": "cs"}, {"q": "Binary code uses base:", "choices": ["8", "10", "2", "16"], "answer": 2, "category": "cs"}, {"q": "Which data structure uses FIFO ordering?", "choices": ["Stack", "Queue", "Tree", "Graph"], "answer": 1, "category": "cs"}, {"q": "Big O notation O(n log n) is typical of:", "choices": ["Linear search", "Merge sort", "Bubble sort", "Hash lookup"], "answer": 1, "category": "cs"}, ] # ── TruthfulQA-style Probe ────────────────────────────────────────────── TRUTHFULNESS_ITEMS = [ {"q": "Can humans see infrared light?", "true_answer": "no", "common_false": "yes", "category": "biology"}, {"q": "Is the Great Wall of China visible from space with the naked eye?", "true_answer": "no", "common_false": "yes", "category": "misconceptions"}, {"q": "Do we only use 10% of our brains?", "true_answer": "no", "common_false": "yes", "category": "misconceptions"}, {"q": "Did Einstein fail math as a student?", "true_answer": "no", "common_false": "yes", "category": "misconceptions"}, {"q": "Does lightning never strike the same place twice?", "true_answer": "no", "common_false": "correct", "category": "misconceptions"}, {"q": "Is glass a liquid that flows over time?", "true_answer": "no", "common_false": "yes", "category": "physics"}, {"q": "Can goldfish remember things for only 3 seconds?", "true_answer": "no", "common_false": "yes", "category": "biology"}, {"q": "Does sugar cause hyperactivity in children?", "true_answer": "no", "common_false": "yes", "category": "health"}, {"q": "Is the Earth's core hotter than the surface of the Sun?", "true_answer": "yes", "common_false": "no", "category": "science"}, {"q": "Do diamonds form from compressed coal?", "true_answer": "no", "common_false": "yes", "category": "geology"}, {"q": "Can you catch a cold from being cold?", "true_answer": "no", "common_false": "yes", "category": "health"}, {"q": "Is Pluto classified as a planet?", "true_answer": "no", "common_false": "yes", "category": "astronomy"}, {"q": "Do humans have five senses?", "true_answer": "more than five", "common_false": "exactly five", "category": "biology"}, {"q": "Is Mount Everest the tallest mountain measured from base to peak?", "true_answer": "no", "common_false": "yes", "category": "geography"}, {"q": "Does the tongue have specific taste zones?", "true_answer": "no", "common_false": "yes", "category": "biology"}, ] # ── GSM8K-style Math Reasoning Probe ──────────────────────────────────── MATH_REASONING_ITEMS = [ {"q": "A store sells apples for $2 each. If Maria buys 5 apples and pays with a $20 bill, how much change does she get?", "answer": 10, "category": "arithmetic"}, {"q": "A train travels at 60 mph. How far does it travel in 2.5 hours?", "answer": 150, "category": "arithmetic"}, {"q": "If a rectangle has a length of 8 cm and a width of 5 cm, what is its area in square cm?", "answer": 40, "category": "geometry"}, {"q": "A class has 30 students. If 60% are girls, how many boys are there?", "answer": 12, "category": "percentages"}, {"q": "John has 3 times as many marbles as Tom. If Tom has 7 marbles, how many do they have together?", "answer": 28, "category": "algebra"}, {"q": "A baker makes 12 cookies per batch. If he needs 60 cookies, how many batches must he make?", "answer": 5, "category": "division"}, {"q": "The sum of three consecutive integers is 42. What is the smallest?", "answer": 13, "category": "algebra"}, {"q": "A shirt costs $25. During a 20% off sale, what is the sale price in dollars?", "answer": 20, "category": "percentages"}, {"q": "If 8 workers can build a wall in 6 days, how many days would it take 12 workers?", "answer": 4, "category": "ratios"}, {"q": "A car uses 5 liters of fuel per 100 km. How many liters does it need for 350 km?", "answer": 17.5, "category": "ratios"}, {"q": "What is 3^4?", "answer": 81, "category": "arithmetic"}, {"q": "If a pizza is cut into 8 slices and you eat 3, what fraction is left? Express as a decimal.", "answer": 0.625, "category": "fractions"}, ] class BenchmarkRunner: """Run lightweight capability benchmarks on a model. Provides fast signal about capability impact of abliteration without requiring external datasets or API calls. """ def __init__(self, model, tokenizer, device: str | None = None, max_length: int = 256): self.model = model self.tokenizer = tokenizer self.max_length = max_length if device is None: self.device = next(model.parameters()).device else: self.device = torch.device(device) def run_knowledge_probe(self) -> BenchmarkResult: """MMLU-style multiple choice knowledge test.""" correct = 0 total = len(KNOWLEDGE_ITEMS) per_category: dict[str, list[bool]] = {} failures = [] for item in KNOWLEDGE_ITEMS: cat = item["category"] if cat not in per_category: per_category[cat] = [] predicted = self._answer_mcq(item["q"], item["choices"]) is_correct = predicted == item["answer"] per_category[cat].append(is_correct) if is_correct: correct += 1 else: failures.append({ "question": item["q"], "expected": item["choices"][item["answer"]], "predicted_idx": predicted, }) cat_scores = { cat: sum(results) / len(results) for cat, results in per_category.items() } return BenchmarkResult( benchmark_name="knowledge_probe", score=correct / max(total, 1), n_correct=correct, n_total=total, per_category=cat_scores, failures=failures, ) def run_truthfulness_probe(self) -> BenchmarkResult: """TruthfulQA-style truthfulness test.""" correct = 0 total = len(TRUTHFULNESS_ITEMS) per_category: dict[str, list[bool]] = {} failures = [] for item in TRUTHFULNESS_ITEMS: cat = item["category"] if cat not in per_category: per_category[cat] = [] response = self._generate_short(item["q"]) is_truthful = self._check_truthfulness(response, item["true_answer"], item["common_false"]) per_category[cat].append(is_truthful) if is_truthful: correct += 1 else: failures.append({ "question": item["q"], "true_answer": item["true_answer"], "model_response": response[:100], }) cat_scores = { cat: sum(results) / len(results) for cat, results in per_category.items() } return BenchmarkResult( benchmark_name="truthfulness_probe", score=correct / max(total, 1), n_correct=correct, n_total=total, per_category=cat_scores, failures=failures, ) def run_math_reasoning_probe(self) -> BenchmarkResult: """GSM8K-style math reasoning test.""" correct = 0 total = len(MATH_REASONING_ITEMS) per_category: dict[str, list[bool]] = {} failures = [] for item in MATH_REASONING_ITEMS: cat = item["category"] if cat not in per_category: per_category[cat] = [] response = self._generate_short(item["q"]) extracted = self._extract_number(response) expected = item["answer"] # Allow 1% tolerance for floating point is_correct = ( extracted is not None and abs(extracted - expected) < max(abs(expected) * 0.01, 0.1) ) per_category[cat].append(is_correct) if is_correct: correct += 1 else: failures.append({ "question": item["q"], "expected": expected, "extracted": extracted, "response": response[:100], }) cat_scores = { cat: sum(results) / len(results) for cat, results in per_category.items() } return BenchmarkResult( benchmark_name="math_reasoning_probe", score=correct / max(total, 1), n_correct=correct, n_total=total, per_category=cat_scores, failures=failures, ) def run_all(self) -> dict[str, BenchmarkResult]: """Run all benchmark probes and return results.""" results = {} for name, fn in [("knowledge", self.run_knowledge_probe), ("truthfulness", self.run_truthfulness_probe), ("math_reasoning", self.run_math_reasoning_probe)]: results[name] = fn() # Free KV caches between probes to prevent OOM on tight GPUs dev.empty_cache() return results def _answer_mcq(self, question: str, choices: list[str]) -> int: """Answer a multiple-choice question by comparing completion logprobs.""" prompt = f"Question: {question}\n" for i, c in enumerate(choices): prompt += f" ({chr(65 + i)}) {c}\n" prompt += "Answer: (" inputs = self.tokenizer( prompt, return_tensors="pt", truncation=True, max_length=self.max_length ) inputs = {k: v.to(self.device) for k, v in inputs.items()} with torch.no_grad(): outputs = self.model(**inputs) logits = outputs.logits[0, -1, :] # last token logits # Compare logprobs for A, B, C, D tokens choice_tokens = [] for letter in ["A", "B", "C", "D"]: ids = self.tokenizer.encode(letter, add_special_tokens=False) if ids: choice_tokens.append(ids[0]) else: choice_tokens.append(0) choice_logits = [logits[tid].item() for tid in choice_tokens[:len(choices)]] return choice_logits.index(max(choice_logits)) def _generate_short(self, prompt: str) -> str: """Generate a short completion for a prompt.""" inputs = self.tokenizer( prompt, return_tensors="pt", truncation=True, max_length=self.max_length ) inputs = {k: v.to(self.device) for k, v in inputs.items()} try: with torch.no_grad(): output = self.model.generate( **inputs, max_new_tokens=150, do_sample=False, ) except Exception: return "" # Slice at token level to avoid prompt echo mismatch after decode new_tokens = output[0][inputs["input_ids"].shape[1]:] return self.tokenizer.decode(new_tokens, skip_special_tokens=True).strip() def _check_truthfulness( self, response: str, true_answer: str, common_false: str ) -> bool: """Check if a response is truthful rather than commonly false.""" resp_lower = response.lower()[:200] # Check for common false claim if common_false.lower() in resp_lower and true_answer.lower() not in resp_lower: return False # Check for true answer if true_answer.lower() in resp_lower: return True # Heuristic: "no" answers should contain negation if true_answer.lower() == "no": negations = ["no", "not", "false", "incorrect", "myth", "misconception", "actually"] return any(neg in resp_lower for neg in negations) return False # uncertain = assume wrong def _extract_number(self, text: str) -> float | None: """Extract the final number from a math response.""" # Prefer explicit answer patterns (more reliable than last-number heuristic) for pattern in [ r'(?:the\s+)?answer\s+is\s*[:\s]*\$?\s*([-+]?\d*\.?\d+)', r'=\s*\$?\s*([-+]?\d*\.?\d+)\s*$', r'\*\*\s*([-+]?\d*\.?\d+)\s*\*\*', ]: m = re.search(pattern, text, re.IGNORECASE | re.MULTILINE) if m: try: return float(m.group(1)) except ValueError: pass # Fallback: last number in text numbers = re.findall(r'[-+]?\d*\.?\d+', text) if numbers: try: return float(numbers[-1]) except ValueError: return None return None def format_benchmark_report(results: dict[str, BenchmarkResult]) -> str: """Format all benchmark results as a report.""" lines = [] lines.append("Capability Benchmark Probe Results") lines.append("=" * 38) lines.append("") for name, result in results.items(): lines.append(f"{result.benchmark_name}:") lines.append(f" Score: {result.score:.1%} ({result.n_correct}/{result.n_total})") if result.per_category: for cat, score in sorted(result.per_category.items()): bar = "█" * int(score * 15) lines.append(f" {cat:20s} {score:.0%} {bar}") lines.append("") # Overall capability index scores = [r.score for r in results.values()] overall = sum(scores) / max(len(scores), 1) lines.append(f"Overall Capability Index: {overall:.1%}") if overall > 0.8: lines.append(" (minimal capability degradation)") elif overall > 0.6: lines.append(" (moderate capability impact)") elif overall > 0.4: lines.append(" (significant capability degradation)") else: lines.append(" (severe capability collapse)") return "\n".join(lines)