| |
| import regex as re |
| import torch |
|
|
| from raw_word_aligned import word_aligned_ppsn_spans |
|
|
|
|
| TOKEN_RE = re.compile(r"[A-Za-z0-9]+|[^\w\s]", re.UNICODE) |
| PHONE_RE = re.compile(r"^(?:\+353\s?(?:\(0\))?\s?\d(?:[\s-]?\d){7,8}|0\d(?:[\s-]?\d){7,8})$") |
| PASSPORT_RE = re.compile(r"^[A-Z]{2}\s?\d{7}$") |
| SORT_RE = re.compile(r"^(?:\d{6}|\d{2}[ -]\d{2}[ -]\d{2})$") |
| IBAN_IE_RE = re.compile(r"^IE\d{2}(?:\s?[A-Z]{4})(?:\s?\d{4}){3}\s?\d{2}$") |
| BIC_RE = re.compile(r"^[A-Z]{4}[A-Z]{2}[A-Z0-9]{2}(?:[A-Z0-9]{3})?$") |
| EIRCODE_RE = re.compile(r"^(?:[ACDEFHKNPRTVWXY]\d{2}|D6W)\s?[0-9ACDEFHKNPRTVWXY]{4}$", re.IGNORECASE) |
|
|
| DEFAULT_LABEL_THRESHOLDS = { |
| "PHONE_NUMBER": 0.35, |
| "PASSPORT_NUMBER": 0.11, |
| "BANK_ROUTING_NUMBER": 0.35, |
| "ACCOUNT_NUMBER": 0.40, |
| "CREDIT_DEBIT_CARD": 0.08, |
| "SWIFT_BIC": 0.50, |
| } |
|
|
| FORMAT_LABELS = set(DEFAULT_LABEL_THRESHOLDS) |
| 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 tokenize_with_spans(text: str): |
| return [(m.group(0), m.start(), m.end()) for m in TOKEN_RE.finditer(text)] |
|
|
|
|
| def normalize_label(label: str) -> str: |
| label = (label or "").strip() |
| if label.startswith("B-") or label.startswith("I-"): |
| label = label[2:] |
| return label.upper() |
|
|
|
|
| def luhn_ok(value: str) -> bool: |
| digits = "".join(ch for ch in value if ch.isdigit()) |
| if not (13 <= len(digits) <= 19): |
| return False |
| total = 0 |
| double = False |
| for ch in reversed(digits): |
| number = int(ch) |
| if double: |
| number *= 2 |
| if number > 9: |
| number -= 9 |
| total += number |
| double = not double |
| return total % 10 == 0 |
|
|
|
|
| def plausible_label_text(label: str, value: str) -> bool: |
| value = value.strip() |
| if label == "PHONE_NUMBER": |
| return PHONE_RE.match(value) is not None |
| if label == "PASSPORT_NUMBER": |
| return PASSPORT_RE.match(value) is not None |
| if label == "BANK_ROUTING_NUMBER": |
| return SORT_RE.match(value) is not None |
| if label == "ACCOUNT_NUMBER": |
| compact = value.replace(" ", "") |
| return IBAN_IE_RE.match(value) is not None or (compact.isdigit() and len(compact) == 8) |
| if label == "CREDIT_DEBIT_CARD": |
| return luhn_ok(value) |
| if label == "SWIFT_BIC": |
| return BIC_RE.match(value) is not None |
| if label == "POSTCODE": |
| return EIRCODE_RE.match(value) is not None |
| return True |
|
|
|
|
| def label_ids_from_mapping(id2label, label: str): |
| target = label.upper() |
| ids = [] |
| for raw_id, raw_label in id2label.items(): |
| if normalize_label(str(raw_label)) == target: |
| ids.append(int(raw_id)) |
| return ids |
|
|
|
|
| def label_ids(model, label: str): |
| return label_ids_from_mapping(model.config.id2label, label) |
|
|
|
|
| def word_scores_for_label(text: str, model, tokenizer, label: str): |
| pieces = tokenize_with_spans(text) |
| if not pieces: |
| return pieces, [] |
| words = [word for word, _, _ in pieces] |
| encoded = tokenizer(words, is_split_into_words=True, return_tensors="pt", truncation=True) |
| word_ids = encoded.word_ids(batch_index=0) |
| device = next(model.parameters()).device |
| encoded = {key: value.to(device) for key, value in encoded.items()} |
| with torch.no_grad(): |
| logits = model(**encoded).logits[0] |
| probs = torch.softmax(logits, dim=-1) |
| ids = label_ids(model, label) |
| scores = [] |
| for word_index in range(len(pieces)): |
| score = 0.0 |
| for token_index, wid in enumerate(word_ids): |
| if wid != word_index: |
| continue |
| for label_id in ids: |
| score = max(score, float(probs[token_index, label_id])) |
| scores.append(score) |
| return pieces, scores |
|
|
|
|
| def word_scores_for_label_onnx(text: str, session, tokenizer, config, label: str): |
| from onnx_token_classifier import _run_onnx, _softmax |
|
|
| pieces = tokenize_with_spans(text) |
| if not pieces: |
| return pieces, [] |
| words = [word for word, _, _ in pieces] |
| encoded = tokenizer(words, is_split_into_words=True, return_tensors="np", truncation=True) |
| word_ids = encoded.word_ids(batch_index=0) |
| logits = _run_onnx(session, encoded)[0] |
| probs = _softmax(logits, axis=-1) |
| ids = label_ids_from_mapping(config.id2label, label) |
| scores = [] |
| for word_index in range(len(pieces)): |
| score = 0.0 |
| for token_index, wid in enumerate(word_ids): |
| if wid != word_index: |
| continue |
| for label_id in ids: |
| score = max(score, float(probs[token_index, label_id])) |
| scores.append(score) |
| return pieces, scores |
|
|
|
|
| def _word_aligned_label_spans_from_scores(text: str, label: str, threshold: float, pieces, scores): |
| spans = [] |
| active = None |
| for (word, start, end), score in zip(pieces, scores): |
| keep = score >= threshold |
| if label in {"PHONE_NUMBER", "BANK_ROUTING_NUMBER", "CREDIT_DEBIT_CARD"} and word in {"-", "/"}: |
| keep = active is not None and score >= threshold / 2.0 |
| if keep: |
| if active is None: |
| active = {"start": start, "end": end, "label": label} |
| else: |
| if start - active["end"] <= 1: |
| active["end"] = end |
| else: |
| spans.append(active) |
| active = {"start": start, "end": end, "label": label} |
| elif active is not None: |
| spans.append(active) |
| active = None |
| if active is not None: |
| spans.append(active) |
| out = [] |
| for span in spans: |
| value = text[span["start"] : span["end"]] |
| if plausible_label_text(label, value): |
| out.append( |
| { |
| "label": label, |
| "start": span["start"], |
| "end": span["end"], |
| "text": value, |
| } |
| ) |
| return out |
|
|
|
|
| def word_aligned_label_spans( |
| text: str, |
| model, |
| tokenizer, |
| label: str, |
| threshold: float, |
| ): |
| pieces, scores = word_scores_for_label(text, model, tokenizer, label) |
| return _word_aligned_label_spans_from_scores(text, label, threshold, pieces, scores) |
|
|
|
|
| def word_aligned_label_spans_onnx( |
| text: str, |
| session, |
| tokenizer, |
| config, |
| label: str, |
| threshold: float, |
| ): |
| pieces, scores = word_scores_for_label_onnx(text, session, tokenizer, config, label) |
| return _word_aligned_label_spans_from_scores(text, label, threshold, pieces, scores) |
|
|
|
|
| def pipeline_to_spans(text: str, outputs: list[dict], min_score: float): |
| spans = [] |
| for output in outputs: |
| label = normalize_label(output.get("entity_group") or output.get("entity") or "") |
| if not label: |
| continue |
| score = float(output.get("score", 0.0)) |
| if score < min_score: |
| continue |
| spans.append( |
| { |
| "label": label, |
| "start": int(output["start"]), |
| "end": int(output["end"]), |
| "score": score, |
| "text": text[int(output["start"]) : int(output["end"])], |
| } |
| ) |
| return spans |
|
|
|
|
| def overlaps(a: dict, b: dict) -> bool: |
| return not (a["end"] <= b["start"] or b["end"] <= a["start"]) |
|
|
|
|
| def span_length(span: dict) -> int: |
| return int(span["end"]) - int(span["start"]) |
|
|
|
|
| def normalize_simple_span(span: dict): |
| label = normalize_label(span["label"]) |
| value = span["text"] |
| if label == "PHONE_NUMBER" and plausible_label_text("CREDIT_DEBIT_CARD", value): |
| label = "CREDIT_DEBIT_CARD" |
| if label in FORMAT_LABELS or label == "POSTCODE": |
| if not plausible_label_text(label, value): |
| return None |
| return { |
| "label": label, |
| "start": int(span["start"]), |
| "end": int(span["end"]), |
| "score": float(span.get("score", 0.0)), |
| "text": value, |
| } |
|
|
|
|
| def dedupe_and_sort(spans: list[dict]): |
| ordered = sorted( |
| spans, |
| key=lambda span: ( |
| int(span["start"]), |
| -span_length(span), |
| OUTPUT_PRIORITY.get(str(span["label"]).upper(), 99), |
| ), |
| ) |
| kept = [] |
| for span in ordered: |
| if any(overlaps(span, other) for other in kept): |
| continue |
| kept.append(span) |
| return kept |
|
|
|
|
| def repair_irish_core_spans( |
| text: str, |
| model, |
| tokenizer, |
| general_outputs: list[dict], |
| other_min_score: float, |
| ppsn_min_score: float, |
| label_thresholds: dict[str, float] | None = None, |
| ): |
| thresholds = dict(DEFAULT_LABEL_THRESHOLDS) |
| if label_thresholds: |
| thresholds.update({key.upper(): value for key, value in label_thresholds.items()}) |
|
|
| spans = [] |
| for span in pipeline_to_spans(text, general_outputs, min_score=other_min_score): |
| normalized = normalize_simple_span(span) |
| if normalized is not None and normalized["label"] != "PPSN": |
| spans.append(normalized) |
|
|
| ppsn_spans = word_aligned_ppsn_spans(text, model, tokenizer, threshold=ppsn_min_score) |
| for span in ppsn_spans: |
| spans.append( |
| { |
| "label": "PPSN", |
| "start": int(span["start"]), |
| "end": int(span["end"]), |
| "score": float(span.get("score", 0.0)), |
| "text": text[int(span["start"]) : int(span["end"])], |
| } |
| ) |
|
|
| repairs = [] |
| for label, threshold in thresholds.items(): |
| repairs.extend(word_aligned_label_spans(text, model, tokenizer, label, threshold)) |
|
|
| for candidate in repairs: |
| updated = [] |
| replaced = False |
| for span in spans: |
| if not overlaps(candidate, span): |
| updated.append(span) |
| continue |
| if candidate["label"] == span["label"] and span_length(candidate) > span_length(span): |
| replaced = True |
| continue |
| if candidate["label"] in FORMAT_LABELS and span["label"] in FORMAT_LABELS and span_length(candidate) > span_length(span): |
| replaced = True |
| continue |
| updated.append(span) |
| spans = updated |
| if replaced or not any(overlaps(candidate, span) for span in spans): |
| spans.append(candidate) |
|
|
| return dedupe_and_sort(spans) |
|
|
|
|
| def repair_irish_core_spans_onnx( |
| text: str, |
| session, |
| tokenizer, |
| config, |
| other_min_score: float, |
| ppsn_min_score: float, |
| label_thresholds: dict[str, float] | None = None, |
| general_outputs: list[dict] | None = None, |
| ): |
| from onnx_token_classifier import simple_aggregate_spans_onnx, word_aligned_ppsn_spans_onnx |
|
|
| thresholds = dict(DEFAULT_LABEL_THRESHOLDS) |
| if label_thresholds: |
| thresholds.update({key.upper(): value for key, value in label_thresholds.items()}) |
|
|
| spans = [] |
| if general_outputs is None: |
| general_outputs = simple_aggregate_spans_onnx( |
| text, |
| session, |
| tokenizer, |
| config, |
| min_score=other_min_score, |
| ) |
| for span in pipeline_to_spans(text, general_outputs, min_score=other_min_score): |
| normalized = normalize_simple_span(span) |
| if normalized is not None and normalized["label"] != "PPSN": |
| spans.append(normalized) |
|
|
| ppsn_spans = word_aligned_ppsn_spans_onnx(text, session, tokenizer, config, threshold=ppsn_min_score) |
| for span in ppsn_spans: |
| spans.append( |
| { |
| "label": "PPSN", |
| "start": int(span["start"]), |
| "end": int(span["end"]), |
| "score": float(span.get("score", 0.0)), |
| "text": text[int(span["start"]) : int(span["end"])], |
| } |
| ) |
|
|
| repairs = [] |
| for label, threshold in thresholds.items(): |
| repairs.extend(word_aligned_label_spans_onnx(text, session, tokenizer, config, label, threshold)) |
|
|
| for candidate in repairs: |
| updated = [] |
| replaced = False |
| for span in spans: |
| if not overlaps(candidate, span): |
| updated.append(span) |
| continue |
| if candidate["label"] == span["label"] and span_length(candidate) > span_length(span): |
| replaced = True |
| continue |
| if candidate["label"] in FORMAT_LABELS and span["label"] in FORMAT_LABELS and span_length(candidate) > span_length(span): |
| replaced = True |
| continue |
| updated.append(span) |
| spans = updated |
| if replaced or not any(overlaps(candidate, span) for span in spans): |
| spans.append(candidate) |
|
|
| return dedupe_and_sort(spans) |
|
|
