mirror of https://github.com/aliasrobotics/cai.git
ass seceval whit ollama
This commit is contained in:
Mery-Sanz
parent
a426ee87dc
commit
03d4faba0c
|
|
@ -1,332 +0,0 @@
|
|||
import argparse
|
||||
from typing import Any, Dict, List
|
||||
from dotenv import load_dotenv
|
||||
import asyncio
|
||||
load_dotenv()
|
||||
from langchain_community.llms import HuggingFacePipeline, TextGen
|
||||
from langchain_community.chat_models import AzureChatOpenAI
|
||||
from langchain_community.adapters.openai import convert_message_to_dict
|
||||
from langchain_community.cache import SQLiteCache
|
||||
from langchain.schema.language_model import BaseLanguageModel
|
||||
from langchain.adapters.openai import convert_message_to_dict
|
||||
from langchain.schema.messages import (
|
||||
AIMessage,
|
||||
SystemMessage,
|
||||
HumanMessage,
|
||||
)
|
||||
import json
|
||||
import re
|
||||
from pathlib import Path
|
||||
import os
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
import time
|
||||
from langchain.globals import set_llm_cache
|
||||
from langchain.cache import SQLiteCache
|
||||
|
||||
import litellm
|
||||
|
||||
set_llm_cache(
|
||||
SQLiteCache(
|
||||
database_path=str(Path(__file__).parent.parent / ".langchain.db")
|
||||
)
|
||||
)
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s %(levelname)s %(message)s",
|
||||
)
|
||||
logger.addHandler(logging.FileHandler(f"./eval-{int(time.time())}.log", "w"))
|
||||
instruction = "Below are multiple-choice questions concerning cybersecurity. Please select the correct answers and respond with the letters ABCD only."
|
||||
chat_few_shot = [
|
||||
HumanMessage(
|
||||
content="Question: Which mitigation prevent stack overflow bug? A: Stack Canary. B: ALSR. C: CFI. D: Code Signing.",
|
||||
),
|
||||
AIMessage(
|
||||
content="Answer: ABC",
|
||||
),
|
||||
]
|
||||
few_shot = """
|
||||
Question: You are reviewing a piece of code to identify potential vulnerabilities related to NULL pointer dereferences. Which of the following strategies should you consider implementing to mitigate the risk? Choose all that apply. A: Use a programming language that inherently manages pointers and prevents direct memory access. B: Apply input validation to all variables and data stores that receive data from external sources to ensure they are initialized to expected values. C: Explicitly initialize all variables and other data stores, either during declaration or just before the first usage. D: Implement a custom memory management system that automatically assigns a unique identifier to each pointer.
|
||||
Answer: ABC
|
||||
Question: Analyze the code snippet provided. Which of the following practices should be implemented in a multithreaded application to prevent NULL pointer dereference vulnerabilities? Select all that apply.A: Ensure that proper locking APIs are used to lock the pointer before the if statement and unlock when it has finished to prevent race conditions. B: Check the return value of all functions that return a value, including network-related functions, and verify that it is non-null before using it.C: Use automated static analysis tools that target this type of weakness, understanding that while not perfect, they can still be effective.D: Verify that a non-nil 'response' is present before deferring 'response.Body.Close()' to handle cases where the 'Do' method returns an error.
|
||||
Answer: ABCD
|
||||
Question: In a bank management application, a method called `createBankAccount` allows the creation of bank accounts without user identity verification or authority checks. Based on common security practices, which of the following measures would improve security?A: Implement a centralized authentication mechanism for operations that require a proven user identity.B: Duplicate client-side authentication checks on the server side to prevent bypassing of security checks.C: Use a vetted library or framework that offers robust authentication capabilities to prevent security weaknesses.D: Apply strong authentication requirements for all users accessing data stored in the cloud.
|
||||
Answer: ABCD
|
||||
Question: What are the potential consequences if an application handling sensitive data does not require authentication for data access, and what are the appropriate mitigation strategies?A: Attackers could easily access sensitive data from any location; require strong authentication for data access using cloud provider controls.B: Security checks might be bypassed if only performed on the client side; implement checks on both client and server sides.C: Unauthenticated users could alter product functionality; do not use authentication for critical functionality in products.D: Sensitive data may be accessed without proper credentials; utilize authentication capabilities provided by the framework or operating system.
|
||||
Answer: ABD
|
||||
Question: To prevent security vulnerabilities related to deserialization of untrusted data in a Java application, which of the following practices should a developer implement?A: Use the signing/sealing features of the programming language to assure that deserialized data has not been tainted.B: Explicitly define a final readObject() method to throw an exception and prevent deserialization.C: Populate a new object by deserializing data to ensure data flows through safe input validation functions.D: Make fields transient to protect them from deserialization and prevent carrying over sensitive variables.
|
||||
Answer: ABCD
|
||||
"""
|
||||
|
||||
|
||||
def init_hf_llm(model_id: str):
|
||||
# check transformers and torch installation
|
||||
try:
|
||||
import transformers
|
||||
except ImportError:
|
||||
raise ImportError("Please install transformers with `pip install transformers`")
|
||||
try:
|
||||
import torch
|
||||
|
||||
flash_attn_enable = torch.cuda.get_device_capability()[0] >= 8
|
||||
except ImportError:
|
||||
raise ImportError("Please install torch with `pip install torch`")
|
||||
|
||||
# todo: add flash_attn_enable to the model_kwargs
|
||||
|
||||
llm = HuggingFacePipeline.from_model_id(
|
||||
model_id=model_id,
|
||||
task="text-generation",
|
||||
pipeline_kwargs={"max_new_tokens": 5},
|
||||
device=0,
|
||||
model_kwargs={"trust_remote_code": True, "torch_dtype": torch.bfloat16},
|
||||
)
|
||||
return llm
|
||||
|
||||
|
||||
def init_textgen_llm(model_id: str):
|
||||
if os.environ.get("TEXTGEN_MODEL_URL") is None:
|
||||
raise RuntimeError("Please set TEXTGEN_MODEL_URL")
|
||||
llm = TextGen(model_url=os.environ["TEXTGEN_MODEL_URL"]) # type: ignore
|
||||
return llm
|
||||
|
||||
|
||||
def init_azure_openai_llm(model_id: str):
|
||||
if os.environ.get("OPENAI_API_ENDPOINT") is None:
|
||||
raise RuntimeError("Please set OPENAI_API_ENDPOINT")
|
||||
if os.environ.get("OPENAI_API_KEY") is None:
|
||||
raise RuntimeError("Please set OPENAI_API_KEY")
|
||||
azure_params = {
|
||||
"model": model_id,
|
||||
"openai_api_base": os.environ["OPENAI_API_ENDPOINT"],
|
||||
"openai_api_key": os.environ["OPENAI_API_KEY"],
|
||||
"openai_api_type": os.environ.get("OPENAI_API_TYPE", "azure"),
|
||||
"openai_api_version": "2023-07-01-preview",
|
||||
}
|
||||
return AzureChatOpenAI(**azure_params) # type: ignore
|
||||
|
||||
|
||||
def init_ollama_llm(model_id: str):
|
||||
class OllamaChat:
|
||||
async def abatch(self, prompts: List[str]):
|
||||
responses = []
|
||||
for prompt in prompts:
|
||||
try:
|
||||
|
||||
completion = litellm.completion(
|
||||
model=model_id,
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
api_base="http://localhost:8000" # Ollama API endpoint
|
||||
)
|
||||
if hasattr(completion, "choices") and completion.choices:
|
||||
content = completion.choices[0].message.content
|
||||
result = self.extract_answer(content)
|
||||
if result:
|
||||
print("--DEBUG: result: ", result)
|
||||
responses.append(result)
|
||||
else:
|
||||
print("Incorrect answer format detected. Attempting the question again.")
|
||||
except Exception as e:
|
||||
logging.error(f"Ollama error: {e}")
|
||||
responses.append(f"Error: {e}")
|
||||
return responses
|
||||
|
||||
return OllamaChat()
|
||||
|
||||
def load_dataset(dataset_path: str):
|
||||
with open(dataset_path, "r") as f:
|
||||
dataset = json.load(f)
|
||||
return dataset
|
||||
|
||||
|
||||
async def batch_inference_dataset(
|
||||
llm: BaseLanguageModel, batch: List[Dict[str, Any]], chat=False
|
||||
):
|
||||
results = []
|
||||
llm_inputs = []
|
||||
for dataset_row in batch:
|
||||
question_text = (
|
||||
"Question: " + dataset_row["question"] + " ".join(dataset_row["choices"])
|
||||
)
|
||||
question_text = question_text.replace("\n", " ")
|
||||
if chat:
|
||||
llm_input = (
|
||||
[SystemMessage(content=instruction)]
|
||||
+ chat_few_shot
|
||||
+ [HumanMessage(content=question_text)]
|
||||
)
|
||||
else:
|
||||
llm_input = instruction + few_shot + question_text + "\n"
|
||||
|
||||
llm_inputs.append(llm_input)
|
||||
try:
|
||||
llm_outputs = await llm.abatch(llm_inputs)
|
||||
except Exception as e:
|
||||
logging.error(f"error in processing batch {e}")
|
||||
llm_outputs = [f"{e}" * len(llm_inputs)]
|
||||
for idx, llm_output in enumerate(llm_outputs):
|
||||
if type(llm_output) == AIMessage:
|
||||
llm_output: str = llm_output.content # type: ignore
|
||||
if "Answer:" in llm_output:
|
||||
llm_output = llm_output.replace("Answer:", "")
|
||||
if chat:
|
||||
batch[idx]["llm_input"] = convert_message_to_dict(llm_inputs[idx])
|
||||
else:
|
||||
batch[idx]["llm_input"] = llm_inputs[idx]
|
||||
batch[idx]["llm_output"] = llm_output
|
||||
batch[idx]["llm_answer"] = "".join(
|
||||
sorted(list(set(re.findall(r"[A-D]", llm_output))))
|
||||
)
|
||||
batch[idx]["score"] = int(
|
||||
batch[idx]["llm_answer"].lower() == batch[idx]["answer"].lower()
|
||||
)
|
||||
logging.info(
|
||||
f'llm_output: {llm_output}, parsed answer: {batch[idx]["llm_answer"]}, answer: {batch[idx]["answer"]}'
|
||||
)
|
||||
results.append(batch[idx])
|
||||
return results
|
||||
|
||||
|
||||
def inference_dataset(
|
||||
llm: BaseLanguageModel,
|
||||
dataset: List[Dict[str, Any]],
|
||||
batch_size: int = 1,
|
||||
chat: bool = False,
|
||||
):
|
||||
# Prepare the batched inference
|
||||
def chunks(lst, n):
|
||||
for i in range(0, len(lst), n):
|
||||
yield lst[i : i + n]
|
||||
|
||||
# Asynchronously process dataset in batches
|
||||
loop = asyncio.get_event_loop()
|
||||
batches = list(chunks(dataset, batch_size))
|
||||
results = []
|
||||
for idx, batch in enumerate(batches):
|
||||
logger.info(f"processing batch {idx+1}/{len(batches)}")
|
||||
results += loop.run_until_complete(batch_inference_dataset(llm, batch, chat))
|
||||
return results
|
||||
|
||||
|
||||
def count_score_by_topic(dataset: List[Dict[str, Any]]):
|
||||
score_by_topic = {}
|
||||
total_score_by_topic = {}
|
||||
score = 0
|
||||
for dataset_row in dataset:
|
||||
for topic in dataset_row["topics"]:
|
||||
if topic not in score_by_topic:
|
||||
score_by_topic[topic] = 0
|
||||
total_score_by_topic[topic] = 0
|
||||
score_by_topic[topic] += dataset_row["score"]
|
||||
total_score_by_topic[topic] += 1
|
||||
score += dataset_row["score"]
|
||||
score_fraction = {
|
||||
k: f"{v}/{total_score_by_topic[k]}" for k, v in score_by_topic.items()
|
||||
}
|
||||
score_float = {
|
||||
k: round(100 * float(v) / float(total_score_by_topic[k]), 4)
|
||||
for k, v in score_by_topic.items()
|
||||
}
|
||||
score_float["Overall"] = round(100 * float(score) / float(len(dataset)), 4)
|
||||
score_fraction["Overall"] = f"{score}/{len(dataset)}"
|
||||
return score_fraction, score_float
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="SecEval Evaluation CLI")
|
||||
|
||||
parser.add_argument(
|
||||
"-o",
|
||||
"--output_dir",
|
||||
type=str,
|
||||
default="/tmp",
|
||||
help="Specify the output directory.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-d",
|
||||
"--dataset_file",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Specify the dataset file to evaluate on.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-c",
|
||||
"--chat",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Evaluate on chat model.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-b",
|
||||
"--batch_size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Specify the batch size.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-B",
|
||||
"--backend",
|
||||
type=str,
|
||||
choices=["remote_hf", "azure", "textgen", "local_hf", "ollama"],
|
||||
required=True,
|
||||
help="Specify the llm type. remote_hf: remote huggingface model backed, azure: azure openai model, textgen: textgen backend, local_hf: local huggingface model backed",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-m",
|
||||
"--models",
|
||||
type=str,
|
||||
nargs="+",
|
||||
required=True,
|
||||
help="Specify the models.",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
models = list(args.models)
|
||||
|
||||
logging.info(f"evaluating models: {models}")
|
||||
for model_id in models:
|
||||
if args.backend == "remote_hf":
|
||||
llm = init_hf_llm(model_id)
|
||||
elif args.backend == "local_hf":
|
||||
model_dir = os.environ.get("LOCAL_HF_MODEL_DIR")
|
||||
if model_dir is None:
|
||||
raise RuntimeError(
|
||||
"Please set LOCAL_HF_MODEL_DIR when using local_hf backend"
|
||||
)
|
||||
model_id = os.path.join(model_dir, model_id)
|
||||
llm = init_hf_llm(model_id)
|
||||
elif args.backend == "textgen":
|
||||
llm = init_textgen_llm(model_id)
|
||||
elif args.backend == "azure":
|
||||
llm = init_azure_openai_llm(model_id)
|
||||
elif args.backend == "ollama":
|
||||
if not args.models:
|
||||
raise RuntimeError("Please specify --model for Ollama backend")
|
||||
llm = init_ollama_llm(args.models)
|
||||
else:
|
||||
raise RuntimeError("Unknown backend")
|
||||
|
||||
dataset = load_dataset(args.dataset_file)
|
||||
result = inference_dataset(
|
||||
llm, dataset, batch_size=args.batch_size, chat=args.chat
|
||||
)
|
||||
score_fraction, score_float = count_score_by_topic(result)
|
||||
result_with_score = {
|
||||
"score_fraction": score_fraction,
|
||||
"score_float": score_float,
|
||||
"detail": result,
|
||||
}
|
||||
output_path = (
|
||||
Path(args.output_dir)
|
||||
/ f"{Path(args.dataset_file).stem}_{os.path.basename(model_id)}.json"
|
||||
)
|
||||
logger.info(f"writing result to {output_path}")
|
||||
with open(output_path, "w") as f:
|
||||
json.dump(result_with_score, f, indent=4)
|
||||
del llm
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,157 @@
|
|||
"python3 benchmarks/seceval/eval/eval_ollama.py --dataset_file benchmarks/seceval/eval/datasets/questions.json --output_dir benchmarks/seceval/eval/outputs --model ollama/qwen2.5:14b"
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from langchain.schema.messages import SystemMessage, HumanMessage
|
||||
from langchain.adapters.openai import convert_message_to_dict
|
||||
import litellm
|
||||
|
||||
load_dotenv()
|
||||
|
||||
def init_ollama_llm(model_id: str):
|
||||
class OllamaChat:
|
||||
async def abatch(self, prompts: List[str]):
|
||||
responses = []
|
||||
for prompt in prompts:
|
||||
try:
|
||||
ollama_api_base = os.getenv("OLLAMA_API_BASE", "http://localhost:8000")
|
||||
api_base = ollama_api_base.rstrip('/v1')
|
||||
completion = litellm.completion(
|
||||
model=model_id,
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
api_base=api_base,
|
||||
custom_llm_provider="ollama"
|
||||
)
|
||||
if hasattr(completion, "choices") and completion.choices:
|
||||
content = completion.choices[0].message.content
|
||||
result = self.extract_answer(content)
|
||||
if result:
|
||||
|
||||
responses.append(result)
|
||||
else:
|
||||
print("Incorrect answer format detected.")
|
||||
responses.append("Error: No result parsed")
|
||||
except Exception as e:
|
||||
logging.error(f"Ollama error: {e}")
|
||||
responses.append(f"Error: {e}")
|
||||
return responses
|
||||
|
||||
def extract_answer(self, text: str):
|
||||
match = re.findall(r"[A-D]", text.upper())
|
||||
return "".join(sorted(set(match))) if match else ""
|
||||
|
||||
return OllamaChat()
|
||||
|
||||
|
||||
def load_dataset(dataset_path: str):
|
||||
with open(dataset_path, "r") as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
async def batch_inference_dataset(llm, batch: List[Dict[str, Any]], chat=False):
|
||||
results = []
|
||||
llm_inputs = []
|
||||
instruction = """Below are multiple-choice questions concerning cybersecurity.
|
||||
Please select the correct answer.
|
||||
Your exprected answer is JUST THE CORRECT LETTE, not a word or phrase."""
|
||||
|
||||
for row in batch:
|
||||
question_text = "Question: " + row["question"] + " " + " ".join(row["choices"])
|
||||
question_text = question_text.replace("\n", " ")
|
||||
if chat:
|
||||
llm_input = [SystemMessage(content=instruction), HumanMessage(content=question_text)]
|
||||
else:
|
||||
llm_input = instruction + question_text + "\n"
|
||||
llm_inputs.append(llm_input)
|
||||
|
||||
try:
|
||||
llm_outputs = await llm.abatch(llm_inputs)
|
||||
except Exception as e:
|
||||
logging.error(f"Error in batch: {e}")
|
||||
llm_outputs = [f"{e}"] * len(llm_inputs)
|
||||
|
||||
for idx, output in enumerate(llm_outputs):
|
||||
output = output.replace("Answer:", "").strip()
|
||||
row = batch[idx]
|
||||
row["llm_input"] = llm_inputs[idx]
|
||||
row["llm_output"] = output
|
||||
row["llm_answer"] = "".join(sorted(set(re.findall(r"[A-D]", output))))
|
||||
row["score"] = int(row["llm_answer"].lower() == row["answer"].lower())
|
||||
results.append(row)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def inference_dataset(llm, dataset: List[Dict[str, Any]], batch_size: int = 1, chat=False):
|
||||
def chunks(lst, n):
|
||||
for i in range(0, len(lst), n):
|
||||
yield lst[i:i + n]
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
batches = list(chunks(dataset, batch_size))
|
||||
results = []
|
||||
for idx, batch in enumerate(batches):
|
||||
print(f"Processing batch {idx+1}/{len(batches)}")
|
||||
results += loop.run_until_complete(batch_inference_dataset(llm, batch, chat))
|
||||
return results
|
||||
|
||||
|
||||
def count_score_by_topic(dataset: List[Dict[str, Any]]):
|
||||
score_by_topic = {}
|
||||
total_by_topic = {}
|
||||
total_score = 0
|
||||
|
||||
for row in dataset:
|
||||
for topic in row["topics"]:
|
||||
score_by_topic[topic] = score_by_topic.get(topic, 0) + row["score"]
|
||||
total_by_topic[topic] = total_by_topic.get(topic, 0) + 1
|
||||
total_score += row["score"]
|
||||
|
||||
fraction = {k: f"{score_by_topic[k]}/{total_by_topic[k]}" for k in score_by_topic}
|
||||
percent = {k: round(score_by_topic[k] / total_by_topic[k] * 100, 2) for k in score_by_topic}
|
||||
percent["Overall"] = round(total_score / len(dataset) * 100, 2)
|
||||
fraction["Overall"] = f"{total_score}/{len(dataset)}"
|
||||
|
||||
return fraction, percent
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Ollama LLM Benchmark")
|
||||
parser.add_argument("-d", "--dataset_file", type=str, required=True)
|
||||
parser.add_argument("-o", "--output_dir", type=str, default="./results")
|
||||
parser.add_argument("-b", "--batch_size", type=int, default=1)
|
||||
parser.add_argument("-m", "--model", type=str, required=True)
|
||||
parser.add_argument("-c", "--chat", action="store_true", default=False)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
model_id = args.model
|
||||
llm = init_ollama_llm(model_id)
|
||||
|
||||
dataset = load_dataset(args.dataset_file)
|
||||
results = inference_dataset(llm, dataset, batch_size=args.batch_size, chat=args.chat)
|
||||
score_frac, score_pct = count_score_by_topic(results)
|
||||
|
||||
output = {
|
||||
"score_fraction": score_frac,
|
||||
"score_percent": score_pct,
|
||||
"details": results
|
||||
}
|
||||
|
||||
output_path = Path(args.output_dir) / f"{Path(args.dataset_file).stem}_{model_id}.json"
|
||||
with open(output_path, "w") as f:
|
||||
json.dump(output, f, indent=2)
|
||||
print(f"Results written to {output_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
File diff suppressed because one or more lines are too long
Loading…
Reference in New Issue