Use the reference environment's uv.lock

.gitignore

@@ -50,4 +50,3 @@ __pycache__
 
 docs/build
 site/
-uv.lock

oe_eval/configs/task_suites.py

@@ -399,6 +399,10 @@
     "tasks": [f"mmlu_pro_{cat}:cot::none" for cat in MMLU_PRO_CATEGORIES],
     "primary_metric": "micro",
 }
+TASK_SUITE_CONFIGS["mmlu_pro:cot::qwen3_stopfix"] = {
+    "tasks": [f"mmlu_pro_{cat}:cot::qwen3_stopfix" for cat in MMLU_PRO_CATEGORIES],
+    "primary_metric": "micro",
+}
 TASK_SUITE_CONFIGS["mmlu_pro:cot::llama3.1"] = {
     "tasks": [f"mmlu_pro_{cat}:cot::llama3.1" for cat in MMLU_PRO_CATEGORIES],
     "primary_metric": "micro",

oe_eval/configs/tasks.py

@@ -1226,6 +1226,68 @@
                 "regimes": ["Llama-3"],
             },
         },
+        "gpqa:5shot_cot::qwen3_official_prompt": {
+            "task_name": "gpqa",
+            "dataset_name": "gpqa_main",
+            "split": "train",
+            "primary_metric": "exact_match",
+            "num_shots": 5,
+            "fewshot_source": "Original:GPQA",
+            "generation_kwargs": {
+                "max_gen_toks": 2048,
+                "do_sample": False,
+                "temperature": 0.0,
+                "stop_sequences": [
+                    "\n\nQuestion:",
+                    "\nQuestion:",
+                    "You are an AI assistant",
+                    "</s>",
+                    "<|im_end|>",
+                ],
+            },
+            "context_kwargs": {
+                "answer_shuffling_seed": 111,
+                "description": "Here are some example questions from experts. An explanation is given before the final answer. Answer the final question yourself, giving your reasoning beforehand.\n",
+                "final_description": "\nGive step by step reasoning before you answer, and when you're ready to answer, please use the format \"The correct answer is (insert answer here)\":\n",
+            },
+            "metric_kwargs": {
+                "answer_format_regex": "The correct answer is \\(([A-D])\\)",
+            },
+            "metadata": {
+                "regimes": [],
+            },
+        },
+        "gpqa_diamond:5shot_cot::qwen3_official_prompt": {
+            "task_name": "gpqa",
+            "dataset_name": "gpqa_diamond",
+            "split": "train",
+            "primary_metric": "exact_match",
+            "num_shots": 5,
+            "fewshot_source": "Original:GPQA",
+            "generation_kwargs": {
+                "max_gen_toks": 2048,
+                "do_sample": False,
+                "temperature": 0.0,
+                "stop_sequences": [
+                    "\n\nQuestion:",
+                    "\nQuestion:",
+                    "You are an AI assistant",
+                    "</s>",
+                    "<|im_end|>",
+                ],
+            },
+            "context_kwargs": {
+                "answer_shuffling_seed": 111,
+                "description": "Here are some example questions from experts. An explanation is given before the final answer. Answer the final question yourself, giving your reasoning beforehand.\n",
+                "final_description": "\nGive step by step reasoning before you answer, and when you're ready to answer, please use the format \"The correct answer is (insert answer here)\":\n",
+            },
+            "metric_kwargs": {
+                "answer_format_regex": "The correct answer is \\(([A-D])\\)",
+            },
+            "metadata": {
+                "regimes": [],
+            },
+        },
         "gpqa:0shot_cot::tulu3": {
             "task_name": "gpqa",
             "split": "train",
@@ -1653,6 +1715,33 @@
             "num_shots": 3,  # like deepseek
             "limit": 500,
         },
+        "mbpp:3shot::qwen3_signature_tests": {
+            "task_name": "mbpp",
+            "primary_metric": "pass_at_1",
+            "generation_kwargs": {
+                "stop_sequences": ["```", '\n"""', "\nassert", "\n#", "\n<|/", "<|eot_id|>"],
+                "do_sample": True,
+                "top_p": 0.95,
+                "temperature": 0.2,
+                "repeats": 20,
+            },
+            "use_chat_format": False,
+            "context_kwargs": {
+                "prompt_variant": "qwen3_signature_tests",
+                "assistant_prefix": None,
+            },
+            "metric_kwargs": {
+                "pass_at_ks": [1],
+                "n_exe_workers": 20,
+                "rich_exec_info": True,
+            },
+            "compute_gold_bpb": False,
+            "num_shots": 3,
+            "limit": 500,
+            "metadata": {
+                "regimes": [],
+            },
+        },
         "mbpp:3shot::olmo3": {
             "task_name": "mbpp",
             "primary_metric": "pass_at_1",
@@ -4942,6 +5031,20 @@
             "regimes": [],
         },
     }
