| """Zap tokenization, windowing, label space and training-time augmentation.""" |
| import random |
|
|
| from zap.extract import SPECIAL_TOKENS, serialize_ctx, serialize_field |
|
|
| FIELD_LABELS = ["username", "email", "current-password", "new-password", "one-time-code", "tel", "name", |
| "given-name", "family-name", "organization", "street-address", "address-line2", "address-level2", |
| "address-level1", "postal-code", "country", "cc-name", "cc-number", "cc-exp", "cc-exp-month", |
| "cc-exp-year", "cc-csc", "bday", "search", "captcha", "other"] |
| FORM_LABELS = ["login", "signup", "password-change", "password-recovery", "otp", "payment", "address", "newsletter", |
| "contact", "search", "other"] |
| |
| ID2LABEL = [f"form:{x}" for x in FORM_LABELS] + [f"field:{x}" for x in FIELD_LABELS] |
| LABEL2ID = {l: i for i, l in enumerate(ID2LABEL)} |
| N_FORM = len(FORM_LABELS) |
| MAX_LEN = 512 |
|
|
|
|
| def windows(tok, ctx, fields, max_len=MAX_LEN): |
| """Split a form into model inputs: each window = [CLS] ctx + as many whole fields as fit + [SEP]. |
| |
| Returns a list of (input_ids, field_positions, field_indices). Shared by training and inference; |
| the JS/extension port must do the same. |
| """ |
| ctx_ids = tok(serialize_ctx(ctx), add_special_tokens=False)["input_ids"][:160] |
| f_ids = [tok(serialize_field(f), add_special_tokens=False)["input_ids"][:120] for f in fields] |
| cls, sep = tok.cls_token_id, tok.sep_token_id |
| out, cur, cur_idx = [], [], [] |
| budget = max_len - 2 - len(ctx_ids) |
| used = 0 |
| for i, ids in enumerate(f_ids): |
| if cur and used + len(ids) > budget: |
| out.append((cur, cur_idx)) |
| cur, cur_idx, used = [], [], 0 |
| cur.append(ids) |
| cur_idx.append(i) |
| used += len(ids) |
| if cur or not out: |
| out.append((cur, cur_idx)) |
| res = [] |
| for chunk, idx in out: |
| ids = [cls] + ctx_ids |
| pos = [] |
| for f in chunk: |
| pos.append(len(ids)) |
| ids += f |
| ids.append(sep) |
| res.append((ids, pos, idx)) |
| return res |
|
|
|
|
| def augment(ctx, fields, rng: random.Random): |
| """Hint dropout: teaches the model to cope with missing autocomplete/type/labels (e.g. revealed passwords).""" |
| ctx = dict(ctx) |
| for k, p in (("title", 0.1), ("url", 0.1), ("txt", 0.15), ("link", 0.1), ("hid", 0.2), ("head", 0.1)): |
| if rng.random() < p: |
| ctx[k] = "" |
| out = [] |
| for f in fields: |
| f = dict(f) |
| if rng.random() < 0.25: |
| f["ac"] = "" |
| if f["type"] == "input password" and rng.random() < 0.08: |
| f["type"] = "input text" |
| for k, p in (("name", 0.05), ("id", 0.05), ("class", 0.1), ("lbl", 0.1), ("ph", 0.07), ("near", 0.1), |
| ("after", 0.1), ("data", 0.1)): |
| if rng.random() < p: |
| f[k] = "" |
| out.append(f) |
| return ctx, out |
|
|
|
|
| def train_tokenizer(texts_iter, out_dir, vocab_size=16384): |
| import os |
| from tokenizers import BertWordPieceTokenizer |
| os.makedirs(out_dir, exist_ok=True) |
| t = BertWordPieceTokenizer(lowercase=True, strip_accents=False, clean_text=True, handle_chinese_chars=True) |
| t.train_from_iterator(texts_iter, vocab_size=vocab_size, min_frequency=3, |
| special_tokens=["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"] + SPECIAL_TOKENS, |
| limit_alphabet=6000, wordpieces_prefix="##") |
| from transformers import PreTrainedTokenizerFast |
| hf = PreTrainedTokenizerFast(tokenizer_object=t._tokenizer, unk_token="[UNK]", pad_token="[PAD]", |
| cls_token="[CLS]", sep_token="[SEP]", mask_token="[MASK]", |
| model_max_length=MAX_LEN) |
| hf.save_pretrained(out_dir) |
| return hf |
|
|
