File size: 2,247 Bytes
b07771a | 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 | """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 # bytes 0..2 are reserved; real bytes are shifted by +3
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: # keep head and tail of very long values (e.g. PEM blocks)
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)
# left context is right-aligned so the candidate always starts at the same offset
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 # candidate contains the secret (e.g. 'Bearer xxx', URL with creds)
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"
|