cai/benchmarks/cybermetric/CyberMetric_evaluator.py

398 lines
17 KiB
Python

"""
CyberMetric Evaluator for LLMs
This script evaluates the performance of language models on the CyberMetric benchmark.
It supports OpenRouter-hosted models, local Ollama models via LiteLLM proxy, and OpenAI models.
Usage:
python CyberMetric_evaluator.py --model_name MODEL_NAME [--file_path FILE_PATH] [--api_key API_KEY]
Arguments:
--model_name: Required. Model name with prefix (openrouter/, ollama/, or openai/)
Examples: openrouter/anthropic/claude-3-opus, ollama/llama3, openai/gpt-4o
--file_path: Optional. Path to the CyberMetric JSON file (default: CyberMetric-2-v1.json)
--api_key: Optional. API key for OpenRouter or OpenAI (can also use env vars)
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)
Examples:
# Run with an OpenRouter model
python CyberMetric_evaluator.py --model_name openrouter/qwen/qwen3-32b:free
# Run with a local Ollama model (requires LiteLLM proxy running)
python CyberMetric_evaluator.py --model_name ollama/qwen2.5:14b
# Run with an OpenAI model
python CyberMetric_evaluator.py --model_name openai/gpt-4o
# Specify a different benchmark file
python CyberMetric_evaluator.py --model_name openai/gpt-4o --file_path CyberMetric-10000-v1.json
"""
import json
import re
import time
import os
import datetime
import random
import string
import argparse
from tqdm import tqdm
import litellm
import requests
# Enable debug mode for litellm
#litellm._turn_on_debug()
# Default API bases
OPENROUTER_API_BASE = "https://openrouter.ai/api/v1"
OLLAMA_LITELLM_API_BASE = os.environ.get("OLLAMA_API_BASE", "http://localhost:8000/v1")
OPENAI_API_BASE = os.environ.get("OPENAI_API_BASE", "https://api.openai.com/v1")
class CyberMetricEvaluator:
def __init__(self, model_name, file_path, api_key=None, openrouter_api_base=None, ollama_litellm_api_base=None, openai_api_base=None):
self.model_name = model_name
self.file_path = file_path
self.failed_questions = []
self.failed_count = 0
# Set API configurations
self.openrouter_api_base = openrouter_api_base or os.environ.get("OPENROUTER_API_BASE", OPENROUTER_API_BASE)
self.ollama_litellm_api_base = ollama_litellm_api_base or os.environ.get("OLLAMA_LITELLM_API_BASE", OLLAMA_LITELLM_API_BASE)
self.openai_api_base = openai_api_base or os.environ.get("OPENAI_API_BASE", OPENAI_API_BASE)
# Set API key based on model type
self.api_key = None
if self.model_name.startswith("openrouter/"):
self.api_key = api_key or os.environ.get("OPENROUTER_API_KEY")
elif self.model_name.startswith("openai/"):
self.api_key = api_key or os.environ.get("OPENAI_API_KEY")
self.start_time = datetime.datetime.now()
# Create output directory structure
self.output_dir = self.create_output_directory()
self.info_file = os.path.join(self.output_dir, "information.txt")
self.report_file = os.path.join(self.output_dir, "report_failed_questions.json")
# Initialize info file
self.initialize_info_file()
print("--DEBUG: model_name: ", self.model_name)
def create_output_directory(self):
# Create base directory if it doesn't exist
base_dir = "outputs"
if not os.path.exists(base_dir):
os.makedirs(base_dir)
# Format model name for directory (replace / with -)
model_dir_name = self.model_name.replace("/", "-")
# Get current date
current_date = datetime.datetime.now().strftime("%Y-%m-%d")
# Create directory name
dir_name = f"{model_dir_name}-{current_date}"
# If directory already exists, add random string
full_path = os.path.join(base_dir, dir_name)
if os.path.exists(full_path):
random_str = ''.join(random.choices(string.ascii_lowercase + string.digits, k=5))
dir_name = f"{model_dir_name}-{current_date}-{random_str}"
full_path = os.path.join(base_dir, dir_name)
os.makedirs(full_path)
return full_path
def initialize_info_file(self):
with open(self.info_file, 'w') as file:
file.write(f"CyberMetric Evaluation\n")
file.write(f"=====================\n\n")
file.write(f"Model: {self.model_name}\n")
file.write(f"Dataset: {self.file_path}\n")
file.write(f"Start Time: {self.start_time.strftime('%Y-%m-%d %H:%M:%S')}\n")
file.write(f"Status: Running\n")
file.write(f"Questions Processed: 0\n")
file.write(f"Correct Answers: 0\n")
file.write(f"Accuracy: 0.00%\n")
def update_info_file(self, questions_processed, correct_count, status="Running"):
accuracy = correct_count / questions_processed * 100 if questions_processed > 0 else 0
with open(self.info_file, 'w') as file:
file.write(f"CyberMetric Evaluation\n")
file.write(f"=====================\n\n")
file.write(f"Model: {self.model_name}\n")
file.write(f"Dataset: {self.file_path}\n")
file.write(f"Start Time: {self.start_time.strftime('%Y-%m-%d %H:%M:%S')}\n")
file.write(f"Status: {status}\n")
file.write(f"Questions Processed: {questions_processed}\n")
file.write(f"Correct Answers: {correct_count}\n")
file.write(f"Accuracy: {accuracy:.2f}%\n")
if status == "Completed":
end_time = datetime.datetime.now()
duration = end_time - self.start_time
file.write(f"End Time: {end_time.strftime('%Y-%m-%d %H:%M:%S')}\n")
file.write(f"Duration: {duration}\n")
def read_json_file(self):
with open(self.file_path, 'r') as file:
return json.load(file)
def save_failed_questions(self):
report_data = {"questions": self.failed_questions}
with open(self.report_file, 'w') as file:
json.dump(report_data, file, indent=4)
print(f"Updated failed questions report in {self.report_file}")
@staticmethod
def extract_answer(response):
if response.strip():
match = re.search(r"ANSWER:?\s*([A-D])", response, re.IGNORECASE)
if match:
return match.group(1).upper()
return None
def ask_ollama_litellm(self, prompt, max_retries=5):
for attempt in range(max_retries):
try:
response = litellm.completion(
model=self.model_name,
messages=[
{"role": "system", "content": "You are a security expert who answers questions."},
{"role": "user", "content": prompt},
],
api_base="http://localhost:8000"
)
if hasattr(response, "choices") and response.choices:
content = response.choices[0].message.content
result = self.extract_answer(content)
if result:
print("--DEBUG: result: ", result)
return result
else:
print("Incorrect answer format detected. Attempting the question again.")
except Exception as e:
print(f"Error: {e}. Attempting the question again in {2 ** attempt} seconds.")
time.sleep(2 ** attempt)
return None
def ask_openrouter(self, prompt, max_retries=5):
if not self.api_key:
raise ValueError("API key is required for OpenRouter models")
for attempt in range(max_retries):
try:
response = litellm.completion(
model=self.model_name,
messages=[
{"role": "system", "content": "You are a security expert who answers questions."},
{"role": "user", "content": prompt},
],
api_base=self.openrouter_api_base,
api_key=self.api_key,
headers={
"HTTP-Referer": "https://your-site-url.com",
"X-Title": "CyberMetric Evaluator"
}
)
if hasattr(response, "choices") and response.choices:
content = response.choices[0].message.content
result = self.extract_answer(content)
if result:
print("--DEBUG: result: ", result)
return result
else:
print("Incorrect answer format detected. Attempting the question again.")
except Exception as e:
print(f"Error: {e}. Attempting the question again in {2 ** attempt} seconds.")
time.sleep(2 ** attempt)
return None
def ask_openai(self, prompt, max_retries=5):
if not self.api_key:
raise ValueError("API key is required for OpenAI models")
for attempt in range(max_retries):
try:
# Extract the actual model name without the 'openai/' prefix
actual_model = self.model_name.replace("openai/", "")
response = litellm.completion(
model=actual_model, # Use the actual OpenAI model name
messages=[
{"role": "system", "content": "You are a security expert who answers questions."},
{"role": "user", "content": prompt},
],
api_base=self.openai_api_base,
api_key=self.api_key
)
if hasattr(response, "choices") and response.choices:
content = response.choices[0].message.content
result = self.extract_answer(content)
if result:
print("--DEBUG: result: ", result)
return result
else:
print("Incorrect answer format detected. Attempting the question again.")
except Exception as e:
print(f"Error: {e}. Attempting the question again in {2 ** attempt} seconds.")
time.sleep(2 ** attempt)
return None
def ask_llm(self, question, answers, max_retries=5):
options = ', '.join([f"{key}) {value}" for key, value in answers.items()])
prompt = (
f"Question: {question}\n"
f"Options: {options}\n\n"
"Choose the correct answer (A, B, C, or D) only. Always return in this format: 'ANSWER: X'"
)
print("--DEBUG: prompt: ", prompt)
if self.model_name.startswith("openrouter/"):
return self.ask_openrouter(prompt, max_retries)
elif self.model_name.startswith("ollama/"):
return self.ask_ollama_litellm(prompt, max_retries)
elif self.model_name.startswith("openai/"):
return self.ask_openai(prompt, max_retries)
else:
print(f"Error: Unsupported model prefix: {self.model_name}")
return None
def run_evaluation(self):
if not (self.model_name.startswith("openrouter/") or
self.model_name.startswith("ollama/") or
self.model_name.startswith("openai/")):
print("Error: Model name must start with 'ollama/', 'openrouter/', or 'openai/'")
return
json_data = self.read_json_file()
questions_data = json_data['questions']
correct_count = 0
incorrect_answers = []
with tqdm(total=len(questions_data), desc="Processing Questions") as progress_bar:
for i, item in enumerate(questions_data):
question = item['question']
answers = item['answers']
correct_answer = item['solution']
llm_answer = self.ask_llm(question, answers)
if llm_answer == correct_answer:
correct_count += 1
else:
self.failed_questions.append({
'question': question,
'answers': answers,
'solution': correct_answer,
'llm_answer': llm_answer
})
self.failed_count += 1
if self.failed_count % 2 == 0:
self.save_failed_questions()
incorrect_answers.append({
'question': question,
'correct_answer': correct_answer,
'llm_answer': llm_answer
})
# Update progress and information file
questions_processed = i + 1
accuracy_rate = correct_count / questions_processed * 100
progress_bar.set_postfix_str(f"Accuracy: {accuracy_rate:.2f}%")
progress_bar.update(1)
# Update info file every 5 questions
if questions_processed % 5 == 0 or questions_processed == len(questions_data):
self.update_info_file(questions_processed, correct_count)
# Final update with completed status
self.update_info_file(len(questions_data), correct_count, "Completed")
print(f"\nFinal Accuracy: {correct_count / len(questions_data) * 100:.2f}%")
if self.failed_questions:
self.save_failed_questions() # final failed questions
if incorrect_answers:
print("\nIncorrect Answers:")
for item in incorrect_answers:
print(f"Question: {item['question']}")
print(f"Expected Answer: {item['correct_answer']}, LLM Answer: {item['llm_answer']}\n")
if __name__ == "__main__":
#litellm._turn_on_debug()
# Create argument parser
parser = argparse.ArgumentParser(description='CyberMetric Evaluator for LLMs')
parser.add_argument('--model_name', type=str, required=True,
help='Model name with prefix (openrouter/, ollama/, or openai/)')
parser.add_argument('--file_path', type=str, default='CyberMetric-2-v1.json',
help='Path to the CyberMetric JSON file')
parser.add_argument('--api_key', type=str,
help='API key for OpenRouter or OpenAI (can also use env vars)')
args = parser.parse_args()
model_name = args.model_name
file_path = args.file_path
api_key = args.api_key
if model_name.startswith("ollama/"):
# Ollama configuration
evaluator = CyberMetricEvaluator(
model_name=model_name,
file_path=file_path
)
print(f"Using Ollama configuration with LiteLLM proxy on port 8000")
elif model_name.startswith("openrouter/"):
# OpenRouter configuration
api_key = api_key or os.environ.get("OPENROUTER_API_KEY")
if not api_key:
raise ValueError("API key must be provided via --api_key or OPENROUTER_API_KEY environment variable for OpenRouter models")
evaluator = CyberMetricEvaluator(
model_name=model_name,
file_path=file_path,
api_key=api_key,
openrouter_api_base=os.environ.get("OPENROUTER_API_BASE", "https://openrouter.ai/api/v1")
)
print("Using OpenRouter configuration")
elif model_name.startswith("openai/"):
# OpenAI configuration
api_key = api_key or os.environ.get("OPENAI_API_KEY")
if not api_key:
raise ValueError("API key must be provided via --api_key or OPENAI_API_KEY environment variable for OpenAI models")
print(f"API key provided: {api_key[:4]}...{api_key[-4:] if len(api_key) > 8 else ''}")
evaluator = CyberMetricEvaluator(
model_name=model_name,
file_path=file_path,
api_key=api_key,
openai_api_base=os.environ.get("OPENAI_API_BASE", "https://api.openai.com/v1")
)
print("Using OpenAI configuration")
else:
raise ValueError("Model name must start with 'ollama/', 'openrouter/', or 'openai/'")
# Run the evaluation
evaluator.run_evaluation()