| |
| """CERTA certify-or-escalate inference handler (weights-aware). |
| |
| Load behaviour: |
| * If the repository contains baked quantized weights (`*.safetensors`), load them **directly** from the |
| repo folder. A model saved after bitsandbytes NF4/INT8 quantization carries its `quantization_config` |
| in `config.json`, so `from_pretrained(path)` reconstructs the quantized model without re-downloading or |
| re-quantizing the base. |
| * Otherwise fall back to the wrapper form: download the base model named in `certa_config.json` and |
| quantize it at load time. |
| |
| For each accept/reject decision it serves the quantized decision only when its margin clears the CERTA |
| threshold for its predicted class at the requested risk level alpha; otherwise it returns escalate=True |
| (refer the decision to the full-precision reference). Requires a CUDA GPU + bitsandbytes for NF4/INT8. |
| |
| Scope: the thresholds certify the clinical accept/reject verification decision on the PMC-Patients |
| calibration distribution. They carry NO finite-sample guarantee on a different task or population; |
| re-calibrate for a new distribution. RTN variants (bfloat16 storage) are not runnable through this handler. |
| """ |
| import json, os, glob, torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig |
|
|
|
|
| class EndpointHandler: |
| def __init__(self, path="."): |
| self.cfg = json.load(open(os.path.join(path, "certa_config.json"))) |
| self.th = json.load(open(os.path.join(path, "certa_thresholds.json"))) |
| self.alpha = str(self.cfg.get("default_alpha", 0.05)) |
| base, q, bits = self.cfg["base_model"], self.cfg["quantizer"], self.cfg["bits"] |
| has_local_weights = bool(glob.glob(os.path.join(path, "*.safetensors"))) |
|
|
| if has_local_weights: |
| |
| self.model = AutoModelForCausalLM.from_pretrained(path, device_map="auto") |
| self.tok = AutoTokenizer.from_pretrained(path) |
| elif q == "bnb" and bits == 4: |
| qc = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", |
| bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16) |
| self.model = AutoModelForCausalLM.from_pretrained(base, quantization_config=qc, device_map="auto") |
| self.tok = AutoTokenizer.from_pretrained(base) |
| elif q == "bnb" and bits == 8: |
| self.model = AutoModelForCausalLM.from_pretrained( |
| base, quantization_config=BitsAndBytesConfig(load_in_8bit=True), device_map="auto") |
| self.tok = AutoTokenizer.from_pretrained(base) |
| else: |
| raise NotImplementedError( |
| "This is an RTN evaluation variant (bfloat16 storage, not memory-reducing); " |
| "it is provided for reproducibility and is not runnable through this handler.") |
|
|
| self.model.eval() |
| yes, no = self.th["answer_tokens"]["YES"], self.th["answer_tokens"]["NO"] |
| self.yid = self.tok.encode(yes, add_special_tokens=False)[-1] |
| self.nid = self.tok.encode(no, add_special_tokens=False)[-1] |
|
|
| @torch.no_grad() |
| def __call__(self, data): |
| prompt = data["inputs"] if isinstance(data, dict) else data |
| alpha = str(data.get("alpha", self.alpha)) if isinstance(data, dict) else self.alpha |
| th = self.th["alpha"][alpha]["class_conditional"] |
| ids = self.tok(prompt, return_tensors="pt").to(self.model.device) |
| logits = self.model(**ids).logits[0, -1] |
| zy, zn = float(logits[self.yid]), float(logits[self.nid]) |
| decision = "YES" if zy >= zn else "NO" |
| margin = abs(zy - zn) |
| tau = th["tau_yes"] if decision == "YES" else th["tau_no"] |
| escalate = (tau is None) or (margin < tau) |
| return {"decision": decision, "margin": round(margin, 4), "threshold": tau, "alpha": float(alpha), |
| "escalate": escalate, "served_by": "reference" if escalate else "quantized", |
| "note": "escalate=True: refer this decision to the full-precision reference model"} |
|
|
