Point mlx-lm to fix/float32-logprobs branch for float32 log_softmax fix

EXO passes repetition_context_size=None to mlx_lm which does
tokens[-None:] causing TypeError. Now automatically sets context
size to 64 when repetition_penalty is provided.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

bench/exo_eval.py

@@ -85,6 +85,9 @@ _MC_PATTERNS: list[re.Pattern[str]] = [
 # Code extraction: last ```python ... ``` block (AA regex)
 _CODE_BLOCK_RE = re.compile(r"```(?:python|Python)?\s*\n(.*?)```", re.DOTALL)
 
+# LCB-compatible extraction mode (--lcb-compat): use line-based extraction
+# matching official LiveCodeBench extract_code() behavior.
+
 
 # ---------------------------------------------------------------------------
 # Model config loading
@@ -148,12 +151,22 @@ def extract_boxed_answer(text: str) -> str | None:
     return matches[-1].strip() if matches else None
 
 
-def extract_code_block(text: str, preserve_indent: bool = False) -> str | None:
+def extract_code_block(text: str, preserve_indent: bool = False, lcb_compat: bool = False) -> str | None:
     """Extract the last Python code block from markdown response.
 
     If preserve_indent is True, only strip trailing whitespace (keeps leading
     indentation intact — needed for HumanEval function-body completions).
+
+    If lcb_compat is True, use the official LiveCodeBench extract_code() logic:
+    line-based search for ```, extract between last two backtick lines.
     """
+    if lcb_compat:
+        lines = text.split("\n")
+        backtick_lines = [i for i, line in enumerate(lines) if "```" in line]
+        if len(backtick_lines) < 2:
+            return ""
+        return "\n".join(lines[backtick_lines[-2] + 1 : backtick_lines[-1]])
+
     matches = _CODE_BLOCK_RE.findall(text)
     if matches:
         raw = matches[-1]
@@ -384,7 +397,7 @@ _LCB_WITH_STARTER = (
     "### Format: You will use the following starter code to write the "
     "solution to the problem and enclose your code within delimiters.\n"
     "```python\n{starter_code}\n```\n\n"
-    "### Answer: (use the provided format with backticks)\n"
+    "### Answer: (use the provided format with backticks)\n\n"
 )
 
 _LCB_WITHOUT_STARTER = (
@@ -395,7 +408,7 @@ _LCB_WITHOUT_STARTER = (
     "python program runs, it reads the inputs, runs the algorithm and "
     "writes output to STDOUT.\n"
     "```python\n# YOUR CODE HERE\n```\n\n"
-    "### Answer: (use the provided format with backticks)\n"
+    "### Answer: (use the provided format with backticks)\n\n"
 )
 
 
@@ -529,6 +542,11 @@ async def _call_api(
     system_message: str | None = None,
     reasoning_effort: str | None = None,
     top_p: float | None = None,
+    enable_thinking: bool | None = None,
+    top_k: int | None = None,
+    min_p: float | None = None,
+    repetition_penalty: float | None = None,
+    repetition_context_size: int | None = None,
 ) -> ApiResult:
     messages = []
     if system_message:
@@ -545,6 +563,15 @@ async def _call_api(
         body["reasoning_effort"] = reasoning_effort
     if top_p is not None:
         body["top_p"] = top_p
+    if enable_thinking is not None:
+        body["enable_thinking"] = enable_thinking
+    if top_k is not None:
+        body["top_k"] = top_k
+    if min_p is not None:
+        body["min_p"] = min_p
+    if repetition_penalty is not None:
+        body["repetition_penalty"] = repetition_penalty
+        body["repetition_context_size"] = repetition_context_size or 64
 
     resp = await client.post(
         f"{base_url}/v1/chat/completions",
@@ -577,6 +604,11 @@ async def call_with_retries(
     system_message: str | None = None,
     reasoning_effort: str | None = None,
     top_p: float | None = None,
+    enable_thinking: bool | None = None,
+    top_k: int | None = None,
+    min_p: float | None = None,
+    repetition_penalty: float | None = None,
+    repetition_context_size: int | None = None,
 ) -> ApiResult | None:
     for attempt in range(MAX_RETRIES):
         try:
@@ -591,6 +623,11 @@ async def call_with_retries(
                 system_message,
                 reasoning_effort,
                 top_p,
+                enable_thinking,
+                top_k,
+                min_p,
+                repetition_penalty,
+                repetition_context_size,
             )
         except Exception as e:
             if attempt < MAX_RETRIES - 1:
@@ -621,6 +658,14 @@ async def evaluate_benchmark(
     reasoning_effort: str | None = None,
     top_p: float | None = None,
     difficulty: str | None = None,
+    start_index: int | None = None,
+    end_index: int | None = None,
+    lcb_compat: bool = False,
+    enable_thinking: bool | None = None,
+    top_k: int | None = None,
+    min_p: float | None = None,
+    repetition_penalty: float | None = None,
+    repetition_context_size: int | None = None,
 ) -> list[QuestionResult]:
     """Run a benchmark. Returns per-question results."""
     import datasets
@@ -635,6 +680,8 @@ async def evaluate_benchmark(
                 data_files="hf://datasets/livecodebench/code_generation_lite/*.jsonl",
                 split="train",
             )
+            # Sort by question_id to match official LCB runner ordering
+            ds = ds.sort("question_id")
         else:
             ds = datasets.load_dataset(
                 config.dataset_name,
@@ -651,6 +698,12 @@ async def evaluate_benchmark(
         ds = ds.filter(lambda x: x["difficulty"] == difficulty)
         logger.info(f"Filtered to {len(ds)} {difficulty} problems")
 
+    if start_index is not None or end_index is not None:
+        si = start_index or 0
+        ei = end_index or len(ds)
+        ds = ds.select(range(si, min(ei, len(ds))))
+        logger.info(f"Sliced to [{si}:{ei}] → {len(ds)} problems")
+
     total = len(ds)
     if limit and limit < total:
         ds = ds.select(range(limit))
@@ -708,6 +761,11 @@ async def evaluate_benchmark(
                 system_message=system_msg,
                 reasoning_effort=reasoning_effort,
                 top_p=top_p,
+                enable_thinking=enable_thinking,
+                top_k=top_k,
+                min_p=min_p,
+                repetition_penalty=repetition_penalty,
+                repetition_context_size=repetition_context_size,
             )
             elapsed = time.monotonic() - t0
 
@@ -759,8 +817,9 @@ async def evaluate_benchmark(
             elif config.kind == "code":
                 # HumanEval needs preserved indentation (function body completion)
                 keep_indent = benchmark_name == "humaneval"
-                code = extract_code_block(response, preserve_indent=keep_indent)
-                if code is None:
+                code = extract_code_block(response, preserve_indent=keep_indent,
+                                          lcb_compat=(lcb_compat and benchmark_name == "livecodebench"))
+                if not code:
                     result = QuestionResult(
                         question_id=idx,
                         prompt=prompt,
@@ -1093,6 +1152,18 @@ def main() -> int:
         default=None,
         help="Max questions per benchmark (for fast iteration).",
     )
+    ap.add_argument(
+        "--start-index",
+        type=int,
+        default=None,
+        help="Start index for problem range (inclusive). Applied after difficulty filter.",
+    )
+    ap.add_argument(
+        "--end-index",
+        type=int,
+        default=None,
+        help="End index for problem range (exclusive). Applied after difficulty filter.",
+    )
 
     reasoning_group = ap.add_mutually_exclusive_group()
     reasoning_group.add_argument(
@@ -1155,6 +1226,26 @@ def main() -> int:
         action="store_true",
         help="Skip exo instance management (assumes model is already running).",
     )
+    ap.add_argument(
+        "--lcb-compat",
+        action="store_true",
+        help="Use LiveCodeBench-compatible code extraction and prompt format.",
+    )
+    ap.add_argument(
+        "--enable-thinking",
+        action="store_true",
+        default=False,
+        help="Send enable_thinking=true in API request (for Qwen/DeepSeek thinking mode).",
+    )
+    ap.add_argument("--top-k", type=int, default=None, help="Override top_k sampling.")
+    ap.add_argument("--min-p", type=float, default=None, help="Override min_p sampling.")
+    ap.add_argument(
+        "--repetition-penalty", type=float, default=None, help="Override repetition_penalty."
+    )
+    ap.add_argument(
+        "--repetition-context-size", type=int, default=None,
+        help="Context window for repetition penalty (default: 64 when repetition_penalty is set).",
+    )
 
     args, _ = ap.parse_known_args()
 
@@ -1286,12 +1377,29 @@ def main() -> int:
         reasoning_effort = str(cfg["reasoning_effort"])
     else:
         reasoning_effort = "high" if is_reasoning else None
+
+    enable_thinking = True if args.enable_thinking else None
+
+    top_k: int | None = args.top_k
+    min_p: float | None = args.min_p
+    repetition_penalty: float | None = args.repetition_penalty
+    repetition_context_size: int | None = args.repetition_context_size
+
     base_url = f"http://{args.host}:{args.port}"
 
     logger.info(f"Model: {full_model_id}")
+    extra_params = ""
+    if top_p is not None:
+        extra_params += f"top_p={top_p}, "
+    if top_k is not None:
+        extra_params += f"top_k={top_k}, "
+    if min_p is not None:
+        extra_params += f"min_p={min_p}, "
+    if repetition_penalty is not None:
+        extra_params += f"repetition_penalty={repetition_penalty}, "
     logger.info(
         f"Settings: temperature={temperature}, max_tokens={max_tokens}, "
-        + (f"top_p={top_p}, " if top_p is not None else "")
+        + extra_params
         + f"reasoning={'yes' if is_reasoning else 'no'}"
         + (f", reasoning_effort={reasoning_effort}" if reasoning_effort else "")
     )
@@ -1317,6 +1425,14 @@ def main() -> int:
                             reasoning_effort=reasoning_effort,
                             top_p=top_p,
                             difficulty=args.difficulty,
+                            start_index=args.start_index,
+                            end_index=args.end_index,
+                            lcb_compat=args.lcb_compat,
+                            enable_thinking=enable_thinking,
+                            top_k=top_k,
+                            min_p=min_p,
+                            repetition_penalty=repetition_penalty,
+                            repetition_context_size=repetition_context_size,
                         )
                     )
                     if results:
@@ -1347,6 +1463,14 @@ def main() -> int:
                         reasoning_effort=reasoning_effort,
                         top_p=top_p,
                         difficulty=args.difficulty,
+                        start_index=args.start_index,
+                        end_index=args.end_index,
+                        lcb_compat=args.lcb_compat,
+                        enable_thinking=enable_thinking,
+                        top_k=top_k,
+                        min_p=min_p,
+                        repetition_penalty=repetition_penalty,
+                        repetition_context_size=repetition_context_size,
                     )
                 )
                 if results:

pyproject.toml

@@ -62,7 +62,7 @@ members = ["rust/exo_pyo3_bindings", "bench"]
 [tool.uv.sources]
 exo_pyo3_bindings = { workspace = true }
 mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
-mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/eval-left-padding-in-batched-rotation" }
+mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "fix/float32-logprobs" }
 # Uncomment to use local mlx/mlx-lm development versions:
 # mlx = { path = "/Users/Shared/mlx", editable=true }
 # mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true }