test_vector_db.py•3.63 kB
#!/usr/bin/env python3
"""
Test Vector Database with Real 14K MMLU Questions
"""
import json
from pathlib import Path
from benchmark_vector_db import BenchmarkVectorDB, BenchmarkQuestion
def load_real_mmlu_data():
"""Load the 14K real MMLU questions"""
print("Loading 14,042 real MMLU questions...")
with open("./data/benchmark_results/mmlu_real_results.json") as f:
data = json.load(f)
questions = []
for qid, q in data['questions'].items():
questions.append(BenchmarkQuestion(
question_id=q['question_id'],
source_benchmark=q['source_benchmark'],
domain=q['domain'],
question_text=q['question_text'],
correct_answer="", # Not needed for vector DB
choices=q.get('choices'),
success_rate=q['success_rate'],
difficulty_score=1.0 - q['success_rate'],
difficulty_label=q['difficulty_label'],
num_models_tested=q['num_models_tested']
))
print(f"✓ Loaded {len(questions)} questions")
return questions
def build_and_test_vector_db():
"""Build vector DB with real data and run tests"""
# Initialize fresh database
db = BenchmarkVectorDB(
db_path=Path("./data/benchmark_vector_db"),
embedding_model="all-MiniLM-L6-v2"
)
# Load real data
questions = load_real_mmlu_data()
# Index questions (this takes 1-2 minutes)
print("\nIndexing into vector database...")
db.index_questions(questions)
# Get stats
stats = db.get_statistics()
print(f"\n📊 Database Statistics:")
print(f" Total Questions: {stats['total_questions']}")
print(f" Difficulty Distribution: {stats.get('difficulty_levels', {})}")
# Test with diverse prompts
test_prompts = [
# Should be HARD (physics/abstract math)
"Calculate the quantum correction to the partition function for a 3D harmonic oscillator",
"Prove that there are infinitely many prime numbers",
"Find all zeros of the polynomial x^3 + 2x + 2 in the finite field Z_7",
# Should be MODERATE (reasoning)
"Diagnose a patient with acute chest pain and shortness of breath",
"Explain the legal doctrine of precedent in common law systems",
"Implement a binary search tree with insert and search operations",
# Should be EASY (basic knowledge)
"What is 2 + 2?",
"What is the capital of France?",
"Who wrote Romeo and Juliet?",
"What is the boiling point of water in Celsius?",
# Cross-domain abstract reasoning
"Statement 1 | Every field is also a ring. Statement 2 | Every ring has a multiplicative identity."
]
print(f"\n🧪 Testing {len(test_prompts)} diverse prompts:")
print("="*80)
for prompt in test_prompts:
result = db.query_similar_questions(prompt, k=5)
print(f"\n📝 '{prompt}'")
print(f" 🎯 Risk: {result['risk_level']}")
print(f" 📊 Success Rate: {result['weighted_success_rate']:.1%}")
print(f" 📌 Top Match: {result['similar_questions'][0]['question_text'][:80]}...")
if result['similar_questions'][0]['success_rate'] < 0.5:
print(f" 🔍 Found similar hard question (success: {result['similar_questions'][0]['success_rate']:.0%})")
print(f" 💡 {result['recommendation']}")
print("\n" + "="*80)
print("✅ Real data test complete!")
print("="*80)
if __name__ == "__main__":
build_and_test_vector_db()