| |
| import json, random, argparse, ast |
| import numpy as np, torch |
| from datasets import load_dataset, Dataset |
| from transformers import AutoTokenizer, AutoModelForTokenClassification, TrainingArguments, Trainer, DataCollatorForTokenClassification |
| import evaluate |
|
|
| CATEGORIES = ["account_number","private_address","private_date","private_email","private_person","private_phone","private_url","secret","fax_number","credit_card_last4","company_contact_block"] |
| LABELS = ["O"] |
| for cat in CATEGORIES: |
| for p in ("B","I","E","S"): LABELS.append(f"{p}-{cat}") |
| label2id = {l:i for i,l in enumerate(LABELS)} |
| id2label = {i:l for l,i in label2id.items()} |
| NUM_LABELS = len(LABELS) |
| seqeval = evaluate.load("seqeval") |
|
|
| def compute_metrics(p): |
| predictions, labels = p |
| predictions = np.argmax(predictions, axis=2) |
| true_preds = [[id2label[pred] for pred,lab in zip(pred_row,lab_row) if lab!=-100] for pred_row,lab_row in zip(predictions,labels)] |
| true_labs = [[id2label[lab] for pred,lab in zip(pred_row,lab_row) if lab!=-100] for pred_row,lab_row in zip(predictions,labels)] |
| results = seqeval.compute(predictions=true_preds, references=true_labs) |
| return {"precision":results["overall_precision"],"recall":results["overall_recall"],"f1":results["overall_f1"],"accuracy":results["overall_accuracy"]} |
|
|
| from faker import Faker |
| fake = Faker() |
|
|
| def generate_synthetic_examples(n=5000, seed=42): |
| random.seed(seed); fake.seed_instance(seed) |
| examples = [] |
| def add(text, spans): examples.append({"text":text,"spans":spans}) |
| for _ in range(n): |
| r = random.random() |
| if r < 0.25: |
| fax = fake.numerify(text="(###) ###-####") |
| tmpl = random.choice([f"Please fax documents to {fax}.",f"Fax: {fax}\nAttn: Legal",f"Secure fax line: {fax}",f"You can reach us at phone (555) 123-4567 or fax {fax}.",f"Facsimile: {fax}"]) |
| s = tmpl.find(fax); add(tmpl, [(s,s+len(fax),"fax_number")]) |
| elif r < 0.5: |
| last4 = fake.numerify(text="####") |
| tmpl = random.choice([f"Card ending in {last4} charged.",f"Visa ****-****-****-{last4}",f"Last 4 digits: {last4}",f"Card on file ...{last4}",f"XXXX-XXXX-XXXX-{last4}"]) |
| s = tmpl.find(last4); add(tmpl, [(s,s+len(last4),"credit_card_last4")]) |
| elif r < 0.75: |
| company = fake.company(); addr = fake.street_address() + ", " + fake.city() + ", " + fake.state_abbr() + " " + fake.zipcode() |
| phone = fake.numerify(text="(###) ###-####"); email = fake.company_email() |
| tmpl = random.choice([f"{company}\n{addr}\nPhone: {phone}\nEmail: {email}",f"Contact:\n{company}\n{addr}\nTel: {phone}\n{email}",f"{company} HQ\n{addr}\nMain: {phone}\nInquiries: {email}"]) |
| s = tmpl.find(company); e = tmpl.find(email)+len(email); add(tmpl, [(s,e,"company_contact_block")]) |
| else: |
| person = fake.name(); email = fake.email(); fax = fake.numerify(text="(###) ###-####") |
| phone = fake.numerify(text="(###) ###-####"); last4 = fake.numerify(text="####") |
| company = fake.company(); addr = fake.street_address() + ", " + fake.city() + ", " + fake.state_abbr() + " " + fake.zipcode() |
| tmpl = random.choice([f"From: {person} <{email}>\nTo: Legal\nFax: {fax}\nPhone: {phone}\nCard: {last4}\n{company}\n{addr}",f"Client: {person}\nEmail: {email}\nFax: {fax}\nTel: {phone}\nPayment: ****{last4}\nEmployer: {company}\n{addr}"]) |
| spans = [] |
| for sub,lab in [(person,"private_person"),(email,"private_email"),(fax,"fax_number"),(phone,"private_phone"),(last4,"credit_card_last4"),(company,"company_contact_block"),(addr,"private_address")]: |
| idx = tmpl.find(sub) |
| if idx >= 0: spans.append((idx,idx+len(sub),lab)) |
| add(tmpl, spans) |
| return examples |
|
|
| NEMOTRON_MAP = {"first_name":"private_person","last_name":"private_person","full_name":"private_person","name":"private_person","email":"private_email","phone_number":"private_phone","street_address":"private_address","address":"private_address","date_of_birth":"private_date","date":"private_date","credit_card_number":"account_number","ssn":"account_number","company_name":"company_contact_block","url":"private_url","secret":"secret","api_key":"secret","password":"secret","token":"secret"} |
|
|
| def load_nemotron_split(split, max_examples=10000): |
| ds = load_dataset("nvidia/Nemotron-PII", split=split) |
| examples = [] |
| for ex in ds: |
| if len(examples) >= max_examples: break |
| text = ex["text"] |
| spans_raw = ex["spans"] |
| if isinstance(spans_raw, str): |
| try: |
| spans_raw = json.loads(spans_raw) |
| except json.JSONDecodeError: |
| spans_raw = ast.literal_eval(spans_raw) |
| spans = [] |
| for sp in spans_raw: |
| lab = NEMOTRON_MAP.get(sp["label"]) |
| if lab: spans.append((sp["start"],sp["end"],lab)) |
| if spans: examples.append({"text":text,"spans":spans}) |
| return examples |
|
|
| def tokenize_and_align(examples, tokenizer): |
| tokenized = tokenizer([ex["text"] for ex in examples], truncation=True, max_length=512, return_offsets_mapping=True) |
| all_labels = [] |
| for i,ex in enumerate(examples): |
| offsets = tokenized["offset_mapping"][i]; labels = ["O"]*len(offsets) |
| for start,end,lab in ex["spans"]: |
| covered = [j for j,(ts,te) in enumerate(offsets) if ts is not None and te is not None and ts < end and te > start] |
| if not covered: continue |
| if len(covered)==1: labels[covered[0]] = f"S-{lab}" |
| else: |
| labels[covered[0]] = f"B-{lab}" |
| for idx in covered[1:-1]: labels[idx] = f"I-{lab}" |
| labels[covered[-1]] = f"E-{lab}" |
| label_ids = [] |
| for j,(ts,te) in enumerate(offsets): |
| if ts is None and te is None: label_ids.append(-100) |
| else: label_ids.append(label2id.get(labels[j],0)) |
| all_labels.append(label_ids) |
| tokenized["labels"] = all_labels |
| tokenized.pop("offset_mapping") |
| return tokenized |
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--model", default="openai/privacy-filter") |
| parser.add_argument("--output_model", default="narcolepticchicken/privacy-filter-enhanced") |
| parser.add_argument("--epochs", type=int, default=3) |
| parser.add_argument("--batch_size", type=int, default=8) |
| parser.add_argument("--grad_accum", type=int, default=4) |
| parser.add_argument("--lr", type=float, default=2e-5) |
| parser.add_argument("--max_synthetic", type=int, default=5000) |
| parser.add_argument("--max_nemotron_train", type=int, default=15000) |
| parser.add_argument("--max_nemotron_eval", type=int, default=2000) |
| parser.add_argument("--seed", type=int, default=42) |
| args = parser.parse_args() |
| random.seed(args.seed); np.random.seed(args.seed); torch.manual_seed(args.seed) |
| tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) |
| model = AutoModelForTokenClassification.from_pretrained(args.model, num_labels=NUM_LABELS, id2label=id2label, label2id=label2id, trust_remote_code=True, ignore_mismatched_sizes=True) |
| print("Generating synthetic data...") |
| synth = generate_synthetic_examples(args.max_synthetic, args.seed) |
| print(f"Synthetic examples: {len(synth)}") |
| print("Loading Nemotron-PII...") |
| nemotron_train = load_nemotron_split("train", args.max_nemotron_train) |
| nemotron_eval = load_nemotron_split("test", args.max_nemotron_eval) |
| print(f"Nemotron train: {len(nemotron_train)}, eval: {len(nemotron_eval)}") |
| train_examples = synth + nemotron_train |
| eval_examples = nemotron_eval |
| print("Tokenizing...") |
| train_tok = tokenize_and_align(train_examples, tokenizer) |
| eval_tok = tokenize_and_align(eval_examples, tokenizer) |
| train_ds = Dataset.from_dict(train_tok) |
| eval_ds = Dataset.from_dict(eval_tok) |
| data_collator = DataCollatorForTokenClassification(tokenizer) |
| training_args = TrainingArguments( |
| output_dir="/app/privacy-filter-checkpoints", |
| learning_rate=args.lr, |
| per_device_train_batch_size=args.batch_size, |
| per_device_eval_batch_size=args.batch_size, |
| num_train_epochs=args.epochs, |
| weight_decay=0.01, |
| eval_strategy="epoch", |
| save_strategy="epoch", |
| load_best_model_at_end=True, |
| metric_for_best_model="f1", |
| greater_is_better=True, |
| logging_strategy="steps", |
| logging_steps=50, |
| logging_first_step=True, |
| disable_tqdm=True, |
| push_to_hub=True, |
| hub_model_id=args.output_model, |
| report_to="trackio", |
| run_name=f"privacy-filter-enhanced-lr{args.lr}-bs{args.batch_size}-ep{args.epochs}", |
| project="privacy-filter-enhanced", |
| seed=args.seed, |
| bf16=True, |
| gradient_accumulation_steps=args.grad_accum, |
| dataloader_num_workers=4, |
| ) |
| trainer = Trainer(model=model, args=training_args, train_dataset=train_ds, eval_dataset=eval_ds, processing_class=tokenizer, data_collator=data_collator, compute_metrics=compute_metrics) |
| print("Starting training...") |
| trainer.train() |
| print("Pushing to hub...") |
| trainer.push_to_hub() |
| print("Done!") |
|
|
| if __name__ == "__main__": |
| main() |
|
|
