File size: 3,902 Bytes
e4805d0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | """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"]
# One token-classification label space: [CLS] carries form:*, each [FLD] carries field:*.
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)) # f starts with the [FLD] token
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
|