Pythonrun_benchmark.py
Copy
1
import os
2
import json
3
from dataclasses import dataclass, field
4
 
5
import openai
6
 
7
client = openai.OpenAI(api_key=os.environ["OPENAI_API_KEY"])
8
 
9
# Custom domain benchmark (replace with your domain questions)
10
BENCHMARK = [
11
    {
12
        "id": "q1",
13
        "category": "python",
14
        "question": "What does the walrus operator (:=) do in Python 3.8+?",
15
        "choices": {
16
            "A": "Assigns a value within an expression",
17
            "B": "Compares two values",
18
            "C": "Creates a new class",
19
            "D": "Defines a lambda function",
20
        },
21
        "answer": "A",
22
    },
23
    {
24
        "id": "q2",
25
        "category": "ml",
26
        "question": "Which metric is most appropriate for imbalanced classification?",
27
        "choices": {
28
            "A": "Accuracy",
29
            "B": "F1 Score",
30
            "C": "Mean Squared Error",
31
            "D": "R-squared",
32
        },
33
        "answer": "B",
34
    },
35
    # Add more questions...
36
]
37
 
38
@dataclass
39
class BenchmarkRun:
40
    model: str
41
    correct: int = 0
42
    total: int = 0
43
    by_category: dict = field(default_factory=dict)
44
 
45
    @property
46
    def accuracy(self) -> float:
47
        return self.correct / self.total if self.total > 0 else 0.0
48
 
49
def run_benchmark(model: str) -> BenchmarkRun:
50
    run = BenchmarkRun(model=model)
51
    
52
    for q in BENCHMARK:
53
        choices_str = "
54
".join(f"{k}) {v}" for k, v in q["choices"].items())
55
        prompt = f"{q['question']}
56
 
57
{choices_str}
58
 
59
Answer with just the letter (A/B/C/D):"
60
        
61
        response = client.chat.completions.create(
62
            model=model,
63
            messages=[{"role": "user", "content": prompt}],
64
            max_tokens=5,
65
            temperature=0,  # Deterministic for benchmarks
66
        )
67
        predicted = response.choices[0].message.content.strip().upper()[:1]
68
        correct = predicted == q["answer"]
69
        
70
        run.total += 1
71
        if correct:
72
            run.correct += 1
73
        
74
        cat = q["category"]
75
        if cat not in run.by_category:
76
            run.by_category[cat] = {"correct": 0, "total": 0}
77
        run.by_category[cat]["total"] += 1
78
        if correct:
79
            run.by_category[cat]["correct"] += 1
80
    
81
    return run
82
 
83
# Run against two models
84
for model in ["gpt-4o-mini", "gpt-4o"]:
85
    result = run_benchmark(model)
86
    print(f"\n{model}: accuracy={result.accuracy:.1%} ({result.correct}/{result.total})")
87
    for cat, stats in result.by_category.items():
88
        acc = stats["correct"] / stats["total"]
89
        print(f"  {cat}: {acc:.1%}")
90