fix: filter truncated candidates from batched best-of-n sampling

optillm/bon.py

@@ -37,7 +37,11 @@ def best_of_n_sampling(system_prompt: str, initial_query: str, client, model: st
         if response is None or not response.choices:
             raise Exception("Response is None or has no choices")
 
-        completions = [choice.message.content for choice in response.choices if choice.message.content is not None]
+        completions = [
+            choice.message.content for choice in response.choices
+            if choice.message.content and choice.message.content.strip()
+            and choice.finish_reason != "length"
+        ]
         logger.info(f"Generated {len(completions)} initial completions using n parameter. Tokens used: {response.usage.completion_tokens}")
         bon_completion_tokens += response.usage.completion_tokens
 
@@ -68,7 +72,8 @@ def best_of_n_sampling(system_prompt: str, initial_query: str, client, model: st
                 # Check for valid response with None-checking
                 if (response is None or
                     not response.choices or
-                    response.choices[0].message.content is None or
+                    not response.choices[0].message.content or
+                    not response.choices[0].message.content.strip() or
                     response.choices[0].finish_reason == "length"):
                     logger.warning(f"Completion {i+1}/{n} truncated or empty, skipping")
                     continue

tests/test_bon_candidate_filtering.py

@@ -0,0 +1,57 @@
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+
+from optillm.bon import best_of_n_sampling
+
+
+def response(*choices):
+    return SimpleNamespace(
+        choices=[
+            SimpleNamespace(message=SimpleNamespace(content=text), finish_reason=reason)
+            for text, reason in choices
+        ],
+        usage=SimpleNamespace(completion_tokens=1),
+    )
+
+
+@pytest.mark.parametrize(
+    "invalid_text, reason", [("incomplete answer", "length"), ("   ", "stop")]
+)
+def test_batched_invalid_candidates_are_never_sent_to_judge(invalid_text, reason):
+    rated = []
+
+    def create(**kwargs):
+        if kwargs.get("n") == 2:
+            return response((invalid_text, reason), ("complete answer", "stop"))
+        candidate = kwargs["messages"][-2]["content"]
+        rated.append(candidate)
+        return response(("10" if candidate == invalid_text else "8", "stop"))
+
+    client = SimpleNamespace(
+        chat=SimpleNamespace(completions=SimpleNamespace(create=create))
+    )
+    answer, tokens = best_of_n_sampling("system", "question", client, "model", n=2)
+
+    assert answer == "complete answer"
+    assert rated == ["complete answer"]
+    assert tokens == 2
+
+
+def test_all_invalid_batch_uses_existing_sequential_fallback():
+    create = Mock(
+        side_effect=[
+            response(("partial", "length")),
+            response(("complete answer", "stop")),
+            response(("9", "stop")),
+        ]
+    )
+    client = SimpleNamespace(
+        chat=SimpleNamespace(completions=SimpleNamespace(create=create))
+    )
+    answer, tokens = best_of_n_sampling("system", "question", client, "model", n=1)
+
+    assert answer == "complete answer"
+    assert "n" not in create.call_args_list[1].kwargs
+    assert tokens == 3