diff --git a/benchmarks/seceval/eval/eval.py b/benchmarks/seceval/eval/eval.py index cd73df71..02aea7c1 100644 --- a/benchmarks/seceval/eval/eval.py +++ b/benchmarks/seceval/eval/eval.py @@ -61,8 +61,10 @@ logging.basicConfig( logger.addHandler(logging.FileHandler(f"./eval-{int(time.time())}.log", "w")) # Instruction template for the LLM -instruction = """Below are multiple-choice questions concerning cybersecurity. - Please select the correct letter. Your exprected answer is JUST THE CORRECT LETTER, not a word or phrase.""" +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. +""" def init_hf_llm(model_id: str) -> HuggingFacePipeline: @@ -161,15 +163,6 @@ def init_ollama_llm(model_id: str) -> 'OllamaChat': """ class OllamaChat: async def abatch(self, prompts: List[str]) -> List[str]: - """ - Process a batch of prompts asynchronously. - - Args: - prompts: List of prompts to process - - Returns: - List[str]: List of model responses - """ responses = [] for prompt in prompts: try: @@ -195,21 +188,52 @@ def init_ollama_llm(model_id: str) -> 'OllamaChat': return responses def extract_answer(self, text: str) -> str: - """ - Extract answer letter from model response. - - Args: - text: Model response text - - Returns: - str: Extracted answer letter(s) - """ match = re.findall(r"[A-D]", text.upper()) return "".join(sorted(set(match))) if match else "" return OllamaChat() +def init_openrouter_llm(model_id: str): + class OpenRouterChat: + async def abatch(self, prompts: List[str]): + responses = [] + for prompt in prompts: + try: + api_base = os.getenv("OPENROUTER_API_BASE", "https://openrouter.ai/api/v1/chat/completions") + api_key = os.getenv("OPENROUTER_API_KEY") + + if not api_key: + raise ValueError("OPENROUTER_API_KEY is not defined in the environment variables.") + + completion = litellm.completion( + model=model_id, + messages=[{"role": "user", "content": prompt}], + api_base=api_base, + api_key=api_key, + custom_llm_provider="openrouter" + ) + + 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("Formato de respuesta incorrecto.") + responses.append("Error: No se pudo extraer resultado") + except Exception as e: + logging.error(f"OpenRouter 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 OpenRouterChat() + + def load_dataset(dataset_path: str) -> List[Dict[str, Any]]: """ Load evaluation dataset from JSON file. @@ -338,7 +362,7 @@ def main(): 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, ollama: ollama model") + parser.add_argument("-B", "--backend", type=str, choices=["remote_hf", "azure", "textgen", "local_hf", "ollama", "openrouter"], 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, ollama: ollama model, openrouter: openrouter model") parser.add_argument("-m", "--models", type=str, nargs="+", required=True, help="Specify the models.") args = parser.parse_args() @@ -362,6 +386,8 @@ def main(): llm = init_azure_openai_llm(model_id) elif args.backend == "ollama": llm = init_ollama_llm(model_id) + elif args.backend == "openrouter": + llm = init_openrouter_llm(model_id) else: raise RuntimeError("Unknown backend")