| """Restricted next-token decision probabilities for OR-decision-4B.""" |
|
|
| import string |
|
|
| import torch |
|
|
|
|
| @torch.inference_mode() |
| def decision_probabilities(model, tokenizer, context, instructions, options, kind="choice"): |
| """Options are (label, description) pairs; return probabilities by label.""" |
| if not 1 <= len(options) <= 62: |
| raise ValueError("Provide 1 to 62 options") |
| if len({label for label, _ in options}) != len(options): |
| raise ValueError("Option labels must be unique") |
| identifiers = (string.ascii_uppercase + string.ascii_lowercase + string.digits)[:len(options)] |
| system = ( |
| "Choose the option that best answers the decision using the supplied context and " |
| "instructions. Reply with only its choice identifier, such as choice#A." |
| ) |
| choices = "\n".join( |
| f"choice#{identifier}: {label}\n{description}" |
| for identifier, (label, description) in zip(identifiers, options, strict=True) |
| ) |
| user = ( |
| f"Decision type: {kind}\n\nContext:\n{context}" |
| f"\n\nInstructions:\n{instructions}\n\nOptions:\n{choices}" |
| ) |
| prompt = tokenizer.apply_chat_template( |
| [{"role": "system", "content": system}, {"role": "user", "content": user}], |
| tokenize=False, add_generation_prompt=True, enable_thinking=False, |
| ) + "choice#" |
| input_ids = tokenizer.encode(prompt, add_special_tokens=False) |
| if len(input_ids) + 1 > 4096: |
| raise ValueError("Decision exceeds the evaluated 4096-token window") |
| choice_ids = [] |
| for identifier in identifiers: |
| tokens = tokenizer.encode(identifier, add_special_tokens=False) |
| if len(tokens) != 1 or tokenizer.encode(prompt + identifier, add_special_tokens=False) != input_ids + tokens: |
| raise ValueError("Tokenizer does not preserve single-token choice identifiers") |
| choice_ids.append(tokens[0]) |
| inputs = torch.tensor([input_ids], device=model.device) |
| logits = model(input_ids=inputs, attention_mask=torch.ones_like(inputs)).logits[0, -1].float() |
| probabilities = torch.softmax(logits[choice_ids], dim=-1).cpu().tolist() |
| return {label: probability for (label, _), probability in zip(options, probabilities, strict=True)} |
|
|
