File size: 3,187 Bytes
dfdef9d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | #!/usr/bin/env python3
"""Run calibrated retrieval-value inference for this model."""
from __future__ import annotations
import argparse
import json
import math
from pathlib import Path
from typing import Any
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
DEFAULT_MODEL = "jansowa/dev-knowledge-bullshit-detector-v0.1"
def _sigmoid(value: float) -> float:
if value >= 0:
exponent = math.exp(-value)
return 1.0 / (1.0 + exponent)
exponent = math.exp(value)
return exponent / (1.0 + exponent)
def _calibrate_boundary(score: float, calibration: dict[str, Any]) -> float:
boundary = calibration.get("boundary_ge_3", {})
if boundary.get("method") != "platt-logit":
return score
clipped = min(1.0 - 1e-12, max(1e-12, score))
logit = math.log(clipped / (1.0 - clipped))
return _sigmoid(float(boundary["slope"]) * logit + float(boundary["intercept"]))
def predict(texts: list[str], model_id: str = DEFAULT_MODEL) -> list[dict[str, Any]]:
"""Predict ordered retrieval value and calibrated usefulness measures."""
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSequenceClassification.from_pretrained(model_id)
model.eval()
model_path = Path(model_id)
calibration_path = model_path / "calibration.json"
if calibration_path.is_file():
calibration = json.loads(calibration_path.read_text(encoding="utf-8"))
else:
from huggingface_hub import hf_hub_download
downloaded = hf_hub_download(repo_id=model_id, filename="calibration.json")
calibration = json.loads(Path(downloaded).read_text(encoding="utf-8"))
encoded = tokenizer(
texts,
padding=True,
truncation=True,
max_length=256,
return_tensors="pt",
)
with torch.inference_mode():
logits = model(**encoded).logits
temperature = float(calibration.get("temperature", 1.0))
probabilities = torch.softmax(logits / temperature, dim=-1).tolist()
results = []
for text, distribution in zip(texts, probabilities, strict=True):
retrieval_value = sum(index * probability for index, probability in enumerate(distribution))
raw_ge_3 = sum(distribution[3:])
results.append(
{
"text": text,
"predicted_retrieval_value": max(range(5), key=distribution.__getitem__),
"retrieval_value": retrieval_value,
"usefulness_score": retrieval_value / 4.0,
"probability_retrieval_value_ge_3": _calibrate_boundary(raw_ge_3, calibration),
"probabilities": {str(i): value for i, value in enumerate(distribution)},
}
)
return results
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("text", nargs="+", help="One or more technical comments")
parser.add_argument("--model", default=DEFAULT_MODEL, help="Hub model ID or local directory")
args = parser.parse_args()
print(json.dumps(predict(args.text, args.model), ensure_ascii=False, indent=2))
if __name__ == "__main__":
main()
|