File size: 4,591 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 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | #!/usr/bin/env python3
"""
Class balance audit for Arcspan NER datasets.
Analyzes entity distribution across fixed/deleaked training sets.
"""
import json
from collections import defaultdict
from pathlib import Path
CLASSES = ["Malware", "Indicator", "Organization", "System", "Vulnerability"]
FILES = [
"/home/ubuntu/alkyline/data/processed/enriched_5class_train_cleaned_trimmed.jsonl",
"/home/ubuntu/alkyline/data/processed/enriched_5class_train_cleaned_deleaked.jsonl",
"/home/ubuntu/alkyline/data/processed/aptner_5class_train_deleaked.jsonl",
"/home/ubuntu/alkyline/data/processed/securebert2_5class_train_deleaked.jsonl",
"/home/ubuntu/alkyline/data/processed/defanged_augmented.jsonl",
]
def analyze_file(filepath):
"""Analyze a single JSONL file."""
if not Path(filepath).exists():
return None
stats = {
"total_examples": 0,
"class_counts": defaultdict(int),
"all_o_examples": 0,
"total_spans": 0,
}
with open(filepath, 'r') as f:
for line in f:
try:
record = json.loads(line.strip())
stats["total_examples"] += 1
spans = record.get("spans", {})
# Check if this is an all-O example (no entities)
if not spans:
stats["all_o_examples"] += 1
else:
# Count entities by class
for label_key in spans:
# Parse label key format: "Label: entity_text" or just class name
if ": " in label_key:
class_name = label_key.split(": ")[0]
else:
class_name = label_key
# Count occurrences of this class in this example
offsets = spans[label_key]
if isinstance(offsets, list) and len(offsets) > 0:
count = len(offsets) if isinstance(offsets[0], list) else 1
stats["class_counts"][class_name] += count
stats["total_spans"] += count
except json.JSONDecodeError as e:
print(f" ⚠ JSON error in {Path(filepath).name}: {e}")
continue
return stats
def format_report(filename, stats):
"""Format stats for a single file."""
if stats is None:
return f" ✗ {filename}: NOT FOUND\n"
total = stats["total_examples"]
all_o_pct = 100.0 * stats["all_o_examples"] / total if total > 0 else 0
# Get min/max for imbalance ratio
class_counts = {c: stats["class_counts"].get(c, 0) for c in CLASSES}
nonzero_counts = [c for c in class_counts.values() if c > 0]
if len(nonzero_counts) < 2:
imbalance_ratio = 1.0
else:
imbalance_ratio = max(nonzero_counts) / min(nonzero_counts)
lines = [f" {filename}"]
lines.append(f" Examples: {total:,} | All-O: {stats['all_o_examples']:,} ({all_o_pct:.1f}%)")
lines.append(f" Total spans: {stats['total_spans']:,} | Imbalance ratio: {imbalance_ratio:.2f}x")
for cls in CLASSES:
count = class_counts[cls]
lines.append(f" {cls}: {count:,}")
return "\n".join(lines) + "\n"
# Main
print("=" * 80)
print("ARCSPAN NER CLASS BALANCE AUDIT")
print(f"Classes: {', '.join(CLASSES)}")
print("=" * 80)
print()
all_stats = {}
combined = {
"total_examples": 0,
"class_counts": defaultdict(int),
"all_o_examples": 0,
"total_spans": 0,
}
for filepath in FILES:
filename = Path(filepath).name
stats = analyze_file(filepath)
all_stats[filename] = stats
if stats:
print(format_report(filename, stats))
combined["total_examples"] += stats["total_examples"]
combined["all_o_examples"] += stats["all_o_examples"]
combined["total_spans"] += stats["total_spans"]
for cls in CLASSES:
combined["class_counts"][cls] += stats["class_counts"][cls]
else:
print(f" ✗ {filename}: NOT FOUND\n")
print("\n" + "=" * 80)
print("COMBINED TOTAL (all files)")
print("=" * 80)
print(format_report("COMBINED", combined))
# Class imbalance for combined
combined_class_counts = {c: combined["class_counts"][c] for c in CLASSES}
nonzero = [c for c in combined_class_counts.values() if c > 0]
if len(nonzero) >= 2:
combined_imbalance = max(nonzero) / min(nonzero)
print(f" Overall imbalance ratio: {combined_imbalance:.2f}x")
print(f" Most common: {max(nonzero):,} | Least common: {min(nonzero):,}")
|