diff --git a/optillm/bon.py b/optillm/bon.py index e22ee18..f410d6e 100644 --- a/optillm/bon.py +++ b/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 diff --git a/tests/test_bon_candidate_filtering.py b/tests/test_bon_candidate_filtering.py new file mode 100644 index 0000000..8cb6fda --- /dev/null +++ b/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