| from __future__ import annotations
|
|
|
| from typing import Any, Sequence
|
|
|
| import torch
|
| from transformers import AutoModel, BitsAndBytesConfig, MistralCommonBackend
|
|
|
|
|
| class MinistralDecisionInference:
|
| def __init__(self, model, tokenizer):
|
| self.model = model.eval()
|
| self.tokenizer = tokenizer
|
| self.device = model.get_input_embeddings().weight.device
|
| self.temperature = float(model.decision_temperature)
|
|
|
| @classmethod
|
| def from_pretrained(cls, model_id: str = ".", *, quantized: bool = True):
|
| quantization = None
|
| if quantized:
|
| quantization = BitsAndBytesConfig(
|
| load_in_4bit=True,
|
| bnb_4bit_quant_type="nf4",
|
| bnb_4bit_use_double_quant=True,
|
| bnb_4bit_compute_dtype=torch.float16,
|
| )
|
| model = AutoModel.from_pretrained(
|
| model_id,
|
| trust_remote_code=True,
|
| quantization_config=quantization,
|
| device_map={"": 0} if torch.cuda.is_available() else "cpu",
|
| low_cpu_mem_usage=True,
|
| )
|
| tokenizer = MistralCommonBackend.from_pretrained(model_id)
|
| return cls(model, tokenizer)
|
|
|
| @staticmethod
|
| def _pair_text(state: str, question: str, option: dict[str, str]) -> str:
|
| criteria = option.get("criteria", "").strip()
|
| criterion_text = f"\nCandidate criteria: {criteria}" if criteria else ""
|
| return (
|
| "Choose the best answer for this decision.\n"
|
| f"Task: {question}\n"
|
| f"Context: {state}\n"
|
| f"Candidate answer: {option['label']}{criterion_text}\n"
|
| "Does this candidate correctly answer the task?"
|
| )
|
|
|
| def score(self, state: str, question: str, options: Sequence[dict[str, str]], *, abstain_below: float | None = None) -> dict[str, Any]:
|
| if len(options) < 2:
|
| raise ValueError("Provide at least two options")
|
| ids = [option["id"] for option in options]
|
| if len(ids) != len(set(ids)):
|
| raise ValueError("Option ids must be unique")
|
| rows = []
|
| for option in options:
|
| encoded = self.tokenizer.encode(self._pair_text(state, question, option), return_tensors="pt")
|
| input_ids = encoded.input_ids if hasattr(encoded, "input_ids") else encoded
|
| rows.append(input_ids[0] if input_ids.ndim == 2 else input_ids)
|
| pad_id = getattr(self.tokenizer, "pad_token_id", None)
|
| if pad_id is None:
|
| pad_id = getattr(self.tokenizer, "eos_token_id", 0) or 0
|
| scores = []
|
| for start in range(0, len(rows), 4):
|
| batch = rows[start:start + 4]
|
| lengths = torch.tensor([row.numel() for row in batch], dtype=torch.long)
|
| width = int(lengths.max().item())
|
| input_ids = torch.full((len(batch), width), int(pad_id), dtype=torch.long)
|
| mask = torch.zeros((len(batch), width), dtype=torch.long)
|
| for index, row in enumerate(batch):
|
| input_ids[index, :row.numel()] = row
|
| mask[index, :row.numel()] = 1
|
| input_ids, mask = input_ids.to(self.device), mask.to(self.device)
|
| with torch.inference_mode():
|
| scores.extend(self.model.score_candidate_pairs(input_ids, mask).float().cpu().tolist())
|
| probabilities = torch.softmax(torch.tensor(scores) / self.temperature, dim=0).tolist()
|
| best = max(range(len(options)), key=probabilities.__getitem__)
|
| confidence = probabilities[best]
|
| abstain = abstain_below is not None and confidence < abstain_below
|
| return {
|
| "choice": None if abstain else ids[best],
|
| "confidence": confidence,
|
| "probabilities": dict(zip(ids, probabilities)),
|
| "scores": dict(zip(ids, scores)),
|
| "abstain": abstain,
|
| "calibrated_temperature": self.temperature,
|
| }
|
|
|
|
|
