#!/usr/bin/env python3 from __future__ import annotations import json import tempfile from pathlib import Path from typing import Any import numpy as np from huggingface_hub import HfApi, hf_hub_download from transformers import AutoConfig, AutoTokenizer TOKENIZER_FILES = [ "tokenizer_config.json", "tokenizer.json", "special_tokens_map.json", "vocab.txt", "vocab.json", "merges.txt", "added_tokens.json", "sentencepiece.bpe.model", "spiece.model", ] DEFAULT_LABEL_MAX_SPAN_TOKENS = { # Token-piece limits, not word limits. These need to reflect how the # underlying tokenizer actually fragments compact identifiers. "PPSN": 9, "POSTCODE": 7, "PHONE_NUMBER": 10, "PASSPORT_NUMBER": 8, "BANK_ROUTING_NUMBER": 5, "ACCOUNT_NUMBER": 19, "CREDIT_DEBIT_CARD": 12, "SWIFT_BIC": 8, "EMAIL": 15, "FIRST_NAME": 5, "LAST_NAME": 8, } DEFAULT_LABEL_MIN_NONSPACE_CHARS = { "PPSN": 8, "POSTCODE": 6, "PHONE_NUMBER": 7, "PASSPORT_NUMBER": 7, "BANK_ROUTING_NUMBER": 6, "ACCOUNT_NUMBER": 6, "CREDIT_DEBIT_CARD": 12, "SWIFT_BIC": 8, "EMAIL": 6, "FIRST_NAME": 2, "LAST_NAME": 2, } WHITESPACE_BRIDGE_LABELS = { "PPSN", "POSTCODE", "PHONE_NUMBER", "PASSPORT_NUMBER", "BANK_ROUTING_NUMBER", "ACCOUNT_NUMBER", "CREDIT_DEBIT_CARD", "SWIFT_BIC", "EMAIL", } CONSERVATIVE_BOUNDARY_REFINEMENT_LABELS = { "PPSN", "POSTCODE", "PHONE_NUMBER", "PASSPORT_NUMBER", "BANK_ROUTING_NUMBER", "ACCOUNT_NUMBER", "CREDIT_DEBIT_CARD", "SWIFT_BIC", "EMAIL", } OUTPUT_PRIORITY = { "PPSN": 0, "PASSPORT_NUMBER": 1, "ACCOUNT_NUMBER": 2, "BANK_ROUTING_NUMBER": 3, "CREDIT_DEBIT_CARD": 4, "PHONE_NUMBER": 5, "SWIFT_BIC": 6, "POSTCODE": 7, "EMAIL": 8, "FIRST_NAME": 9, "LAST_NAME": 10, } def normalize_entity_name(label: str) -> str: label = (label or "").strip() if label.startswith("B-") or label.startswith("I-"): label = label[2:] return label.upper() def _sanitize_tokenizer_dir(tokenizer_path: Path) -> str: tokenizer_cfg_path = tokenizer_path / "tokenizer_config.json" if not tokenizer_cfg_path.exists(): return str(tokenizer_path) data = json.loads(tokenizer_cfg_path.read_text(encoding="utf-8")) if "fix_mistral_regex" not in data: return str(tokenizer_path) tmpdir = Path(tempfile.mkdtemp(prefix="openmed_span_tokenizer_")) keep = set(TOKENIZER_FILES) for child in tokenizer_path.iterdir(): if child.is_file() and child.name in keep: (tmpdir / child.name).write_bytes(child.read_bytes()) data.pop("fix_mistral_regex", None) (tmpdir / "tokenizer_config.json").write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") return str(tmpdir) def safe_auto_tokenizer(tokenizer_ref: str): tokenizer_path = Path(tokenizer_ref) if tokenizer_path.exists(): tokenizer_ref = _sanitize_tokenizer_dir(tokenizer_path) else: api = HfApi() files = set(api.list_repo_files(repo_id=tokenizer_ref, repo_type="model")) tmpdir = Path(tempfile.mkdtemp(prefix="openmed_remote_span_tokenizer_")) copied = False for name in TOKENIZER_FILES: if name not in files: continue src = hf_hub_download(repo_id=tokenizer_ref, filename=name, repo_type="model") (tmpdir / Path(name).name).write_bytes(Path(src).read_bytes()) copied = True if copied: tokenizer_ref = _sanitize_tokenizer_dir(tmpdir) try: return AutoTokenizer.from_pretrained(tokenizer_ref, use_fast=True, fix_mistral_regex=True) except Exception: pass try: return AutoTokenizer.from_pretrained(tokenizer_ref, use_fast=True, fix_mistral_regex=False) except TypeError: pass try: return AutoTokenizer.from_pretrained(tokenizer_ref, use_fast=True) except Exception: return AutoTokenizer.from_pretrained(tokenizer_ref, use_fast=False) def label_names_from_config(config) -> list[str]: names = list(getattr(config, "span_label_names", [])) if not names: raise ValueError("Missing span_label_names in config") return [normalize_entity_name(name) for name in names] def label_thresholds_from_config(config, default_threshold: float) -> dict[str, float]: raw = getattr(config, "span_label_thresholds", None) or {} out = {normalize_entity_name(key): float(value) for key, value in raw.items()} for label in label_names_from_config(config): out.setdefault(label, float(default_threshold)) return out def token_label_thresholds_from_config(config, default_threshold: float) -> dict[str, float]: raw = getattr(config, "token_label_thresholds", None) or {} out = {normalize_entity_name(key): float(value) for key, value in raw.items()} for label in label_names_from_config(config): out.setdefault(label, float(default_threshold)) return out def token_extend_thresholds_from_config(config, default_fraction: float = 0.6) -> dict[str, float]: raw = getattr(config, "token_extend_thresholds", None) or {} out = {normalize_entity_name(key): float(value) for key, value in raw.items()} for label in label_names_from_config(config): out.setdefault(label, max(0.0, min(1.0, float(token_label_thresholds_from_config(config, 0.5).get(label, 0.5)) * default_fraction))) return out def boundary_label_thresholds_from_config(config, default_threshold: float = 0.0) -> dict[str, float]: raw = getattr(config, "boundary_label_thresholds", None) or {} out = {normalize_entity_name(key): float(value) for key, value in raw.items()} for label in label_names_from_config(config): out.setdefault(label, float(default_threshold)) return out def label_max_span_tokens_from_config(config) -> dict[str, int]: raw = getattr(config, "span_label_max_span_tokens", None) or {} out = {normalize_entity_name(key): int(value) for key, value in raw.items()} for label, value in DEFAULT_LABEL_MAX_SPAN_TOKENS.items(): out.setdefault(label, value) for label in label_names_from_config(config): out.setdefault(label, 8) return out def label_min_nonspace_chars_from_config(config) -> dict[str, int]: raw = getattr(config, "span_label_min_nonspace_chars", None) or {} out = {normalize_entity_name(key): int(value) for key, value in raw.items()} for label, value in DEFAULT_LABEL_MIN_NONSPACE_CHARS.items(): out.setdefault(label, value) for label in label_names_from_config(config): out.setdefault(label, 1) return out def overlaps(a: dict, b: dict) -> bool: return not (a["end"] <= b["start"] or b["end"] <= a["start"]) def dedupe_spans(spans: list[dict]) -> list[dict]: ordered = sorted( spans, key=lambda item: (-float(item.get("score", 0.0)), item["start"], item["end"], OUTPUT_PRIORITY.get(item["label"], 99)), ) kept = [] for span in ordered: if any(overlaps(span, other) for other in kept): continue kept.append(span) kept.sort(key=lambda item: (item["start"], item["end"], OUTPUT_PRIORITY.get(item["label"], 99))) return kept def _valid_offset(offset: tuple[int, int]) -> bool: return bool(offset) and offset[1] > offset[0] def _has_skippable_bridge(text: str, left: tuple[int, int], right: tuple[int, int], label: str) -> bool: bridge = text[int(left[1]) : int(right[0])] if bridge == "": return True return label in WHITESPACE_BRIDGE_LABELS and bridge.isspace() def _has_left_extension_bridge(text: str, left: tuple[int, int], right: tuple[int, int]) -> bool: bridge = text[int(left[1]) : int(right[0])] return bridge == "" def _nonspace_length(text: str, start: int, end: int) -> int: return sum(0 if ch.isspace() else 1 for ch in text[int(start) : int(end)]) def decode_span_logits( text: str, offsets: list[tuple[int, int]], start_scores: np.ndarray, end_scores: np.ndarray, label_names: list[str], default_threshold: float, label_thresholds: dict[str, float] | None = None, label_max_span_tokens: dict[str, int] | None = None, ) -> list[dict]: thresholds = {label: float(default_threshold) for label in label_names} if label_thresholds: thresholds.update({normalize_entity_name(key): float(value) for key, value in label_thresholds.items()}) max_tokens = dict(DEFAULT_LABEL_MAX_SPAN_TOKENS) if label_max_span_tokens: max_tokens.update({normalize_entity_name(key): int(value) for key, value in label_max_span_tokens.items()}) spans: list[dict] = [] for label_index, label in enumerate(label_names): threshold = thresholds.get(label, float(default_threshold)) max_span = max_tokens.get(label, 8) start_candidates = [idx for idx in range(len(offsets)) if _valid_offset(offsets[idx]) and float(start_scores[idx, label_index]) >= threshold] for start_idx in start_candidates: best = None for end_idx in range(start_idx, min(len(offsets), start_idx + max_span)): if not _valid_offset(offsets[end_idx]): continue end_score = float(end_scores[end_idx, label_index]) if end_score < threshold: continue score = min(float(start_scores[start_idx, label_index]), end_score) if best is None or score > best["score"]: best = { "label": label, "start": int(offsets[start_idx][0]), "end": int(offsets[end_idx][1]), "score": score, } if best is not None and best["start"] < best["end"]: best["text"] = text[best["start"]:best["end"]] spans.append(best) return dedupe_spans(spans) def decode_token_presence_segments( text: str, offsets: list[tuple[int, int]], token_scores: np.ndarray, label_names: list[str], default_threshold: float, label_thresholds: dict[str, float] | None = None, label_extend_thresholds: dict[str, float] | None = None, label_max_span_tokens: dict[str, int] | None = None, label_min_nonspace_chars: dict[str, int] | None = None, boundary_label_thresholds: dict[str, float] | None = None, start_scores: np.ndarray | None = None, end_scores: np.ndarray | None = None, ) -> list[dict]: thresholds = {label: float(default_threshold) for label in label_names} if label_thresholds: thresholds.update({normalize_entity_name(key): float(value) for key, value in label_thresholds.items()}) extend_thresholds = {label: max(0.0, min(1.0, thresholds[label] * 0.6)) for label in label_names} if label_extend_thresholds: extend_thresholds.update({normalize_entity_name(key): float(value) for key, value in label_extend_thresholds.items()}) max_tokens = dict(DEFAULT_LABEL_MAX_SPAN_TOKENS) if label_max_span_tokens: max_tokens.update({normalize_entity_name(key): int(value) for key, value in label_max_span_tokens.items()}) min_nonspace_chars = dict(DEFAULT_LABEL_MIN_NONSPACE_CHARS) if label_min_nonspace_chars: min_nonspace_chars.update({normalize_entity_name(key): int(value) for key, value in label_min_nonspace_chars.items()}) boundary_thresholds = {label: 0.0 for label in label_names} if boundary_label_thresholds: boundary_thresholds.update({normalize_entity_name(key): float(value) for key, value in boundary_label_thresholds.items()}) spans: list[dict] = [] valid = [_valid_offset(offset) for offset in offsets] num_tokens = len(offsets) for label_index, label in enumerate(label_names): threshold = thresholds.get(label, float(default_threshold)) extend_threshold = min(threshold, extend_thresholds.get(label, threshold)) max_span = max_tokens.get(label, 8) idx = 0 while idx < num_tokens: if not valid[idx] or float(token_scores[idx, label_index]) < threshold: idx += 1 continue start_idx = idx end_idx = idx while end_idx + 1 < num_tokens and valid[end_idx + 1] and float(token_scores[end_idx + 1, label_index]) >= threshold and (end_idx + 1 - start_idx + 1) <= max_span: end_idx += 1 while ( start_idx - 1 >= 0 and valid[start_idx - 1] and _has_left_extension_bridge(text, offsets[start_idx - 1], offsets[start_idx]) and float(token_scores[start_idx - 1, label_index]) >= extend_threshold and (end_idx - (start_idx - 1) + 1) <= max_span ): start_idx -= 1 while ( end_idx + 1 < num_tokens and valid[end_idx + 1] and _has_skippable_bridge(text, offsets[end_idx], offsets[end_idx + 1], label) and float(token_scores[end_idx + 1, label_index]) >= extend_threshold and ((end_idx + 1) - start_idx + 1) <= max_span ): end_idx += 1 presence_slice = token_scores[start_idx : end_idx + 1, label_index] score = float(presence_slice.mean()) out_start_idx = start_idx out_end_idx = end_idx if start_scores is not None and end_scores is not None: refine_window = min(3, end_idx - start_idx + 1) start_window = start_scores[start_idx : start_idx + refine_window, label_index] best_start_rel = int(np.argmax(start_window)) best_start_idx = start_idx + best_start_rel end_window_start = max(best_start_idx, end_idx - refine_window + 1) end_window = end_scores[end_window_start : end_idx + 1, label_index] best_end_rel = int(np.argmax(end_window)) best_end_idx = end_window_start + best_end_rel if ( float(start_scores[best_start_idx, label_index]) < boundary_thresholds.get(label, 0.0) or float(end_scores[best_end_idx, label_index]) < boundary_thresholds.get(label, 0.0) ): idx = end_idx + 1 continue out_start_idx = best_start_idx out_end_idx = best_end_idx if label in CONSERVATIVE_BOUNDARY_REFINEMENT_LABELS and ( best_start_idx != start_idx or best_end_idx != end_idx ): outer_boundary = min(float(start_scores[start_idx, label_index]), float(end_scores[end_idx, label_index])) refined_boundary = min( float(start_scores[best_start_idx, label_index]), float(end_scores[best_end_idx, label_index]), ) if refined_boundary < outer_boundary + 0.08: out_start_idx = start_idx out_end_idx = end_idx score = ( 0.65 * score + 0.175 * float(start_scores[out_start_idx, label_index]) + 0.175 * float(end_scores[out_end_idx, label_index]) ) min_chars = int(min_nonspace_chars.get(label, 1)) if _nonspace_length(text, offsets[out_start_idx][0], offsets[out_end_idx][1]) < min_chars: idx = end_idx + 1 continue spans.append( { "label": label, "start": int(offsets[out_start_idx][0]), "end": int(offsets[out_end_idx][1]), "score": score, "text": text[int(offsets[out_start_idx][0]) : int(offsets[out_end_idx][1])], } ) idx = end_idx + 1 return dedupe_spans(spans) def load_onnx_session(model_ref: str, onnx_file: str = "model_quantized.onnx", onnx_subfolder: str = "onnx"): import onnxruntime as ort model_path = Path(model_ref) if model_path.exists(): candidates = [] if onnx_subfolder: candidates.append(model_path / onnx_subfolder / onnx_file) candidates.append(model_path / onnx_file) onnx_path = next((path for path in candidates if path.exists()), candidates[0]) config = AutoConfig.from_pretrained(model_ref) tokenizer = safe_auto_tokenizer(model_ref) else: remote_name = f"{onnx_subfolder}/{onnx_file}" if onnx_subfolder else onnx_file onnx_path = Path(hf_hub_download(repo_id=model_ref, filename=remote_name, repo_type="model")) config = AutoConfig.from_pretrained(model_ref) tokenizer = safe_auto_tokenizer(model_ref) session = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"]) return session, tokenizer, config def run_onnx(session, encoded: dict[str, Any]) -> tuple[np.ndarray, np.ndarray]: feed = {} input_names = {item.name for item in session.get_inputs()} for key, value in encoded.items(): if key == "offset_mapping": continue if key in input_names: feed[key] = value outputs = session.run(None, feed) return outputs[0], outputs[1] def run_onnx_all(session, encoded: dict[str, Any]) -> list[np.ndarray]: feed = {} input_names = {item.name for item in session.get_inputs()} for key, value in encoded.items(): if key == "offset_mapping": continue if key in input_names: feed[key] = value return session.run(None, feed)