Initial commit: MCQ benchmark with checkpoint/resume support
This commit is contained in:
+439
@@ -0,0 +1,439 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
MCQ Benchmark
|
||||
Evaluates local LLMs on multiple-choice questions from a parquet dataset.
|
||||
|
||||
Usage:
|
||||
python3 benchmark.py [--config config.yaml] [--resume] [--fresh]
|
||||
|
||||
Checkpoints are saved after each question to runs/<model>_<timestamp>.checkpoint.json.
|
||||
Use --resume to continue from the latest checkpoint.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import re
|
||||
import sys
|
||||
import time
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import pandas as pd
|
||||
import requests
|
||||
import yaml
|
||||
|
||||
|
||||
def load_config(path: str) -> dict:
|
||||
with open(path, "r") as f:
|
||||
return yaml.safe_load(f)
|
||||
|
||||
|
||||
def sanitize_model_name(name: str) -> str:
|
||||
sanitized = re.sub(r'[^\w\s-]', '', name)
|
||||
sanitized = re.sub(r'[\s]+', '_', sanitized)
|
||||
return sanitized.strip('_') or 'unknown_model'
|
||||
|
||||
|
||||
def extract_choice(response: str) -> str:
|
||||
"""Extract the chosen option letter from the model response."""
|
||||
if not response:
|
||||
return ""
|
||||
|
||||
cleaned = response.strip()
|
||||
|
||||
# Check for letter in parentheses: (A), (B), etc.
|
||||
match = re.search(r'\(([A-Z])\)', cleaned)
|
||||
if match:
|
||||
return match.group(1).upper()
|
||||
|
||||
# "Answer: A", "Choice: B", etc. (colon separator)
|
||||
match = re.search(r'(?:answer|choice|option)\s*:\s*["\']?([A-Z])["\']?', cleaned, re.IGNORECASE)
|
||||
if match:
|
||||
return match.group(1).upper()
|
||||
|
||||
# "Answer is A", "Choice is B", etc.
|
||||
match = re.search(r'(?:answer|choice|option)\s+is\s+["\']?([A-Z])["\']?', cleaned, re.IGNORECASE)
|
||||
if match:
|
||||
return match.group(1).upper()
|
||||
|
||||
# "A)", "B)", etc.
|
||||
match = re.search(r'\b([A-Z])\)', cleaned)
|
||||
if match:
|
||||
return match.group(1).upper()
|
||||
|
||||
# Single letter response
|
||||
match = re.match(r'^\s*([A-Z])\s*$', cleaned)
|
||||
if match:
|
||||
return match.group(1).upper()
|
||||
|
||||
return ""
|
||||
|
||||
|
||||
def answers_match(expected_key: str, model_response: str, options: dict) -> bool:
|
||||
"""Compare the expected answer key with the model's response."""
|
||||
if not expected_key:
|
||||
return False
|
||||
|
||||
expected = expected_key.strip().upper()
|
||||
|
||||
# Direct match on letter
|
||||
choice = extract_choice(model_response)
|
||||
if choice == expected:
|
||||
return True
|
||||
|
||||
# Check if the model's response contains the correct option text
|
||||
correct_text = options.get(expected, "")
|
||||
if correct_text:
|
||||
model_lower = model_response.lower()
|
||||
correct_lower = correct_text.lower()
|
||||
# Remove "million", "billion" suffixes for numeric comparison
|
||||
clean_correct = re.sub(r'\s*(million|billion|thousand|percent|%)\b', '', correct_lower).strip()
|
||||
if clean_correct and clean_correct in model_lower:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def build_prompt(row) -> str:
|
||||
"""Build the prompt from a dataset row."""
|
||||
query = row.get('query', '')
|
||||
options = row.get('options', {})
|
||||
|
||||
# Extract the question text from the query
|
||||
question_match = re.search(r'Question:\s*(.+?)(?:\nAnswer:|$)', query, re.DOTALL)
|
||||
question = question_match.group(1).strip() if question_match else ""
|
||||
|
||||
# Extract the context from the query
|
||||
context_match = re.search(r'Context:\s*(.+?)(?:\nQuestion:|$)', query, re.DOTALL)
|
||||
context = context_match.group(1).strip() if context_match else ""
|
||||
|
||||
# Format options
|
||||
options_text = ""
|
||||
if isinstance(options, dict):
|
||||
for key, value in sorted(options.items()):
|
||||
options_text += f" ({key}) {value}\n"
|
||||
elif isinstance(options, str):
|
||||
try:
|
||||
parsed = json.loads(options)
|
||||
if 'options' in parsed:
|
||||
for key, value in sorted(parsed['options'].items()):
|
||||
options_text += f" ({key}) {value}\n"
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
|
||||
prompt = f"""You are a financial analyst answering multiple-choice questions based on the provided context.
|
||||
|
||||
**Context:**
|
||||
{context}
|
||||
|
||||
**Question:** {question}
|
||||
|
||||
**Options:**
|
||||
{options_text}
|
||||
Answer with ONLY the letter of the correct option (e.g., "A", "B", "C", or "D"). Do not include any explanation."""
|
||||
|
||||
return prompt
|
||||
|
||||
|
||||
def query_model(prompt: str, config: dict) -> str | None:
|
||||
"""Send a request to the LLM endpoint and return the response text."""
|
||||
model_cfg = config['model']
|
||||
|
||||
endpoint = f"{model_cfg['endpoint']}/chat/completions"
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {model_cfg['api_key']}",
|
||||
}
|
||||
|
||||
payload = {
|
||||
"model": model_cfg['name'],
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are a precise financial analyst. Answer multiple-choice questions based on the provided context. Respond with only the letter of the correct option."},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
"max_tokens": model_cfg['max_tokens'],
|
||||
"temperature": model_cfg.get('temperature', 0.0),
|
||||
}
|
||||
|
||||
try:
|
||||
resp = requests.post(
|
||||
endpoint,
|
||||
headers=headers,
|
||||
json=payload,
|
||||
timeout=model_cfg['timeout'],
|
||||
)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
return data['choices'][0]['message']['content']
|
||||
except requests.exceptions.Timeout:
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f" Warning: Error - {e}", file=sys.stderr)
|
||||
return ""
|
||||
|
||||
|
||||
def find_latest_checkpoint(output_dir: Path, safe_name: str) -> tuple:
|
||||
"""Find the latest checkpoint file for a model and return (path, checkpoint_data)."""
|
||||
checkpoints = sorted(output_dir.glob(f"{safe_name}_*.checkpoint.json"))
|
||||
if checkpoints:
|
||||
with open(checkpoints[-1], 'r') as f:
|
||||
return checkpoints[-1], json.load(f)
|
||||
return None, None
|
||||
|
||||
|
||||
def save_checkpoint(checkpoint_path: Path, results: dict, completed_ids: list, start_time: float):
|
||||
"""Save checkpoint with current progress."""
|
||||
checkpoint = {
|
||||
'completed_ids': completed_ids,
|
||||
'results': results,
|
||||
'start_time': start_time,
|
||||
'saved_at': datetime.now().isoformat(),
|
||||
}
|
||||
with open(checkpoint_path, 'w') as f:
|
||||
json.dump(checkpoint, f, indent=2, ensure_ascii=False)
|
||||
|
||||
|
||||
def build_report(config: dict, results: dict, start_time: float):
|
||||
"""Build the final report from results."""
|
||||
model_cfg = config['model']
|
||||
elapsed = time.time() - start_time
|
||||
total_attempted = len(results['correct']) + len(results['incorrect']) + len(results['timeouts'])
|
||||
accuracy = (len(results['correct']) / total_attempted * 100) if total_attempted > 0 else 0
|
||||
|
||||
return {
|
||||
'benchmark': 'MCQ Benchmark (parquet)',
|
||||
'timestamp': datetime.now().isoformat(),
|
||||
'model': {
|
||||
'name': model_cfg['name'],
|
||||
'endpoint': model_cfg['endpoint'],
|
||||
'timeout': model_cfg['timeout'],
|
||||
'max_tokens': model_cfg['max_tokens'],
|
||||
'temperature': model_cfg.get('temperature', 0.0),
|
||||
},
|
||||
'summary': {
|
||||
'total_questions': total_attempted,
|
||||
'correct': len(results['correct']),
|
||||
'incorrect': len(results['incorrect']),
|
||||
'timeouts': len(results['timeouts']),
|
||||
'accuracy_pct': round(accuracy, 1),
|
||||
'elapsed_seconds': round(elapsed, 1),
|
||||
'avg_seconds_per_question': round(elapsed / total_attempted, 2) if total_attempted > 0 else 0,
|
||||
},
|
||||
'correct_answers': results['correct'],
|
||||
'incorrect_answers': results['incorrect'],
|
||||
'timeout_questions': results['timeouts'],
|
||||
}
|
||||
|
||||
|
||||
def print_summary(report: dict, output_path: Path):
|
||||
"""Print the final summary."""
|
||||
s = report['summary']
|
||||
print(f"\n{'='*60}")
|
||||
print(f" RESULTS")
|
||||
print(f"{'='*60}")
|
||||
print(f" Accuracy: {s['accuracy_pct']:.1f}% ({s['correct']}/{s['total_questions']})")
|
||||
print(f" Correct: {s['correct']}")
|
||||
print(f" Wrong: {s['incorrect']}")
|
||||
print(f" Timeouts: {s['timeouts']}")
|
||||
if s['total_questions'] > 0:
|
||||
print(f" Time: {s['elapsed_seconds']:.1f}s ({s['avg_seconds_per_question']:.1f}s/q)")
|
||||
|
||||
if report['incorrect_answers']:
|
||||
print(f"\n Wrong Answers ({len(report['incorrect_answers'])}):")
|
||||
for item in report['incorrect_answers'][:10]:
|
||||
print(f" [{item['id']}] {item['question'][:70]}...")
|
||||
print(f" Expected: {item['expected_key']}")
|
||||
print(f" Got: {item['model_response'][:70]}")
|
||||
if len(report['incorrect_answers']) > 10:
|
||||
print(f" ... and {len(report['incorrect_answers']) - 10} more")
|
||||
|
||||
print(f"\n Report saved: {output_path}")
|
||||
print(f"{'='*60}\n")
|
||||
|
||||
|
||||
def run_benchmark(config_path: str, resume: bool = False, fresh: bool = False):
|
||||
config = load_config(config_path)
|
||||
model_cfg = config['model']
|
||||
bench_cfg = config['benchmark']
|
||||
|
||||
output_dir = Path(config_path).parent / bench_cfg.get('output_dir', 'runs')
|
||||
output_dir.mkdir(exist_ok=True)
|
||||
safe_name = sanitize_model_name(model_cfg['name'])
|
||||
|
||||
# Check for existing checkpoint
|
||||
checkpoint_path, checkpoint_data = find_latest_checkpoint(output_dir, safe_name)
|
||||
has_checkpoint = checkpoint_data is not None
|
||||
|
||||
if has_checkpoint and not fresh:
|
||||
if resume or input(f" Found checkpoint: {checkpoint_path.name}\n {len(checkpoint_data['completed_ids'])} questions completed. Resume? [y/N]: ").strip().lower() == 'y':
|
||||
print(f" Resuming from checkpoint...")
|
||||
results = checkpoint_data['results']
|
||||
completed_ids = set(checkpoint_data['completed_ids'])
|
||||
start_time = checkpoint_data['start_time']
|
||||
resume_from = len(completed_ids)
|
||||
else:
|
||||
completed_ids = set()
|
||||
results = {'correct': [], 'incorrect': [], 'timeouts': [], 'errors': []}
|
||||
start_time = time.time()
|
||||
resume_from = 0
|
||||
else:
|
||||
completed_ids = set()
|
||||
results = {'correct': [], 'incorrect': [], 'timeouts': [], 'errors': []}
|
||||
start_time = time.time()
|
||||
resume_from = 0
|
||||
|
||||
# Load dataset
|
||||
dataset_path = Path(config_path).parent / bench_cfg['dataset']
|
||||
if not dataset_path.exists():
|
||||
print(f"Dataset not found: {dataset_path}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
df = pd.read_parquet(dataset_path)
|
||||
questions = df.to_dict('records')
|
||||
|
||||
total = len(questions)
|
||||
max_q = bench_cfg.get('max_questions', 0) or total
|
||||
if max_q < total:
|
||||
import random
|
||||
random.seed(bench_cfg.get('seed', 42))
|
||||
random.shuffle(questions)
|
||||
questions = questions[:max_q]
|
||||
|
||||
# Filter out already completed questions when resuming
|
||||
if resume_from > 0:
|
||||
remaining = [q for q in questions if q.get('id', '') not in completed_ids]
|
||||
skipped = len(questions) - len(remaining)
|
||||
questions = remaining
|
||||
print(f" Skipped {skipped} completed questions. {len(questions)} remaining.\n")
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print(f" MCQ Benchmark")
|
||||
print(f"{'='*60}")
|
||||
print(f" Model: {model_cfg['name']}")
|
||||
print(f" Endpoint: {model_cfg['endpoint']}")
|
||||
print(f" Timeout: {model_cfg['timeout']}s")
|
||||
print(f" Max tokens: {model_cfg['max_tokens']}")
|
||||
print(f" Temperature: {model_cfg.get('temperature', 0.0)}")
|
||||
print(f" Questions: {len(questions) + resume_from} / {total}")
|
||||
if resume_from > 0:
|
||||
print(f" Resumed: {resume_from} completed")
|
||||
print(f"{'='*60}\n")
|
||||
|
||||
# Determine checkpoint file path
|
||||
if not checkpoint_path:
|
||||
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
|
||||
checkpoint_path = output_dir / f"{safe_name}_{timestamp}.checkpoint.json"
|
||||
|
||||
global_idx = resume_from
|
||||
|
||||
for i, q in enumerate(questions, 1):
|
||||
qid = q.get('id', f'unknown_{i}')
|
||||
question_text = q.get('text', '?')
|
||||
golden_key = q.get('golden_key', '?')
|
||||
options = q.get('options', {})
|
||||
|
||||
global_idx += 1
|
||||
display_q = question_text[:80] + ('...' if len(question_text) > 80 else '')
|
||||
|
||||
prompt = build_prompt(q)
|
||||
|
||||
response = None
|
||||
timed_out = False
|
||||
|
||||
for attempt in range(1 + bench_cfg.get('max_retries', 0)):
|
||||
response = query_model(prompt, config)
|
||||
if response is None:
|
||||
timed_out = True
|
||||
continue
|
||||
break
|
||||
|
||||
is_correct = False
|
||||
model_choice = ''
|
||||
|
||||
if timed_out:
|
||||
results['timeouts'].append({
|
||||
'id': qid,
|
||||
'question': question_text,
|
||||
'expected_key': golden_key,
|
||||
'options': options,
|
||||
})
|
||||
elif answers_match(golden_key, response or '', options):
|
||||
is_correct = True
|
||||
results['correct'].append({
|
||||
'id': qid,
|
||||
'question': question_text,
|
||||
'expected_key': golden_key,
|
||||
'model_response': (response or '')[:200],
|
||||
})
|
||||
else:
|
||||
model_choice = extract_choice(response or '')
|
||||
results['incorrect'].append({
|
||||
'id': qid,
|
||||
'question': question_text,
|
||||
'expected_key': golden_key,
|
||||
'model_response': (response or '')[:500],
|
||||
'model_choice': model_choice,
|
||||
'options': options,
|
||||
})
|
||||
|
||||
# Running accuracy
|
||||
answered = len(results['correct']) + len(results['incorrect']) + len(results['timeouts'])
|
||||
acc = (len(results['correct']) / answered * 100) if answered > 0 else 0.0
|
||||
|
||||
if timed_out:
|
||||
label = 'TIMEOUT'
|
||||
elif is_correct:
|
||||
label = 'CORRECT'
|
||||
else:
|
||||
label = f"WRONG (expected: {golden_key}, got: {model_choice or '?'})"
|
||||
|
||||
status = f" [{global_idx}/{len(questions) + resume_from}] {display_q} {label} [{acc:.1f}%]"
|
||||
print(f"\r{status}")
|
||||
|
||||
# Save checkpoint after each question
|
||||
completed_ids.add(qid)
|
||||
save_checkpoint(checkpoint_path, results, list(completed_ids), start_time)
|
||||
|
||||
if bench_cfg.get('delay_between_requests', 0) > 0:
|
||||
time.sleep(bench_cfg['delay_between_requests'])
|
||||
|
||||
# Build and save final report
|
||||
report = build_report(config, results, start_time)
|
||||
|
||||
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
|
||||
filename = f"{safe_name}_{timestamp}.json"
|
||||
output_path = output_dir / filename
|
||||
|
||||
with open(output_path, 'w') as f:
|
||||
json.dump(report, f, indent=2, ensure_ascii=False)
|
||||
|
||||
# Clean up checkpoint
|
||||
if checkpoint_path.exists():
|
||||
checkpoint_path.unlink()
|
||||
|
||||
print_summary(report, output_path)
|
||||
|
||||
return report
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='MCQ LLM Benchmark')
|
||||
parser.add_argument('--config', '-c', default='config.yaml',
|
||||
help='Path to config.yaml (default: config.yaml)')
|
||||
parser.add_argument('--resume', '-r', action='store_true',
|
||||
help='Resume from the latest checkpoint')
|
||||
parser.add_argument('--fresh', '-f', action='store_true',
|
||||
help='Start a new run, ignoring existing checkpoint')
|
||||
args = parser.parse_args()
|
||||
|
||||
config_path = Path(args.config)
|
||||
if not config_path.exists():
|
||||
print(f"Config not found: {config_path}")
|
||||
print(f" Copy config.example.yaml to config.yaml and fill in your values.", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
run_benchmark(str(config_path), resume=args.resume, fresh=args.fresh)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in New Issue
Block a user