mirror of https://github.com/aliasrobotics/cai.git
327 lines
14 KiB
Python
327 lines
14 KiB
Python
"""
|
|
Benchmark Evaluation Script
|
|
|
|
This script provides utilities to evaluate language models on cybersecurity-related multiple-choice and other question-answering benchmarks.
|
|
|
|
Usage:
|
|
python benchmarks/eval.py --model MODEL_NAME --dataset_file INPUT_FILE --eval EVAL_TYPE --output_dir OUTPUT_DIR --backend BACKEND
|
|
|
|
Arguments:
|
|
-m, --model Specify the model to evaluate (e.g., "gpt-4", "qwen2.5:14b", etc.)
|
|
-d, --dataset_file Path to the dataset file (JSON or TSV) containing questions to evaluate
|
|
-o, --output_dir Specify the output directory for results (default: "benchmarks/outputs/[benchmark_name]")
|
|
-B, --backend Backend to use: "openai", "openrouter", "ollama" (required)
|
|
-e, --eval Specify the evaluation benchmark
|
|
|
|
Example:
|
|
|
|
python benchmarks/eval.py --model ollama/qwen2.5:14b --dataset_file benchmarks/cybermetric/CyberMetric-2-v1.json --eval cybermetric --backend ollama
|
|
python benchmarks/eval.py --model ollama/qwen2.5:14b --dataset_file benchmarks/seceval/eval/datasets/questions-2.json --eval seceval --backend ollama
|
|
python benchmarks/eval.py --model ollama/qwen2.5:14b --dataset_file benchmarks/cti_bench/data/cti-mcq1.tsv --eval cti_bench --backend ollama
|
|
python benchmarks/eval.py --model ollama/qwen2.5:14b --dataset_file benchmarks/cti_bench/data/cti-ate2.tsv --eval cti_bench --backend ollama
|
|
|
|
python benchmarks/eval.py --model qwen/qwen3-32b:free --dataset_file benchmarks/cybermetric/CyberMetric-2-v1.json --eval cybermetric --backend openrouter
|
|
|
|
Environment Variables:
|
|
OPENROUTER_API_KEY: API key for OpenRouter (if using OpenRouter models)
|
|
OPENROUTER_API_BASE: Base URL for OpenRouter API (default: https://openrouter.ai/api/v1)
|
|
OLLAMA_API_BASE: Base URL for Ollama API via LiteLLM proxy (default: http://localhost:8000/v1)
|
|
OPENAI_API_KEY: API key for OpenAI (if using OpenAI models)
|
|
OPENAI_API_BASE: Base URL for OpenAI API (default: https://api.openai.com/v1)
|
|
"""
|
|
|
|
import json
|
|
import re
|
|
import time
|
|
import os
|
|
import datetime
|
|
import random
|
|
import string
|
|
import argparse
|
|
from tqdm import tqdm
|
|
import litellm
|
|
import requests
|
|
import csv
|
|
import os
|
|
import datetime
|
|
|
|
|
|
OPENROUTER_API_BASE = os.environ.get("OPENROUTER_API_BASE", "https://openrouter.ai/api/v1")
|
|
OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY")
|
|
OLLAMA_API_BASE = os.environ.get("OLLAMA_API_BASE", "http://localhost:8000/v1")
|
|
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
|
|
OPENAI_API_BASE = os.environ.get("OPENAI_API_BASE", "https://api.openai.com/v1")
|
|
|
|
|
|
def ask_model(question_obj, instruction, model, api_base, api_key=None, custom_llm_provider=None):
|
|
"""
|
|
Calls the model with the question and choices, returns the extracted answer.
|
|
"""
|
|
# Compose the prompt
|
|
prompt = instruction + "\n" + question_obj["Question"]
|
|
if question_obj.get("Choices"):
|
|
prompt += "\nChoices:\n" + question_obj["Choices"]
|
|
try:
|
|
response = litellm.completion(
|
|
model=model,
|
|
messages=[
|
|
{"role": "system", "content": "You are a security expert who answers questions."},
|
|
{"role": "user", "content": prompt},
|
|
],
|
|
api_base=api_base,
|
|
api_key=api_key,
|
|
custom_llm_provider=custom_llm_provider,
|
|
)
|
|
if hasattr(response, "choices") and response.choices:
|
|
content = response.choices[0].message.content
|
|
|
|
return content
|
|
except Exception as e:
|
|
print(f"Error: {e}.")
|
|
return None
|
|
|
|
def load_dataset(dataset_file, eval_type):
|
|
questions = [] #list of questions: {question: str, answers: dict, solution: str}
|
|
if eval_type == "seceval":
|
|
with open(dataset_file, 'r') as f:
|
|
data = json.load(f)
|
|
for question in data:
|
|
questions.append({
|
|
"Question": question["question"],
|
|
"Choices": "\n".join(question["choices"]),
|
|
"Solution": question["answer"]
|
|
})
|
|
|
|
elif eval_type == "cybermetric":
|
|
with open(dataset_file, 'r') as f:
|
|
data = json.load(f)
|
|
for question in data.get("questions", []):
|
|
questions.append({
|
|
"Question": question.get("question", ""),
|
|
"Choices": "\n".join([f"{k}: {v}" for k, v in question.get("answers", {}).items()]),
|
|
"Solution": question.get("solution", "")
|
|
})
|
|
elif eval_type == "cti_bench":
|
|
with open(dataset_file, 'r') as f:
|
|
reader = csv.reader(f, delimiter='\t')
|
|
header = next(reader, None)
|
|
for row in reader:
|
|
# Handle three possible formats:
|
|
# Format 1: [URL, Question, Option A, Option B, Option C, Option D, Prompt, GT] (8 columns)
|
|
# Format 2: [URL, Platform, Description, Prompt, GT] (5 columns)
|
|
# Format 3: [URL, Text, Prompt] (3 columns) -- see cti-taa.tsv
|
|
if len(row) == 8:
|
|
# MCQ format
|
|
questions.append({
|
|
"Question": row[1],
|
|
"Choices": f"A: {row[2]}\nB: {row[3]}\nC: {row[4]}\nD: {row[5]}",
|
|
"Solution": row[7]
|
|
})
|
|
elif len(row) == 5:
|
|
# ATE format (no choices, just open-ended)
|
|
questions.append({
|
|
"Question": row[2], # Description
|
|
"Choices": "", # No choices for ATE
|
|
"Solution": row[4] # GT
|
|
})
|
|
elif len(row) == 3:
|
|
# TAA format (threat actor attribution, open-ended)
|
|
questions.append({
|
|
"Question": row[1], # Text
|
|
"Choices": "", # No choices
|
|
"Solution": "" # No GT in this format
|
|
})
|
|
return questions
|
|
|
|
def run_evaluation(dataset, instruction, model, api_base=None, api_key=None, custom_llm_provider=None):
|
|
results = []
|
|
for idx, q in enumerate(dataset):
|
|
model_answer = ask_model(q, instruction, model, api_base, api_key, custom_llm_provider)
|
|
print(f"---------------{idx+1}/{len(dataset)}----------------")
|
|
print(f"Evaluating question: {q['Question']}")
|
|
print(f"Choices: {q['Choices']}")
|
|
print(f"Solution: {q['Solution']}")
|
|
print(f"Model Answer: {model_answer}")
|
|
print("--------------------------------")
|
|
results.append({
|
|
"Question": q["Question"],
|
|
"Choices": q["Choices"],
|
|
"ModelAnswer": model_answer,
|
|
"Solution": q["Solution"]
|
|
})
|
|
|
|
return results
|
|
|
|
def compute_accuracy(results, benchmark_name):
|
|
"""
|
|
Compute accuracy for a benchmark result set.
|
|
|
|
Args:
|
|
results (list of dict): Each dict should have the ground truth answer and model answer.
|
|
benchmark_name (str): The name of the benchmark.
|
|
Returns:
|
|
accuracy (float): Accuracy as a percentage (0-100).
|
|
correct_count (int): Number of correct answers.
|
|
total_count (int): Total number of evaluated items.
|
|
"""
|
|
correct_count = 0
|
|
total_count = 0
|
|
for item in results:
|
|
sol = item.get("Solution")
|
|
pred = item.get("ModelAnswer")
|
|
# For cybermetric, parse both gt and pred using parse_result_cybermetric
|
|
if benchmark_name.lower() == "cybermetric":
|
|
from benchmarks.eval import parse_result_cybermetric
|
|
pred_parsed = parse_result_cybermetric(pred)
|
|
if sol is not None and pred_parsed is not None:
|
|
if sol == pred_parsed:
|
|
correct_count += 1
|
|
total_count += 1
|
|
elif benchmark_name.lower() == "seceval":
|
|
|
|
pred_parsed = parse_result_seceval(pred)
|
|
if sol is not None and pred_parsed is not None:
|
|
if sol == pred_parsed:
|
|
correct_count += 1
|
|
total_count += 1
|
|
else:
|
|
if sol is not None and pred is not None:
|
|
# Accept either exact match or case-insensitive match
|
|
if str(sol).strip().upper() == str(pred).strip().upper():
|
|
correct_count += 1
|
|
total_count += 1
|
|
accuracy = (correct_count / total_count * 100) if total_count > 0 else 0.0
|
|
return accuracy, correct_count, total_count
|
|
def parse_result_seceval(result):
|
|
# Expecting format: 'ANSWER: X', 'ANSWER: XY', or 'ANSWER: XYZ' (1, 2, or 3 letters A-D)
|
|
if result is None:
|
|
return None
|
|
match = re.search(r"ANSWER:?\s*([A-D]{1,3})", result, re.IGNORECASE)
|
|
if match:
|
|
return match.group(1).upper()
|
|
return None
|
|
|
|
def parse_result_cybermetric(result):
|
|
# Expecting format: 'ANSWER: X'
|
|
if result is None:
|
|
return None
|
|
match = re.search(r"ANSWER:?\s*([A-D])", result, re.IGNORECASE)
|
|
if match:
|
|
return match.group(1).upper()
|
|
return None
|
|
|
|
def save_benchmark_results(
|
|
benchmark_name,
|
|
model,
|
|
dataset_file,
|
|
start_time,
|
|
end_time,
|
|
questions_processed,
|
|
correct_count,
|
|
accuracy,
|
|
total_count,
|
|
result
|
|
):
|
|
"""
|
|
Save benchmark results in CyberMetric-style format to output_dir/information.txt.
|
|
"""
|
|
output_dir = os.path.join(os.getcwd(), "benchmarks", "outputs", benchmark_name)
|
|
if not os.path.exists(output_dir):
|
|
os.makedirs(output_dir)
|
|
|
|
# Save information file as: <model>_<YYYYMMDD_HHMMSS>.txt
|
|
now_str = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
safe_model = "".join([c if c.isalnum() or c in ('-', '_') else '_' for c in str(model)])
|
|
info_file = os.path.join(output_dir, f"{safe_model}_{now_str}.txt")
|
|
duration = end_time - start_time
|
|
|
|
# Create a subdirectory for this run, named after info_file (without extension)
|
|
run_dir = os.path.splitext(info_file)[0]
|
|
if not os.path.exists(run_dir):
|
|
os.makedirs(run_dir)
|
|
|
|
# Save the info file as .tct inside the run_dir
|
|
info_file="information.txt"
|
|
info_file_tct = os.path.join(run_dir, os.path.basename(os.path.splitext(info_file)[0] + ".txt"))
|
|
with open(info_file_tct, "w") as f:
|
|
f.write(f"{benchmark_name} Evaluation\n")
|
|
f.write("=====================\n\n")
|
|
f.write(f"Model: {model}\n")
|
|
f.write(f"Dataset: {os.path.basename(dataset_file)}\n")
|
|
f.write(f"Start Time: {start_time.strftime('%Y-%m-%d %H:%M:%S')}\n")
|
|
f.write(f"Questions Processed: {questions_processed}\n")
|
|
f.write(f"Correct Answers: {correct_count}\n")
|
|
f.write(f"Accuracy: {accuracy:.2f}%\n")
|
|
f.write(f"End Time: {end_time.strftime('%Y-%m-%d %H:%M:%S')}\n")
|
|
f.write(f"Duration: {duration}\n")
|
|
|
|
# Save the results as answers.json inside the run_dir
|
|
results_file = os.path.join(run_dir, "answers.json")
|
|
with open(results_file, "w", encoding="utf-8") as rf:
|
|
json.dump(result, rf, ensure_ascii=False, indent=2)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="SecEval Evaluation CLI")
|
|
parser.add_argument("-d", "--dataset_file", type=str, required=True, help="Specify the dataset file to evaluate on.")
|
|
parser.add_argument("-B", "--backend", type=str, choices=["openai", "ollama", "openrouter"], required=True, help="Specify the llm type. openai: openai model, ollama: ollama model, openrouter: openrouter model")
|
|
parser.add_argument("-m", "--model", type=str, required=True, help="Specify the models.")
|
|
parser.add_argument("-e", "--eval", type=str, required=True, help="Specify the evaluation benchmark.")
|
|
args = parser.parse_args()
|
|
|
|
model = args.model
|
|
|
|
print(f"Evaluating model: {model}")
|
|
|
|
if args.backend == "openai":
|
|
api_base=OPENAI_API_BASE
|
|
api_key=OPENAI_API_KEY
|
|
custom_llm_provider=None
|
|
elif args.backend == "ollama":
|
|
api_base=OLLAMA_API_BASE
|
|
api_base = api_base.rstrip('/v1')
|
|
api_key=None
|
|
custom_llm_provider="ollama"
|
|
elif args.backend == "openrouter":
|
|
api_base=OPENROUTER_API_BASE
|
|
api_key=OPENROUTER_API_KEY
|
|
custom_llm_provider="openrouter"
|
|
else:
|
|
raise RuntimeError("Unknown backend")
|
|
|
|
if args.eval == "seceval":
|
|
instruction = """Below are multiple-choice questions concerning cybersecurity.
|
|
Please select the correct answers and respond with the correct letters A, B, C, or D.
|
|
You could select more than one letter
|
|
Always return in this format: 'ANSWER: XYZ' or 'ANSWER: X'.
|
|
"""
|
|
if args.dataset_file == "":
|
|
args.dataset_file = "benchmarks/seceval/eval/datasets/questions-2.json"
|
|
dataset = load_dataset(args.dataset_file, "seceval")
|
|
elif args.eval == "cybermetric":
|
|
instruction = """Choose the correct answer (A, B, C, or D) only.
|
|
Always return in this format: 'ANSWER: X'.
|
|
"""
|
|
if args.dataset_file == "":
|
|
args.dataset_file = "benchmarks/cybermetric/CyberMetric-2-v1.json"
|
|
dataset = load_dataset(args.dataset_file, "cybermetric")
|
|
elif args.eval == "cti_bench":
|
|
instruction = ""
|
|
if args.dataset_file == "":
|
|
args.dataset_file = "benchmarks/cti_bench/data/cti-mcq1.tsv"
|
|
dataset = load_dataset(args.dataset_file, "cti_bench")
|
|
|
|
start_time = datetime.datetime.now()
|
|
result = run_evaluation(dataset, instruction, model, api_base, api_key, custom_llm_provider)
|
|
end_time = datetime.datetime.now()
|
|
|
|
print(result)
|
|
|
|
accuracy, correct_count, total_count = compute_accuracy(result, args.eval )
|
|
print(f"Accuracy: {accuracy:.2f}% ({correct_count}/{total_count})")
|
|
|
|
save_benchmark_results(args.eval, model, args.dataset_file, start_time, end_time, len(dataset), correct_count, accuracy, total_count, result)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|