| """Turn (text, candidate span) into a fixed-size byte window for the model.""" |
| import numpy as np |
|
|
| LEFT, CAND, RIGHT = 64, 96, 32 |
| SEQ = LEFT + CAND + RIGHT + 2 |
| PAD, SEP1, SEP2 = 0, 1, 2 |
| VOCAB = 256 + 3 |
| TYPES = ["none", "password", "api_key", "token", "private_key", "generic"] |
|
|
|
|
| def _b(s): |
| return np.frombuffer(s.encode("utf-8", "replace"), dtype=np.uint8).astype(np.int64) + 3 |
|
|
|
|
| def encode(text, start, end): |
| """Return (ids[SEQ], seg[SEQ]) where seg is 0 pad, 1 left ctx, 2 candidate, 3 right ctx.""" |
| left = _b(text[max(0, start - LEFT * 2):start])[-LEFT:] |
| cand = _b(text[start:end]) |
| if len(cand) > CAND: |
| cand = np.concatenate([cand[:CAND // 2], cand[-CAND // 2:]]) |
| right = _b(text[end:end + RIGHT * 2])[:RIGHT] |
| ids = np.zeros(SEQ, dtype=np.int64) |
| seg = np.zeros(SEQ, dtype=np.int64) |
| |
| ids[LEFT - len(left):LEFT] = left; seg[LEFT - len(left):LEFT] = 1 |
| ids[LEFT] = SEP1; seg[LEFT] = 2 |
| o = LEFT + 1 |
| ids[o:o + len(cand)] = cand; seg[o:o + len(cand)] = 2 |
| o += len(cand) |
| ids[o] = SEP2; seg[o] = 2 |
| o += 1 |
| ids[o:o + len(right)] = right; seg[o:o + len(right)] = 3 |
| return ids, seg |
|
|
|
|
| def encode_string(value): |
| """Classify a bare string with no context.""" |
| return encode(value, 0, len(value)) |
|
|
|
|
| def label_candidate(cs, ce, spans): |
| """Label a candidate against ground-truth spans [(s, e, label, kind)]. |
| |
| Returns (label, kind) or None when ambiguous (partial overlap) so the sample is dropped. |
| """ |
| clen = ce - cs |
| for s, e, lab, kind in spans: |
| inter = min(ce, e) - max(cs, s) |
| if inter <= 0: |
| continue |
| if lab == 1: |
| if cs <= s and ce >= e and clen <= (e - s) + 40: |
| return 1, kind |
| if inter >= 0.8 * (e - s) and inter >= 0.6 * clen: |
| return 1, kind |
| return None |
| else: |
| if inter >= 0.6 * clen: |
| return 0, kind |
| return None |
| return 0, "other" |
|
|
