| """NERGAL hybrid PII cleaner: frozen regex ∪ windowed XLM-R BIO head. |
| |
| This file is the public PII island. It does not import the lab training stack and |
| must not call embedding-extension. The packed tokenizer already has the gap ids. |
| """ |
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| import math |
| import re |
| from dataclasses import dataclass |
| from functools import lru_cache |
| from pathlib import Path |
|
|
| import scrub_pii |
| from scrub_pii import LEGACY_PHONE_TAG, PERSON_TAG, PHONE_TAG, PII_TAG |
|
|
| HUB_ID = 'SlayerLab/NERGAL' |
| VERSION = '2.0.9' |
| GAPS = ['[PII_SPACE]', '[PII_BREAK]'] |
| GAP_IDS = [250002, 250003] |
| BIO_LABELS = ['O', 'B-phone', 'I-phone', 'B-pii', 'I-pii'] |
| LABELS = ['phone', 'pii'] |
| THRESHOLD = 0.95 |
| RULES_SHA = '08bcaf1909653a8632df41cc517798a39053a8d89ae34178cb4a1f99e3ce96eb' |
| DTYPES = ('float32', 'float16') |
| MARKERS = (PHONE_TAG, PII_TAG, PERSON_TAG, LEGACY_PHONE_TAG) |
| TAGS = {'person': PERSON_TAG, 'pii': PII_TAG, 'phone': PHONE_TAG} |
| RANK = {label: i for i, label in enumerate(TAGS)} |
| PERSON_LABELS = ('PERSON', 'PERSON_F', 'PERSON_L') |
| _WORD = re.compile(r"[\w'’-]") |
| _PLACEHOLDER = re.compile('|'.join(map(re.escape, MARKERS))) |
| |
| |
| |
| |
| _PARTICIPLE = '(?:y|ego|emu|ym|ych|ymi|a|ej|ą|e)?' |
| _CA = '(?:a|y|ę|ą|o|ów|om|ami|ach)' |
| _NIK = '(?:k|ka|kowi|kiem|ku|cy|ków|kom|kami|kach|czk(?:a|i|ę|ą|o|om|ami|ach)|czce|czek)' |
| _NT = '(?:t|ta|towi|tem|cie|ci|tów|tom|tami|tach|tk(?:a|i|ę|ą|o|om|ami|ach)|tce|tek)' |
| ROLE_NOUNS = tuple(stem + _PARTICIPLE for stem in ('odwołując', 'przystępując', 'zamawiając', 'protestując', 'skarżąc', |
| 'przewodnicząc')) + ( |
| 'pozwan(?:y|ego|emu|ym|i|ych|ymi|a|ej|ą|e)', 'wykonawc' + _CA, 'wnioskodawc' + _CA, |
| 'wnioskodawczy(?:ni|ń|niom|niami|niach)', 'pow(?:ód|od(?:a|owi|em|zie|owie|ów|om|ami|ach))', |
| 'powódk(?:a|i|ce|ę|ą|o|ek|om|ami|ach)', 'inwestor(?:a|owi|em|ze|zy|ów|om|ami|ach)?', |
| 'prezes(?:a|owi|em|ie|i|ów|om|ami|ach)?', 'uczestni' + _NIK, 'pełnomocni' + _NIK, 'oferen' + _NT, 'protokolan' + _NT) |
| _ROLE_NOUN = re.compile(r'\W*(?:%s)_?\d*\W*' % '|'.join(ROLE_NOUNS), re.IGNORECASE) |
| _NUMBER = re.compile(r'\W*\d+\W*') |
| |
| |
| |
| |
| |
| |
| |
| _ROLE_CUT = re.compile(r"\W*(?:(?:odwołując|przystępując|zamawiając|protestując|skarżąc|przewodnicząc)(?:eg|em|yc)|" |
| r"(?:wykonawc|wnioskodawc)(?:ó|am|ac)?)\W*", re.IGNORECASE) |
| _ROLE_BARE = re.compile('(?:%s)' % '|'.join(ROLE_NOUNS), re.IGNORECASE) |
| _GLUED_AFTER = re.compile(r"([a-ząćęłńóśźż]{4,})\W*") |
| _LEAD = re.compile(r'\W*') |
|
|
|
|
| def sha(path): |
| with Path(path).open('rb') as stream: |
| return hashlib.file_digest(stream, 'sha256').hexdigest() |
|
|
|
|
| def verify_rules(path=None): |
| digest = sha(path or scrub_pii.__file__) |
| if digest != RULES_SHA: |
| raise ValueError(f'Unexpected rules sha256 {digest}') |
| return digest |
|
|
|
|
| @dataclass(frozen=True) |
| class Unit: |
| model: str |
| start: int |
| end: int |
| gap: bool |
|
|
|
|
| def unitize(text, encode=None, unk=None): |
| units = [] |
| for match in re.finditer(r'\s+|\S', text): |
| raw = match.group() |
| gap = raw.isspace() |
| model = GAPS[int(any(c in raw for c in '\r\n\v\f\x85\u2028\u2029'))] if gap else raw |
| if not gap and encode is not None and not encode(model): |
| if not unk: |
| raise ValueError('Zero-piece unit without unknown token') |
| model = unk |
| units.append(Unit(model, match.start(), match.end(), gap)) |
| return units |
|
|
|
|
| def windows(units, count, *, max_units=384, limit=512): |
| overlap = 128 |
| result, start = [], 0 |
| while start < len(units): |
| lo, hi = start + 1, min(start + max_units, len(units)) |
| while lo < hi: |
| mid = (lo + hi + 1) // 2 |
| if count(units[start:mid]) <= limit: |
| lo = mid |
| else: |
| hi = mid - 1 |
| end, size = lo, count(units[start:lo]) |
| if size > limit or (end < len(units) and end - start <= overlap): |
| raise ValueError('Token budget cannot fit a progressing window') |
| width = 64 |
| result.append({ |
| 'start': start, 'end': end, 'tokens': size, |
| 'owner_start': start if not result else start + width - 1, |
| 'owner_end': end if end == len(units) else end - width + 1, |
| }) |
| if end == len(units): |
| break |
| start = end - overlap |
| return result |
|
|
|
|
| def raw_span(units, a, b, label, score): |
| if not 0 <= a < b <= len(units) or label not in LABELS or not math.isfinite(score) or not 0 <= score <= 1: |
| raise ValueError('Invalid unit prediction') |
| while a < b and units[a].gap: |
| a += 1 |
| while a < b and units[b - 1].gap: |
| b -= 1 |
| if a == b: |
| return None |
| return {'start': units[a].start, 'end': units[b - 1].end, 'label': label, 'score': score} |
|
|
|
|
| def decode_bio(units, logits): |
| if len(units) != len(logits) or any(len(v) != 5 or any(not math.isfinite(x) for x in v) for v in logits): |
| raise ValueError('Invalid BIO logits') |
| result, active, probabilities = [], None, [] |
|
|
| def finish(end): |
| nonlocal active |
| if active is None: |
| return |
| start, label = active |
| span = raw_span(units, start, end, label, min(probabilities[start:end])) |
| if span is not None: |
| result.append(span) |
| active = None |
|
|
| for i, values in enumerate(logits): |
| tag = max(range(5), key=lambda j: values[j]) |
| exponentials = [math.exp(x - max(values)) for x in values] |
| probabilities.append(exponentials[tag] / sum(exponentials)) |
| label = LABELS[(tag - 1) // 2] if tag else None |
| if tag == 0 or tag in (1, 3) or active is None or active[1] != label: |
| finish(i) |
| active = (i, label) if tag else None |
| finish(len(units)) |
| return result |
|
|
|
|
| def decode(spans, threshold=THRESHOLD): |
| if any(not math.isfinite(s['score']) or not 0 <= s['score'] <= 1 for s in spans): |
| raise ValueError('Nonfinite/invalid confidence') |
| result = [] |
| for span in sorted(spans, key=lambda s: (-s['score'], -(s['end'] - s['start']), s['start'], s['label'])): |
| if span['score'] >= threshold and not any(span['start'] < p['end'] and p['start'] < span['end'] for p in result): |
| result.append(span) |
| return sorted(result, key=lambda s: (s['start'], s['end'], s['label'])) |
|
|
|
|
| class Encoding: |
| def __init__(self, tokenizer): |
| self.tokenizer = tokenizer |
| tokenizer.model_max_length = 512 |
| self.pieces = lru_cache(maxsize=16384)(lambda s: tuple(tokenizer.encode(s, add_special_tokens=False))) |
|
|
| def encode(self, words): |
| encoded = self.tokenizer([words], is_split_into_words=True, truncation=False, padding=False) |
| ids = encoded['input_ids'][0] |
| mapping = encoded.word_ids(0) |
| first, actual = {}, {} |
| for i, word in enumerate(mapping): |
| if word is not None: |
| first.setdefault(word, i) |
| actual.setdefault(word, []).append(ids[i]) |
| if set(first) != set(range(len(words))): |
| raise ValueError('Tokenizer dropped a unit') |
| if any(tuple(actual[j]) != self.pieces(word) for j, word in enumerate(words)): |
| raise ValueError('Unit token IDs change with window context') |
| return encoded, [first[j] for j in range(len(words))] |
|
|
| def count(self, units): |
| """Special tokens + cached unit pieces. Exact because encode() rejects any unit whose pieces change with |
| context; no re-tokenization inside the windows() binary search.""" |
| return self.tokenizer.num_special_tokens_to_add() + sum(len(self.pieces(u.model)) for u in units) |
|
|
| def prepare(self, text): |
| units = unitize(text, self.pieces, self.tokenizer.unk_token) |
| return units, windows(units, self.count) if units else [] |
|
|
|
|
| def spans_from(module, text): |
| """Rule spans on the original text; an unlabelled PESEL carries `bare`. Existing [PHONE]/[PII]/[Telefon] |
| placeholders do not switch the rules off.""" |
| result = [] |
| module.scrub_pii(text, spans=result) |
| return sorted(({k: s[k] for k in ('start', 'end', 'label', 'bare') if k in s} | {'score': 1.0} for s in result), |
| key=lambda s: s['start']) |
|
|
|
|
| def rules(text): |
| verify_rules() |
| return spans_from(scrub_pii, text) |
|
|
|
|
| def expand(text, start, end): |
| """Grow a token range to the whole word: letters, digits, apostrophes and hyphens.""" |
| while start > 0 and _WORD.match(text, start - 1): |
| start -= 1 |
| while end < len(text) and _WORD.match(text, end): |
| end += 1 |
| return start, end |
|
|
|
|
| def whole_digits(text, start, end): |
| """Grow a span that starts or ends inside a run of digits over the whole run (2.0.6). No gold span cuts a digit |
| run, so a model span stopping mid-number leaves the rest of an identifier open.""" |
| if text[start].isdecimal(): |
| while start > 0 and text[start - 1].isdecimal(): |
| start -= 1 |
| if text[end - 1].isdecimal(): |
| while end < len(text) and text[end].isdecimal(): |
| end += 1 |
| return start, end |
|
|
|
|
| def damaged_role_word(text, start, end): |
| """A role noun cut short (_ROLE_CUT) or glued to a lowercase word the text also uses on its own (_GLUED_AFTER).""" |
| if _ROLE_CUT.fullmatch(text, start, end): |
| return True |
| word = _LEAD.match(text, start, end).end() |
| for cut in range(word + 1, end): |
| tail = _GLUED_AFTER.fullmatch(text, cut, end) |
| if tail and _ROLE_BARE.fullmatch(text, word, cut) and re.search( |
| r"(?<![\w'’-])%s(?![\w'’-])" % re.escape(tail[1]), text): |
| return True |
| return False |
|
|
|
|
| def person_spans(text, ranges): |
| """Person spans from FastPDN person-token ranges: whole words, joined across spaces, tabs and no-break spaces |
| (not punctuation or line breaks); words touching a placeholder, role nouns (ROLE_NOUNS), damaged ones |
| (damaged_role_word) and a number right after a skipped role noun (`Odwołujący 1`) are skipped.""" |
| blocked = [m.span() for m in _PLACEHOLDER.finditer(text)] |
| merged, role_end = [], None |
| for start, end in sorted(expand(text, a, b) for a, b in ranges): |
| if any(start < b and a < end for a, b in blocked): |
| continue |
| if _ROLE_NOUN.fullmatch(text, start, end): |
| role_end = end |
| continue |
| if damaged_role_word(text, start, end): |
| continue |
| if role_end is not None and _NUMBER.fullmatch(text, start, end) and not text[role_end:start].strip(' \t\u00a0'): |
| continue |
| if merged and not text[merged[-1][1]:start].strip(' \t\u00a0'): |
| merged[-1][1] = max(merged[-1][1], end) |
| else: |
| merged.append([start, end]) |
| return [{'start': a, 'end': b, 'label': 'person', 'score': 1.0} for a, b in merged] |
|
|
|
|
| def char_labels(text, spans, firm=()): |
| """Each character's highest-ranked label (phone > pii > person), or None; under a `firm` span only firm labels |
| compete.""" |
| labels, best = [None] * len(text), [None] * len(text) |
| for tier, group in ((0, spans), (1, firm)): |
| for span in group: |
| start, end, label = span['start'], span['end'], span['label'] |
| if label not in RANK: |
| raise ValueError(f'Unknown label {label!r}') |
| if not 0 <= start < end <= len(text): |
| raise ValueError('Span outside text') |
| key = (tier, RANK[label]) |
| for i in range(start, end): |
| if best[i] is None or key > best[i]: |
| labels[i], best[i] = label, key |
| return labels |
|
|
|
|
| def apply_union(text, spans, firm=()): |
| """Each run of one label (char_labels) becomes one placeholder.""" |
| labels = char_labels(text, spans, firm) |
| out, chars, counts, i = [], 0, dict.fromkeys(TAGS, 0), 0 |
| while i < len(text): |
| lab = labels[i] |
| if lab is None: |
| out.append(text[i]) |
| i += 1 |
| continue |
| j = i + 1 |
| while j < len(text) and labels[j] == lab: |
| j += 1 |
| out.append(TAGS[lab]) |
| chars += len(TAGS[lab]) |
| counts[lab] += 1 |
| i = j |
| return ''.join(out), chars, counts |
|
|
|
|
| def model_keep(text, model_spans, *, threshold=THRESHOLD, module=scrub_pii): |
| """Decoded model spans grown over whole runs of digits (whole_digits), then each phone span cut by the rules' |
| phone policy: one span per number of 7+ digits; short and emergency numbers alone are dropped.""" |
| keep = [] |
| for s in decode(model_spans, threshold): |
| start, end = whole_digits(text, s['start'], s['end']) |
| parts = module._phone_spans(text, start, end) if s['label'] == 'phone' else [(start, end)] |
| keep += [dict(s, start=a, end=b) for a, b in parts] |
| return keep |
|
|
|
|
| def scrub_spans(text, rule_spans, model_spans, *, threshold=THRESHOLD, module=scrub_pii, name_spans=()): |
| rule_keys = {(s['start'], s['end'], s['label']) for s in rule_spans} |
| keep = model_keep(text, model_spans, threshold=threshold, module=module) |
| extra = sum(1 for s in keep if (s['start'], s['end'], s['label']) not in rule_keys) |
| _, rules_chars, _ = apply_union(text, rule_spans) |
| |
| |
| firm = [s for s in rule_spans if not s.get('bare')] |
| masked, union_chars, counts = apply_union(text, [*rule_spans, *keep, *name_spans], firm) |
| return masked, {**counts, 'rules_placeholder_chars': rules_chars, 'union_placeholder_chars': union_chars, |
| 'model_extra_spans': extra} |
|
|
|
|
| def _resolve(source, *, local_files_only, names=False): |
| path = Path(source) |
| if path.is_dir(): |
| return path |
| from huggingface_hub import snapshot_download |
| return Path(snapshot_download(source, local_files_only=local_files_only, |
| ignore_patterns=None if names else ['names/*'])) |
|
|
|
|
| def _load_rules_module(asset): |
| path = Path(asset) / 'scrub_pii.py' |
| if path.is_file(): |
| import importlib.util |
| spec = importlib.util.spec_from_file_location('_nergal_scrub_pii', path) |
| module = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(module) |
| verify_rules(path) |
| return module |
| verify_rules() |
| return scrub_pii |
|
|
|
|
| _FEEDS = (('input_ids', 'ids'), ('attention_mask', 'attention_mask'), ('token_type_ids', 'type_ids')) |
|
|
|
|
| def pack(lengths, batch_tokens, max_batch): |
| """Indices sorted by length, cut into batches of at most max_batch rows and batch_tokens padded tokens.""" |
| order = sorted(range(len(lengths)), key=lengths.__getitem__) |
| out, i = [], 0 |
| while i < len(order): |
| j = i + 1 |
| while j < len(order) and j - i < max_batch and (j - i + 1) * lengths[order[j]] <= batch_tokens: |
| j += 1 |
| out.append(order[i:j]) |
| i = j |
| return out |
|
|
|
|
| class Names: |
| """FastPDN-PII person names (names/, CC-BY-4.0, see names/NOTICE.md): the FP32 ONNX converted to safetensors, run |
| in torch on NERGAL's device; argmax per token, windows packed like predict_many.""" |
|
|
| def __init__(self, asset, card, device='cpu'): |
| import torch |
| from safetensors.torch import load_file |
| from tokenizers import Tokenizer |
| from transformers import BertConfig, BertForTokenClassification |
| root = Path(asset) / 'names' |
| for name, digest in card['sha256'].items(): |
| if sha(root / name) != digest: |
| raise ValueError(f'Unexpected names/{name} sha256') |
| config = BertConfig.from_pretrained(str(root)) |
| labels = config.id2label |
| self.person_ids = {i for i, label in labels.items() if label[2:] in PERSON_LABELS} |
| if sorted(labels[i] for i in self.person_ids) != sorted(f'{p}-{l}' for p in 'BI' for l in PERSON_LABELS): |
| raise ValueError('names/config.json does not label B/I person tags') |
| self.tokenizer = Tokenizer.from_file(str(root / 'tokenizer.json')) |
| self.tokenizer.enable_truncation(max_length=card['max_length'], stride=card['stride']) |
| self.tokenizer.no_padding() |
| self.pad = {'input_ids': config.pad_token_id, 'attention_mask': 0, 'token_type_ids': 0} |
| self.model = BertForTokenClassification(config) |
| self.model.load_state_dict(load_file(root / 'model.safetensors'), strict=True) |
| self.model.to(device, dtype=getattr(torch, card['dtype'])).eval() |
| self.device, self._torch = device, torch |
|
|
| def spans(self, text): |
| return self.spans_many([text])[0] |
|
|
| def spans_many(self, texts, *, batch_tokens=32768, max_batch=128): |
| torch, texts = self._torch, list(texts) |
| items = [(d, w) for d, text in enumerate(texts) |
| for enc in [self.tokenizer.encode(text)] for w in [enc, *enc.overflowing]] |
| ranges = [[] for _ in texts] |
| with torch.inference_mode(): |
| for part in pack([len(w.ids) for _, w in items], batch_tokens, max_batch): |
| rows = [items[i][1] for i in part] |
| width = max(len(w.ids) for w in rows) |
| feed = {name: torch.tensor([getattr(w, attr) + [self.pad[name]] * (width - len(w.ids)) for w in rows], |
| device=self.device) for name, attr in _FEEDS} |
| logits = self.model(**feed).logits.float().cpu() |
| for row, i in enumerate(part): |
| d, w = items[i] |
| tags = logits[row, :len(w.ids)].argmax(-1).tolist() |
| ranges[d] += [(a, b) for tag, special, (a, b) in zip(tags, w.special_tokens_mask, w.offsets, strict=True) |
| if tag in self.person_ids and not special and a < b] |
| return [person_spans(text, r) for text, r in zip(texts, ranges, strict=True)] |
|
|
|
|
| class Nergal: |
| def __init__(self, asset, device='cpu', dtype='float32', names=False): |
| import torch |
| from transformers import AutoModelForTokenClassification, AutoTokenizer |
| if dtype not in DTYPES: |
| raise ValueError(f'dtype must be one of {DTYPES}') |
| if dtype == 'float16' and torch.device(device).type == 'cpu': |
| raise ValueError('float16 needs a cuda or mps device') |
| self.dtype = dtype |
| self.device = device |
| self._torch = torch |
| asset = Path(asset) |
| self._scrub = _load_rules_module(asset) |
| card = json.loads((asset / 'hybrid.json').read_text()) |
| if card['gap_ids'] != GAP_IDS or card['threshold'] != THRESHOLD: |
| raise ValueError('hybrid.json does not match this NERGAL snapshot') |
| if names and 'names' not in card: |
| raise ValueError('This NERGAL snapshot has no names model') |
| tokenizer = AutoTokenizer.from_pretrained( |
| str(asset), local_files_only=True, use_fast=True, fix_mistral_regex=False, |
| ) |
| if [tokenizer.convert_tokens_to_ids(t) for t in GAPS] != GAP_IDS: |
| raise ValueError('Packed NERGAL tokenizer is missing gap ids') |
| self.model = AutoModelForTokenClassification.from_pretrained(str(asset), local_files_only=True) |
| self.encoding = Encoding(tokenizer) |
| self.threshold = THRESHOLD |
| self.model.to(device, dtype=getattr(torch, dtype)).eval() |
| self.names = Names(asset, card['names'], device) if names else None |
|
|
| @classmethod |
| def from_pretrained(cls, source=HUB_ID, *, device=None, dtype='float32', local_files_only=False, names=False): |
| import torch |
| if device is None: |
| device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu' |
| return cls(_resolve(source, local_files_only=local_files_only, names=names), device=device, dtype=dtype, |
| names=names) |
|
|
| def predict(self, text): |
| return self.predict_many([text])[0] |
|
|
| def predict_many(self, texts, *, batch_tokens=32768, max_batch=128, group=64): |
| """Model spans for each text. Windows of up to `group` texts are sorted by token length and packed into |
| batches of at most `max_batch` rows and `batch_tokens` padded tokens; per text, units, windows, averaging |
| and decode are the same as one window at a time.""" |
| if min(batch_tokens, max_batch, group) < 1: |
| raise ValueError('batch_tokens, max_batch and group must be positive') |
| texts = list(texts) |
| result = [] |
| for i in range(0, len(texts), group): |
| result += self._predict_group(texts[i:i + group], batch_tokens, max_batch) |
| return result |
|
|
| def _predict_group(self, texts, batch_tokens, max_batch): |
| torch, encoding = self._torch, self.encoding |
| docs, items = [], [] |
| for d, text in enumerate(texts): |
| units, chunks = encoding.prepare(text) |
| docs.append((units, torch.zeros(len(units), 5), torch.zeros(len(units), 1))) |
| for window in chunks: |
| encoded, first = encoding.encode([u.model for u in units[window['start']:window['end']]]) |
| items.append((len(encoded['input_ids'][0]), d, window, encoded, first)) |
| with torch.inference_mode(): |
| for part in pack([n for n, *_ in items], batch_tokens, max_batch): |
| part = [items[i] for i in part] |
| batch = encoding.tokenizer.pad([{k: v[0] for k, v in e.items()} for _, _, _, e, _ in part], |
| padding=True, return_tensors='pt') |
| if batch['input_ids'].shape[1] > 512: |
| raise ValueError('Batch exceeds encoder limit') |
| logits = self.model(**{k: v.to(self.device) for k, v in batch.items()}).logits.float().cpu() |
| for row, (_, d, window, _, first) in enumerate(part): |
| a, b = window['start'], window['end'] |
| docs[d][1][a:b] += logits[row, first] |
| docs[d][2][a:b] += 1 |
| result = [] |
| for units, sums, counts in docs: |
| if (counts == 0).any(): |
| raise ValueError('Missing inference units') |
| result.append(decode_bio(units, (sums / counts).tolist()) if units else []) |
| return result |
|
|
| def rule_spans(self, text): |
| return spans_from(self._scrub, text) |
|
|
| def scrub(self, text): |
| return self.scrub_many([text])[0] |
|
|
| def scrub_many(self, texts): |
| texts = list(texts) |
| name_spans = self.names.spans_many(texts) if self.names else [()] * len(texts) |
| return [scrub_spans(text, self.rule_spans(text), model, threshold=self.threshold, module=self._scrub, |
| name_spans=spans) |
| for text, model, spans in zip(texts, self.predict_many(texts), name_spans, strict=True)] |
|
|
|
|
| def main(argv=None): |
| import argparse |
| import sys |
| parser = argparse.ArgumentParser(description='NERGAL hybrid PII cleaner') |
| parser.add_argument('--repo', default=HUB_ID) |
| parser.add_argument('--device', default=None) |
| parser.add_argument('--dtype', default='float32', choices=DTYPES) |
| parser.add_argument('--local', action='store_true') |
| parser.add_argument('--names', action='store_true') |
| args = parser.parse_args(argv) |
| nergal = Nergal.from_pretrained(args.repo, device=args.device, dtype=args.dtype, |
| local_files_only=args.local, names=args.names) |
| text = sys.stdin.read() |
| masked, counts = nergal.scrub(text) |
| sys.stdout.write(masked) |
| print(json.dumps(counts), file=sys.stderr) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
