File size: 3,041 Bytes
3dac39e | 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 | #!/usr/bin/env python3
"""Remove train-test leakage from Arcspan cybersecurity NER training data."""
import json
from pathlib import Path
DATA = Path("/home/ubuntu/alkyline/data/processed")
# --- Test/valid files (DO NOT modify) ---
TEST_VALID_FILES = [
"enriched_5class_test.jsonl",
"enriched_5class_valid_cleaned.jsonl",
"enriched_5class_valid.jsonl",
"cyner_test.jsonl",
"securebert2_5class_test.jsonl",
"securebert2_test.jsonl",
]
# --- Training files to clean ---
TRAIN_FILES = [
"enriched_5class_train_cleaned.jsonl",
"aptner_5class_train.jsonl",
"securebert2_5class_train.jsonl",
]
def load_texts(path: Path):
"""Yield text field from each line of a JSONL file."""
with open(path) as f:
for line in f:
line = line.strip()
if line:
yield json.loads(line)["text"]
def main():
# 1. Collect all test/valid texts and prefixes
exact_texts = set()
prefix_texts = set() # text[:80]
for fname in TEST_VALID_FILES:
p = DATA / fname
if not p.exists():
print(f" SKIP (not found): {fname}")
continue
count = 0
for text in load_texts(p):
exact_texts.add(text)
prefix_texts.add(text[:80])
count += 1
print(f" Loaded {count:,} test/valid examples from {fname}")
print(f"\nTotal unique test/valid texts: {len(exact_texts):,}")
print(f"Total unique test/valid prefixes: {len(prefix_texts):,}\n")
# 2. Clean each training file
for fname in TRAIN_FILES:
p = DATA / fname
if not p.exists():
print(f" SKIP (not found): {fname}")
continue
kept = []
removed_exact = 0
removed_prefix = 0
removed_trivial = 0
with open(p) as f:
for line in f:
line = line.strip()
if not line:
continue
rec = json.loads(line)
text = rec["text"]
# Trivial
if len(text.strip()) <= 2:
removed_trivial += 1
continue
# Exact match
if text in exact_texts:
removed_exact += 1
continue
# Prefix match
if text[:80] in prefix_texts:
removed_prefix += 1
continue
kept.append(line)
total_removed = removed_exact + removed_prefix + removed_trivial
stem = fname.replace(".jsonl", "")
out_path = DATA / f"{stem}_deleaked.jsonl"
with open(out_path, "w") as f:
for line in kept:
f.write(line + "\n")
print(f"{fname}:")
print(f" Removed: {total_removed:,} (exact={removed_exact}, prefix={removed_prefix}, trivial={removed_trivial})")
print(f" Kept: {len(kept):,}")
print(f" Written: {out_path.name}\n")
if __name__ == "__main__":
main()
|