| """Standalone runtime for the merged Pluto-2.1 decision model.""" |
|
|
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| import os |
| import random |
| from pathlib import Path |
| from typing import Any, Mapping, Sequence |
|
|
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| SYSTEM_PROMPT = ( |
| "Choose exactly one supplied decision token. Return only that token. " |
| "Do not explain the answer." |
| ) |
| DECISION_TYPES = frozenset({"choice", "score", "noul"}) |
|
|
|
|
| def _canonical_json(value: Any) -> str: |
| return json.dumps(value, ensure_ascii=False, separators=(",", ":"), sort_keys=True) |
|
|
|
|
| def _compact_criteria(mapped_criteria: Sequence[Mapping[str, Any]]): |
| parsed_descriptions = [] |
| description_fields = None |
| for criterion in mapped_criteria: |
| try: |
| parsed = json.loads(str(criterion["d"])) |
| except json.JSONDecodeError: |
| parsed = None |
| if not isinstance(parsed, dict) or not all(isinstance(key, str) for key in parsed): |
| description_fields = None |
| break |
| fields = sorted(parsed) |
| if description_fields is None: |
| description_fields = fields |
| elif fields != description_fields: |
| description_fields = None |
| break |
| parsed_descriptions.append(parsed) |
| if description_fields is not None and len(parsed_descriptions) == len(mapped_criteria): |
| rows = [ |
| [criterion["t"], criterion["o"], criterion["n"], *(parsed[field] for field in description_fields)] |
| for criterion, parsed in zip(mapped_criteria, parsed_descriptions, strict=True) |
| ] |
| return ["t", "o", "n", *description_fields], rows |
| return ["t", "o", "n", "d"], [ |
| [criterion["t"], criterion["o"], criterion["n"], criterion["d"]] |
| for criterion in mapped_criteria |
| ] |
|
|
|
|
| def render_decision_prompt( |
| *, |
| state: Any, |
| instructions: str, |
| candidates: Sequence[Mapping[str, Any]], |
| decision_type: str, |
| decision_tokens: Sequence[str], |
| seed: int, |
| row_id: str, |
| prompt_schema: str, |
| ): |
| candidate_count = len(candidates) |
| if not 2 <= candidate_count <= 255: |
| raise ValueError("Pluto-2.1 supports 2 to 255 candidates") |
| if decision_type not in DECISION_TYPES: |
| raise ValueError("decision_type must be choice, score, or noul") |
| if len(decision_tokens) < candidate_count: |
| raise ValueError("not enough decision tokens") |
|
|
| presentation_order = list(range(candidate_count)) |
| digest = hashlib.sha256(f"{seed}:{row_id}:candidate-order".encode()).digest() |
| random.Random(int.from_bytes(digest[:8], "big")).shuffle(presentation_order) |
| mapped = [] |
| candidate_to_slot = {} |
| for slot, candidate_index in enumerate(presentation_order): |
| candidate = candidates[candidate_index] |
| candidate_id = str(candidate["id"]) |
| candidate_to_slot[candidate_id] = slot |
| mapped.append( |
| { |
| "t": decision_tokens[slot], |
| "o": candidate_index, |
| "n": str(candidate["label"]), |
| "d": candidate.get("description") or "", |
| } |
| ) |
| fields, criteria = _compact_criteria(mapped) |
| content = _canonical_json( |
| { |
| "schema": prompt_schema, |
| "state": state, |
| "question": { |
| "instructions": instructions, |
| "type": decision_type, |
| "criteria_fields": fields, |
| "criteria": criteria, |
| }, |
| } |
| ) |
| return content, presentation_order, candidate_to_slot |
|
|
|
|
| class Pluto21: |
| def __init__(self, path: str, device: str | None = "cuda"): |
| self.path = Path(path) |
| with open(self.path / "pluto21_config.json", encoding="utf-8") as handle: |
| self.config = json.load(handle) |
| self.model = AutoModelForCausalLM.from_pretrained( |
| path, |
| dtype=torch.bfloat16, |
| device_map=device, |
| ).eval() |
| tokenizer_path = self.path / "tokenizer" |
| self.tokenizer = AutoTokenizer.from_pretrained( |
| str(tokenizer_path if tokenizer_path.is_dir() else path) |
| ) |
| self.tokenizer.padding_side = "left" |
| self.decision_tokens = [ |
| f"<|decision_{index:03d}|>" |
| for index in range(int(self.config["decision_token_count"])) |
| ] |
| ids = [self.tokenizer.convert_tokens_to_ids(token) for token in self.decision_tokens] |
| if any(token_id is None for token_id in ids): |
| raise RuntimeError("tokenizer is missing the decision tokens") |
| self.decision_token_ids = torch.tensor(ids, dtype=torch.long) |
|
|
| def predict( |
| self, |
| *, |
| state: Any, |
| instructions: str, |
| candidates: Sequence[Mapping[str, Any]], |
| decision_type: str = "choice", |
| seed: int = 0, |
| row_id: str = "0", |
| ) -> dict[str, Any]: |
| content, presentation_order, candidate_to_slot = render_decision_prompt( |
| state=state, |
| instructions=instructions, |
| candidates=candidates, |
| decision_type=decision_type, |
| decision_tokens=self.decision_tokens, |
| seed=seed, |
| row_id=row_id, |
| prompt_schema=self.config["prompt_schema"], |
| ) |
| prompt = self.tokenizer.apply_chat_template( |
| [ |
| {"role": "system", "content": SYSTEM_PROMPT}, |
| {"role": "user", "content": content}, |
| ], |
| tokenize=False, |
| add_generation_prompt=True, |
| enable_thinking=self.config.get("decision_prompt_enable_thinking", False), |
| ) |
| encoded = self.tokenizer(str(prompt), return_tensors="pt", truncation=False) |
| device = self.model.device |
| ids = encoded["input_ids"].to(device) |
| mask = encoded["attention_mask"].to(device) |
| with torch.inference_mode(): |
| logits = self.model(input_ids=ids, attention_mask=mask).logits[0, -1] |
| logits = logits.index_select(0, self.decision_token_ids.to(device)) |
| slot_logits = logits[: len(candidates)].float() |
| slot_probs = torch.softmax(slot_logits, dim=-1) |
| probabilities = torch.empty_like(slot_probs) |
| for slot, candidate_index in enumerate(presentation_order): |
| probabilities[candidate_index] = slot_probs[slot] |
| return { |
| "content": content, |
| "prompt": str(prompt), |
| "candidate_to_slot": candidate_to_slot, |
| "slot_logits": slot_logits.detach().cpu().tolist(), |
| "probabilities": probabilities.detach().cpu().tolist(), |
| } |
|
|
|
|
| __all__ = ["Pluto21", "render_decision_prompt"] |
|
|
