#!/usr/bin/env python3 """Smoke-test client for the /v1/scoring endpoint: full decode from raw logits to PII spans and masked text. Label layout (do not reorder): 0 = O; then per category c in 0..8: 1+4c = B, 2+4c = I, 3+4c = E, 4+4c = S. Decode: constrained Viterbi over RAW logits with start mask (O/B/S), end mask (O/E/S) and 0/-inf transitions; span walk B..I..E / S; span score = mean softmax probability of the assigned labels; sensitivity = sigmoid(ranking_scores[0]), threshold 0.5. Usage: uv run --with requests --with tokenizers --with numpy python test_client.py "text" """ import argparse import math from pathlib import Path import numpy as np import requests from tokenizers import Tokenizer CATEGORIES: tuple[str, ...] = ( "private_person", "private_email", "private_phone", "private_address", "private_url", "private_date", "account_number", "secret", "other_pii", ) NUM_LABELS = 1 + 4 * len(CATEGORIES) # 37 SENSITIVITY_THRESHOLD = 0.5 NEG_INF = -1e9 def _tag(idx: int) -> str | None: return None if idx == 0 else "BIES"[(idx - 1) % 4] def _cat(idx: int) -> int | None: return None if idx == 0 else (idx - 1) // 4 def build_tables() -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Build BIOES start/end constraints and the transition matrix.""" start = np.full(NUM_LABELS, NEG_INF, dtype=np.float32) end = np.full(NUM_LABELS, NEG_INF, dtype=np.float32) trans = np.full((NUM_LABELS, NUM_LABELS), NEG_INF, dtype=np.float32) for i in range(NUM_LABELS): if _tag(i) in (None, "B", "S"): start[i] = 0.0 if _tag(i) in (None, "E", "S"): end[i] = 0.0 for j in range(NUM_LABELS): if _tag(i) in (None, "E", "S"): ok = _tag(j) in (None, "B", "S") else: # B-X / I-X must continue or close the same category ok = _cat(i) == _cat(j) and _tag(j) in ("I", "E") if ok: trans[i, j] = 0.0 return start, end, trans def viterbi_path(logits: np.ndarray) -> list[int]: """Decode the highest-scoring valid BIOES path from raw logits.""" start, end, trans = build_tables() seq_len = logits.shape[0] scores = logits[0] + start backptr = np.empty((seq_len - 1, NUM_LABELS), dtype=np.int64) for step in range(1, seq_len): cand = scores[:, None] + trans backptr[step - 1] = cand.argmax(axis=0) scores = cand.max(axis=0) + logits[step] last = int((scores + end).argmax()) path = [last] for step in range(seq_len - 2, -1, -1): last = int(backptr[step, last]) path.append(last) return path[::-1] def extract_spans( path: list[int], offsets: list[tuple[int, int]], probs: np.ndarray, text: str, ) -> list[dict]: """Extract character spans from the decoded BIOES path.""" def span_score(tok_range: range) -> float: return float(np.mean([probs[i, path[i]] for i in tok_range])) def make(start_char: int, end_char: int, cat: int, score: float) -> dict | None: if start_char >= end_char: return None return { "start": start_char, "end": end_char, "pii_type": CATEGORIES[cat], "score": score, "text": text[start_char:end_char], } spans: list[dict] = [] t = 0 while t < len(path): tag, cat = _tag(path[t]), _cat(path[t]) char_start, char_end = offsets[t] if tag == "S": assert cat is not None span = make(char_start, char_end, cat, span_score(range(t, t + 1))) if span: spans.append(span) t += 1 elif tag == "B": assert cat is not None begin_tok, span_end = t, char_end t += 1 while t < len(path): inner_tag, inner_cat = _tag(path[t]), _cat(path[t]) if inner_tag == "I" and inner_cat == cat: span_end = offsets[t][1] t += 1 elif inner_tag == "E" and inner_cat == cat: span_end = offsets[t][1] t += 1 break else: break span = make(char_start, span_end, cat, span_score(range(begin_tok, t))) if span: spans.append(span) else: t += 1 return spans def mask_text(text: str, spans: list[dict]) -> str: masked = text for span in sorted(spans, key=lambda s: s["start"], reverse=True): masked = ( masked[: span["start"]] + f"[{span['pii_type'].upper()}]" + masked[span["end"] :] ) return masked def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("text", nargs="?", default="My name is John Smith and my SSN is 123-45-6789.") parser.add_argument("--url", default="http://localhost:8002") parser.add_argument("--model-dir", default=str(Path(__file__).resolve().parents[1])) args = parser.parse_args() resp = requests.post( f"{args.url}/v1/scoring", json={"model": "pii-masking-latest", "sequences": [args.text]}, headers={"authorization": "bearer local-dev"}, timeout=60, ) resp.raise_for_status() payload = resp.json() logits = np.asarray(payload["pruning_scores"][0], dtype=np.float32) ranking = payload["ranking_scores"][0] server_ids = payload["token_ids"][0] assert logits.shape[1] == NUM_LABELS, logits.shape tokenizer = Tokenizer.from_file(str(Path(args.model_dir) / "tokenizer.json")) encoding = tokenizer.encode(args.text, add_special_tokens=False) if list(encoding.ids) != list(server_ids): raise SystemExit( f"alignment error: local {len(encoding.ids)} ids != server {len(server_ids)}" ) path = viterbi_path(logits) probs = np.exp(logits - logits.max(axis=-1, keepdims=True)) probs /= probs.sum(axis=-1, keepdims=True) spans = extract_spans(path, list(encoding.offsets), probs, args.text) sensitivity = 1 / (1 + math.exp(-ranking[0])) print(f"sensitivity: p={sensitivity:.4f} " f"(logit {ranking[0]:+.3f}, is_sensitive={sensitivity >= SENSITIVITY_THRESHOLD})") print(f"masked: {mask_text(args.text, spans)}") for span in spans: print( f" [{span['start']:>4}:{span['end']:<4}] " f"{span['pii_type']:<16} p={span['score']:.4f} {span['text']!r}" ) if not spans: print(" (no PII spans)") if __name__ == "__main__": main()