| |
| """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) |
| 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: |
| 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() |
|
|
