File size: 4,724 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 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | """Dedup harvested forms, assign host-level splits, and build prioritized teacher-labeling chunks.
Usage: python select.py OUT_DIR PATTERN [PATTERN ...] [--labeled DIR ...]
Signatures already labeled in a previous round (DIR/queue-*.jsonl + DIR/labels-*.jsonl) are not queued again.
Writes OUT_DIR/queue-NNN.jsonl (teacher input, 10k rows each, interleaved priorities),
OUT_DIR/reps.jsonl.gz (one representative per signature, with split), OUT_DIR/select_stats.json.
"""
import gzip
import hashlib
import json
import os
import random
import sys
from collections import Counter
sys.path.insert(0, "/job/work")
from label.common import load_forms # noqa: E402
SEARCHY = {"q", "s", "search", "query", "keyword", "keywords", "search query", "k", "term", "search term"}
MIX = [(0, 0.40), (1, 0.25), (2, 0.17), (3, 0.12), (4, 0.06)]
CAPS = {3: 60000, 4: 20000}
def split_of(host):
h = int(hashlib.md5(host.encode()).hexdigest()[:8], 16) % 100
return "test" if h < 4 else ("val" if h < 6 else "train")
def prio(r):
fs = r["fields"]
types = [f["type"] for f in fs]
if "input password" in types:
return 0
if len(fs) >= 3:
return 1
if len(fs) == 2 or any(t in ("input email", "input tel") for t in types):
return 2
f = fs[0]
if f["type"] == "input search" or f["name"] in SEARCHY or f["id"] in SEARCHY:
return 4
return 3
def labeled_sigs(dirs):
import glob
out = {}
for d in dirs:
uid2sig = {}
for q in glob.glob(os.path.join(d, "queue-*.jsonl")):
for line in open(q):
j = json.loads(line)
uid2sig[j["uid"]] = j["sig"]
for lf in glob.glob(os.path.join(d, "labels-*.jsonl")):
for line in open(lf):
try:
j = json.loads(line)
except json.JSONDecodeError: # partial last line of an interrupted run
continue
if j.get("form") and j.get("fields") and j["uid"] in uid2sig:
out[uid2sig[j["uid"]]] = j
return out
def main(out_dir, patterns, labeled_dirs=()):
os.makedirs(out_dir, exist_ok=True)
done = labeled_sigs(labeled_dirs)
print(f"{len(done)} signatures already labeled", flush=True)
reps, copies, seen = {}, Counter(), set()
n = 0
for pat in patterns:
for r in load_forms(pat):
n += 1
if r["uid"] in seen:
continue
seen.add(r["uid"])
copies[r["sig"]] += 1
cur = reps.get(r["sig"])
if cur is None or (not cur.get("html") and r.get("html")):
reps[r["sig"]] = r
print(f"{n} rows, {len(seen)} unique uids, {len(reps)} unique signatures", flush=True)
rng = random.Random(13)
pools = {p: [] for p, _ in MIX}
with gzip.open(os.path.join(out_dir, "reps.jsonl.gz"), "wt") as g:
for sig, r in reps.items():
r["split"] = split_of(r["host"])
r["copies"] = copies[sig]
r["prio"] = prio(r)
g.write(json.dumps(r, ensure_ascii=False) + "\n")
if r.get("html") and sig not in done:
pools[r["prio"]].append(r)
for p in pools:
rng.shuffle(pools[p])
if p in CAPS:
pools[p] = pools[p][:CAPS[p]]
stats = {"rows": n, "uids": len(seen), "sigs": len(reps), "pools": {p: len(v) for p, v in pools.items()},
"splits": Counter(split_of(r["host"]) for r in reps.values())}
print(stats, flush=True)
order, idx = [], {p: 0 for p in pools}
while any(idx[p] < len(pools[p]) for p in pools):
x, acc = rng.random(), 0.0
live = [(p, w) for p, w in MIX if idx[p] < len(pools[p])]
tot = sum(w for _, w in live)
for p, w in live:
acc += w / tot
if x <= acc:
order.append(pools[p][idx[p]])
idx[p] += 1
break
keep = ("uid", "url", "host", "ctx", "fields", "html", "sig", "split", "prio")
for c in range(0, len(order), 10000):
with open(os.path.join(out_dir, f"queue-{c // 10000:03d}.jsonl"), "w") as f:
for r in order[c:c + 10000]:
f.write(json.dumps({k: r[k] for k in keep}, ensure_ascii=False) + "\n")
stats["queued"] = len(order)
json.dump(stats, open(os.path.join(out_dir, "select_stats.json"), "w"), indent=1)
print("queued", len(order), flush=True)
if __name__ == "__main__":
import argparse
ap = argparse.ArgumentParser()
ap.add_argument("out")
ap.add_argument("patterns", nargs="+")
ap.add_argument("--labeled", action="append", default=[])
a = ap.parse_args()
main(a.out, a.patterns, a.labeled)
|