+    TASK_CONFIGS[f"mmlu_pro_{cat}:cot::qwen3_stopfix"] = {
+        "task_name": f"mmlu_pro_{cat}:cot",
+        "split": "test",
+        "num_shots": 5,
+        "generation_kwargs": {
+            "max_gen_toks": 1024,
+            "do_sample": False,
+            "temperature": 0.0,
+            "stop_sequences": ["Question:"],
+        },
+        "metadata": {
+            "regimes": [],
+        },
+    }
     TASK_CONFIGS[f"mmlu_pro_{cat}:cot::llama3.1"] = {
         "task_name": f"mmlu_pro_{cat}:cot",
         "split": "test",
@@ -7092,6 +7195,127 @@
                 "regimes": [],
             },
         },
+        "cruxeval_output::qwen3_greedy": {
+            "task_name": "cruxeval_output",
+            "primary_metric": "pass_at_1",
+            "generation_kwargs": {
+                "do_sample": False,
+                "temperature": 0.0,
+                "repeats": 1,
+                "stop_sequences": ["\nassert", "\ndef", "\n```", '\n"""', "\n#", "\n\n"],
+            },
+            "metric_kwargs": {
+                "pass_at_ks": [1],
+                "n_exe_workers": 20,
+                "rich_exec_info": True,
+            },
+            "metadata": {
+                "regimes": [],
+            },
+        },
+        "cruxeval_output::official_direct": {
+            "task_name": "cruxeval_output",
+            "primary_metric": "pass_at_1",
+            "num_shots": 0,
+            "generation_kwargs": {
+                "max_gen_toks": 512,
+                "do_sample": False,
+                "temperature": 0.0,
+                "repeats": 1,
+                "stop_sequences": ["[/ANSWER]", "</s>", "<|im_end|>"],
+            },
+            "context_kwargs": {
+                "description": "",
+                "prompt_variant": "official_direct",
+            },
+            "metric_kwargs": {
+                "pass_at_ks": [1],
+                "n_exe_workers": 20,
+                "rich_exec_info": True,
+            },
+            "metadata": {
+                "regimes": [],
+            },
+        },
+        "cruxeval_output::official_cot_1shot": {
+            "task_name": "cruxeval_output",
+            "primary_metric": "pass_at_1",
+            "num_shots": 0,
+            "generation_kwargs": {
+                "max_gen_toks": 2048,
+                "do_sample": True,
+                "temperature": 0.2,
+                "top_p": 0.95,
+                "repeats": 10,
+                "stop_sequences": ["[/ANSWER]", "</s>", "<|im_end|>"],
+            },
+            "context_kwargs": {
+                "description": "",
+                "prompt_variant": "official_cot_output",
+            },
+            "metric_kwargs": {
+                "pass_at_ks": [1, 5],
+                "n_exe_workers": 20,
+                "rich_exec_info": True,
+            },
+            "metadata": {
+                "regimes": [],
+            },
+        },
+        "cruxeval_output::lm_eval_2shot": {
+            "task_name": "cruxeval_output",
+            "primary_metric": "pass_at_1",
+            "num_shots": 2,
+            "fewshot_source": "Original:CruxEvalOutput:lm_eval",
+            "generation_kwargs": {
+                "max_gen_toks": 1024,
+                "do_sample": True,
+                "temperature": 0.2,
+                "top_p": 0.95,
+                "repeats": 10,
+                "stop_sequences": ["[/ANSWER]", "</seed:think>[ANSWER]", "</s>", "<|im_end|>"],
+            },
+            "context_kwargs": {
+                "description": "You are given a Python function and an assertion containing an input to the function. Complete the assertion with a literal (no unsimplified expressions, no function calls) containing the output when executing the provided code on the given input, even if the function is incorrect or incomplete. Do NOT output any extra information. Provide the full assertion with the correct output in [ANSWER] and [/ANSWER] tags, following the examples.\n",
+                "final_description": "\n[ANSWER]\n",
+                "prompt_variant": "lm_eval_direct",
+            },
+            "metric_kwargs": {
+                "pass_at_ks": [1],
+                "n_exe_workers": 20,
+                "rich_exec_info": True,
+            },
+            "metadata": {
+                "regimes": [],
+            },
+        },
+        "cruxeval_output::lm_eval_1shot": {
+            "task_name": "cruxeval_output",
+            "primary_metric": "pass_at_1",
+            "num_shots": 1,
+            "fewshot_source": "Original:CruxEvalOutput:lm_eval",
+            "generation_kwargs": {
+                "max_gen_toks": 1024,
+                "do_sample": True,
+                "temperature": 0.2,
+                "top_p": 0.95,
+                "repeats": 10,
+                "stop_sequences": ["[/ANSWER]", "</seed:think>[ANSWER]", "</s>", "<|im_end|>"],
+            },
+            "context_kwargs": {
+                "description": "You are given a Python function and an assertion containing an input to the function. Complete the assertion with a literal (no unsimplified expressions, no function calls) containing the output when executing the provided code on the given input, even if the function is incorrect or incomplete. Do NOT output any extra information. Provide the full assertion with the correct output in [ANSWER] and [/ANSWER] tags, following the examples.\n",
+                "final_description": "\n[ANSWER]\n",
+                "prompt_variant": "lm_eval_direct",
+            },
+            "metric_kwargs": {
+                "pass_at_ks": [1],
+                "n_exe_workers": 20,
+                "rich_exec_info": True,
+            },
+            "metadata": {
+                "regimes": [],
+            },
+        },
         "cruxeval_output:pass@5": {
             "task_name": "cruxeval_output",
             "primary_metric": "pass_at_5",

oe_eval/models/eleuther_vllm_causallms.py

@@ -7,7 +7,9 @@
 import importlib_metadata
 import torch
 import transformers
+from more_itertools import distribute
 from lm_eval.models.utils import Collator
+from lm_eval.models.utils import undistribute
 from lm_eval.models.vllm_causallms import VLLM
 from packaging import version
 from tqdm import tqdm
@@ -111,6 +113,61 @@ def unload_model(self):
         except Exception:
             pass
 
+    @staticmethod
+    def _as_vllm_token_prompts(requests: List[List[int]]):
+        from vllm import TokensPrompt
+
+        return [TokensPrompt(prompt_token_ids=list(request)) for request in requests]
+
+    def _model_generate(
+        self,
+        requests: List[List[int]] = None,
+        generate: bool = False,
+        max_tokens: int = None,
+        stop: Optional[List[str]] = None,
+        **kwargs,
+    ):
+        from vllm import LLM
+        from vllm import SamplingParams
+
+        requests = requests or []
+        if generate:
+            kwargs = self.modify_gen_kwargs(kwargs)
+            sampling_params = SamplingParams(max_tokens=max_tokens, stop=stop, **kwargs)
+        else:
+            sampling_params = SamplingParams(
+                temperature=0, prompt_logprobs=1, max_tokens=1, detokenize=False
+            )
+
+        if self.data_parallel_size > 1:
+            import ray
+
+            @ray.remote
+            def run_inference_one_model(
+                model_args: dict, sampling_params, requests: List[List[int]]
+            ):
+                llm = LLM(**model_args)
+                return llm.generate(
+                    prompts=VLLM_Verbose._as_vllm_token_prompts(requests),
+                    sampling_params=sampling_params,
+                )
+
+            requests = [list(x) for x in distribute(self.data_parallel_size, requests)]
+            inputs = ((self.model_args, sampling_params, req) for req in requests)
+            object_refs = [run_inference_one_model.remote(*x) for x in inputs]
+            results = ray.get(object_refs)
+            ray.shutdown()
+            return undistribute(results)
+
+        generate_kwargs = {
+            "prompts": self._as_vllm_token_prompts(requests),
+            "sampling_params": sampling_params,
+            "use_tqdm": True if self.batch_size == "auto" else False,
+        }
+        if self.lora_request is not None:
+            generate_kwargs["lora_request"] = self.lora_request
+        return self.model.generate(**generate_kwargs)
+
     def generate_until_verbose(
         self, requests: List[GenerateUntilRequest], disable_tqdm: bool = False
     ) -> List[dict]:

oe_eval/tasks/fewshot_sources.py

@@ -2498,6 +2498,21 @@
     },
 ]
 
+FEWSHOT_SOURCES["Original:CruxEvalOutput:lm_eval"] = [
+    {
+        "id": "1",
+        "code": "def f(n):\n    return n",
+        "input": "17",
+        "output": "17",
+    },
+    {
+        "id": "2",
+        "code": 'def f(s):\n    return s + "a"',
+        "input": '"x9j"',
+        "output": '"x9ja"',
+    },
+]
+
 FEWSHOT_SOURCES["Llama3:squad2"] = [
     {
         "id": "56d43ce42ccc5a1400d830b5",