| """Official-style Qwen3 reranker readout adapted to declared decision options.""" |
|
|
| from __future__ import annotations |
|
|
| import inspect |
| import time |
|
|
| from .core import digest, softmax, validate_row |
|
|
| PREFIX = ( |
| '<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the ' |
| 'Instruct provided. Note that the answer can only be "yes" or "no".<|im_end|>\n' |
| '<|im_start|>user\n' |
| ) |
| SUFFIX = '<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n' |
| DECISION_INSTRUCTION = ( |
| "Given evidence and one possible answer to a question, determine whether the evidence supports that answer " |
| "under the question's criterion. Use only the supplied evidence." |
| ) |
| RETRIEVAL_INSTRUCTION = ( |
| "Given a search query, determine whether the document is relevant and contains evidence that answers the query." |
| ) |
| PROMPT_VERSION = "qwen3-reranker-native-options-v1" |
|
|
|
|
| def _encode(tokenizer, row, option, max_tokens): |
| experiment = row.get("provenance", {}).get("experiment") |
| instruction = RETRIEVAL_INSTRUCTION if experiment in {"code-rag", "company-brain"} else DECISION_INSTRUCTION |
| body = ( |
| f"<Instruct>: {instruction}\n" |
| f"<Query>: Question: {row['question']}\nCandidate answer: {option['description']}\n" |
| f"<Document>: {row['state']}" |
| ) |
| text = PREFIX + body + SUFFIX |
| ids = tokenizer.encode(text, add_special_tokens=False) |
| if not ids or len(ids) > max_tokens: |
| raise ValueError(f"Row {row['id']} option {option['id']}: {len(ids)} tokens exceed limit {max_tokens}") |
| return ids, digest(text) |
|
|
|
|
| def _answer_ids(tokenizer): |
| no = tokenizer.encode("no", add_special_tokens=False) |
| yes = tokenizer.encode("yes", add_special_tokens=False) |
| if len(no) != 1 or len(yes) != 1 or no == yes: |
| raise ValueError("Reranker yes/no answers must be distinct single tokens") |
| if tokenizer.convert_tokens_to_ids("yes") != yes[0] or tokenizer.convert_tokens_to_ids("no") != no[0]: |
| raise ValueError("Tokenizer conversion differs from the official yes/no token contract") |
| return no[0], yes[0] |
|
|
|
|
| def score_pair_batch(model, tokenizer, specs, max_tokens: int = 4096): |
| """Score independent row/option relevance pairs in one native batch.""" |
| import torch |
|
|
| if not specs: |
| raise ValueError("Pair batch is empty") |
| encoded = [_encode(tokenizer, row, option, max_tokens) for row, option in specs] |
| width = max(len(ids) for ids, _ in encoded) |
| pad = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id |
| if pad is None: |
| raise ValueError("Tokenizer requires padding or EOS token") |
| device = next(model.parameters()).device |
| inputs = { |
| "input_ids": torch.tensor([[pad] * (width - len(ids)) + ids for ids, _ in encoded], device=device), |
| "attention_mask": torch.tensor( |
| [[0] * (width - len(ids)) + [1] * len(ids) for ids, _ in encoded], device=device |
| ), |
| } |
| kwargs = dict(inputs, use_cache=False, return_dict=True) |
| if "logits_to_keep" in inspect.signature(model.forward).parameters: |
| kwargs["logits_to_keep"] = 1 |
| no_id, yes_id = _answer_ids(tokenizer) |
| if device.type == "cuda": |
| torch.cuda.synchronize(device) |
| mark = time.perf_counter() |
| with torch.inference_mode(): |
| logits = model(**kwargs).logits[:, -1, :].float() |
| if device.type == "cuda": |
| torch.cuda.synchronize(device) |
| elapsed = time.perf_counter() - mark |
| selected = logits[:, [no_id, yes_id]] |
| odds = selected[:, 1] - selected[:, 0] |
| relevance = selected.softmax(-1)[:, 1] |
| return [ |
| { |
| "option_id": option["id"], |
| "log_odds": float(odds[index]), |
| "binary_relevance": float(relevance[index]), |
| "input_tokens": len(encoded[index][0]), |
| "prompt_sha256": encoded[index][1], |
| } |
| for index, (_, option) in enumerate(specs) |
| ], { |
| "pair_batch_size": len(specs), |
| "forward_seconds": elapsed, |
| "padded_tokens": int(inputs["input_ids"].numel()), |
| } |
|
|
|
|
| def score(model, tokenizer, row: dict, metadata: dict, max_tokens: int = 4096) -> dict: |
| import torch |
|
|
| validate_row(row) |
| started = time.perf_counter() |
| scored, timing = score_pair_batch( |
| model, tokenizer, [(row, option) for option in row["options"]], max_tokens |
| ) |
| log_odds = [item["log_odds"] for item in scored] |
| binary = [item["binary_relevance"] for item in scored] |
| return { |
| "id": row["id"], |
| "option_ids": [option["id"] for option in row["options"]], |
| "probabilities": softmax(log_odds), |
| "option_logits": log_odds, |
| "independent_binary_relevance": binary, |
| "input_tokens": sum(item["input_tokens"] for item in scored), |
| "max_option_input_tokens": max(item["input_tokens"] for item in scored), |
| "forward_seconds": timing["forward_seconds"], |
| "total_seconds": time.perf_counter() - started, |
| "option_prompt_sha256": [item["prompt_sha256"] for item in scored], |
| "pair_batches": [timing], |
| "prompt_version": PROMPT_VERSION, |
| "model": metadata, |
| "readout": "native yes/no log-odds per option, normalized only for relative comparison", |
| "probability_status": "relative option compatibility; uncalibrated as categorical probability", |
| } |
|
|
