class WinnowStrategy(ReadStrategy):
"""Read trained candidate labels without applying a generic chat template."""
def __init__(self, context):
super().__init__(context)
self.tokenizer = context.online_renderer.renderer.get_tokenizer()
if self.tokenizer.bos_token_id is None:
raise ValueError("Winnow requires a BOS token")
self.label_ids: list[int] = []
for label in LABELS:
ids = self.tokenizer.encode(label, add_special_tokens=False)
if (
len(ids) != 1
or ids[0] in self.label_ids
or self.tokenizer.decode(ids) != label
):
raise ValueError("Winnow requires distinct single-token letter labels")
self.label_ids.append(ids[0])
config = context.engine_client.model_config.hf_config
self.state_format = getattr(config, "decision_state_format", "text")
if self.state_format not in {"text", "json"}:
raise ValueError("decision_state_format must be text or json")
self.temperature = getattr(config, "decision_temperature", 1.0)
if not math.isfinite(self.temperature) or self.temperature <= 0:
raise ValueError("decision_temperature must be finite and positive")
def limits(self):
return DecisionLimits(max_questions=64, max_options=len(LABELS))
async def read(
self,
questions,
instructions,
state,
*,
request_id,
chat_template_kwargs,
lora_request,
priority,
):
if instructions or chat_template_kwargs:
raise StructuredDecisionError(
"Winnow uses a fixed trained prompt; chat overrides are unsupported"
)
if self.state_format == "json":
try:
state = json.loads(state)
except (ValueError, TypeError) as exc:
raise StructuredDecisionError(
"Winnow input must contain valid JSON"
) from exc
prompts = [prompt_tokens(self.tokenizer, state, q) for q in questions]
maximum = self.context.engine_client.model_config.max_model_len
if any(len(prompt) + 1 > maximum for prompt in prompts):
raise StructuredDecisionError("Winnow decision exceeds the model context")
slots = [self.label_ids[: len(q.options)] for q in questions]
reads = await next_token_label_reads(
self.context.engine_client,
[tokens_input(p) for p in prompts],
[
SamplingParams(max_tokens=1, temperature=0.0, logprob_token_ids=ids)
for ids in slots
],
request_id,
lora_request=lora_request,
priority=priority,
)
return [
QuestionRead(
probs=label_softmax([lp / self.temperature for lp in read.logprobs]),
label_mass=sum(math.exp(lp) for lp in read.logprobs),
confidence=math.exp(max(read.logprobs)),
argmax_is_label=bool(read.result.outputs[0].token_ids)
and read.result.outputs[0].token_ids[0] in ids,
input_tokens=len(read.result.prompt_token_ids or ()),
output_tokens=len(read.result.outputs[0].token_ids),
cached_tokens=read.result.num_cached_tokens or 0,
cache_write_tokens=read.result.num_cache_creation_tokens or 0,
)
for ids, read in zip(slots, reads)
]