predict.py

#!/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()