inference.py

#!/usr/bin/env python3
"""Standalone inference for the internal three-label PII classifier."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer


LABELS = (
    "privacy_asking_for_pii",
    "privacy_giving_pii",
    "directing_users_off_platform",
)
INSTRUCTION_PREFIX = (
    "Instruct: In the following chat messages from target speaker t and possibly "
    "other speakers s1, s2, etc., detect abuse by speaker t.\nQuery:"
)
INSTRUCTION_SEPARATOR = "\n\n"
TURN_SEPARATOR = " </s> "
TARGET_SPEAKER = "t"
SPEAKER_TEXT_SEPARATOR = ": "
MAX_LENGTH = 512


def _validate_turn(turn: Any, index: int) -> tuple[str, str]:
    if not isinstance(turn, dict):
        raise TypeError(f"conversation turn {index} must be a JSON object")
    if set(turn) != {"speaker", "text"}:
        raise ValueError(
            f"conversation turn {index} must contain exactly 'speaker' and 'text'"
        )
    speaker = turn["speaker"]
    text = turn["text"]
    if not isinstance(speaker, str) or not speaker:
        raise ValueError(f"conversation turn {index} has an invalid speaker")
    if not isinstance(text, str):
        raise ValueError(f"conversation turn {index} has non-string text")
    return speaker, text


def format_conversation(value: str | list[dict[str, str]]) -> str:
    """Format plain text or speaker/text turns exactly like the training formatter."""
    if isinstance(value, str):
        turns: list[dict[str, str]] = [{"speaker": TARGET_SPEAKER, "text": value}]
    elif isinstance(value, list):
        if not value:
            raise ValueError("conversation must contain at least one turn")
        turns = value
    else:
        raise TypeError("input must be plain text or a JSON-list conversation")

    other_speakers: dict[str, str] = {}
    formatted_turns: list[str] = []
    for index, turn in enumerate(turns):
        speaker, text = _validate_turn(turn, index)
        if speaker == TARGET_SPEAKER:
            anonymous_speaker = TARGET_SPEAKER
        else:
            if speaker not in other_speakers:
                other_speakers[speaker] = f"s{len(other_speakers) + 1}"
            anonymous_speaker = other_speakers[speaker]
        formatted_turns.append(
            f"{anonymous_speaker}{SPEAKER_TEXT_SEPARATOR}{text}"
        )

    return (
        INSTRUCTION_PREFIX
        + INSTRUCTION_SEPARATOR
        + TURN_SEPARATOR.join(formatted_turns)
    )


def parse_input(value: str) -> str | list[dict[str, str]]:
    """Interpret a JSON list as turns; otherwise retain the value as plain text."""
    try:
        decoded = json.loads(value)
    except json.JSONDecodeError:
        return value
    if isinstance(decoded, list):
        return decoded
    return value


def predict(
    value: str | list[dict[str, str]],
    model_path: str | Path,
) -> dict[str, float]:
    """Return uncalibrated sigmoid probabilities in the fixed three-label order."""
    formatted = format_conversation(value)
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    tokenizer.truncation_side = "left"
    model = AutoModelForSequenceClassification.from_pretrained(model_path)
    model.eval()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    encoded = tokenizer(
        formatted,
        padding="max_length",
        max_length=MAX_LENGTH,
        truncation=True,
        return_tensors="pt",
    )
    encoded = {name: tensor.to(device) for name, tensor in encoded.items()}

    with torch.inference_mode():
        logits = model(**encoded).logits
        if logits.ndim != 2 or logits.shape[0] != 1 or logits.shape[1] != len(LABELS):
            raise ValueError(
                f"expected logits shape (1, {len(LABELS)}), got {tuple(logits.shape)}"
            )
        probabilities = torch.sigmoid(logits[0]).float().cpu().tolist()
    return {label: probability for label, probability in zip(LABELS, probabilities)}


def _arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run the internal three-label PII classifier."
    )
    inputs = parser.add_mutually_exclusive_group(required=True)
    inputs.add_argument("--text", help="Plain text or a JSON-list conversation.")
    inputs.add_argument(
        "--input-file",
        type=Path,
        help="UTF-8 file containing plain text or a JSON-list conversation.",
    )
    parser.add_argument("--model-path", type=Path, default=Path(__file__).parent)
    return parser.parse_args()


def main() -> None:
    args = _arguments()
    raw_value = (
        args.text
        if args.text is not None
        else args.input_file.read_text(encoding="utf-8")
    )
    result = predict(parse_input(raw_value), args.model_path)
    print(json.dumps(result, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()