Spaces:
Sleeping
Sleeping
File size: 3,310 Bytes
bca5172 | 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 | """
Prepare dataset for fine-tuning: read cleaned JSONL, optionally apply prompt formatting,
create train/val/test splits (stratified if labels exist), and write HF-compatible JSON/JSONL.
Usage (bash):
python scripts/prepare_dataset.py --cleaned data/cleaned/cleaned_extracted_text.jsonl --out_dir data/prepared_run
"""
import argparse
import json
import os
from collections import Counter
from sklearn.model_selection import train_test_split
# Compact prompt template for training (IDs + short names)
COMPACT_PROMPT = (
"CATEGORIES: 1:MaliciousData 2:Ransomware 3:Marketplace 4:MalwareTools 5:ThreatActor "
"6:VulnIntel 7:InfraOps 8:Credentials 9:SocialEng 10:General\n\nPOST:\n{txt}\n\nCATEGORY_ID:"
)
def format_prompt(text):
return COMPACT_PROMPT.format(txt=text)
def read_jsonl(path):
with open(path, 'r', encoding='utf-8') as f:
for line in f:
yield json.loads(line)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--cleaned', required=True)
parser.add_argument('--out_dir', default='data/prepared_run')
parser.add_argument('--use_prompt', action='store_true')
args = parser.parse_args()
os.makedirs(args.out_dir, exist_ok=True)
samples = list(read_jsonl(args.cleaned))
print(f"Loaded {len(samples)} cleaned samples")
# If labels exist, create stratified splits
labels = [s.get('label') for s in samples]
has_labels = all([l is not None for l in labels])
if has_labels:
# map label names to numeric ids
unique = sorted(list(set(labels)))
label2id = {name: i for i, name in enumerate(unique)}
for s in samples:
s['label_id'] = label2id[s['label']]
# stratified split
train, temp = train_test_split(samples, test_size=0.3, stratify=[s['label_id'] for s in samples], random_state=42)
val, test = train_test_split(temp, test_size=0.5, stratify=[s['label_id'] for s in temp], random_state=42)
else:
# no labels: split randomly
train, val_test = train_test_split(samples, test_size=0.3, random_state=42)
val, test = train_test_split(val_test, test_size=0.5, random_state=42)
label2id = {}
# Optionally apply prompt formatting
def maybe_prompt(s):
txt = s['text']
return format_prompt(txt) if args.use_prompt else txt
def write_jsonl(l, path):
with open(path, 'w', encoding='utf-8') as f:
for s in l:
out = {
'text': maybe_prompt(s),
}
if 'label_id' in s:
out['label'] = s['label_id']
f.write(json.dumps(out, ensure_ascii=False) + '\n')
write_jsonl(train, os.path.join(args.out_dir, 'train.jsonl'))
write_jsonl(val, os.path.join(args.out_dir, 'val.jsonl'))
write_jsonl(test, os.path.join(args.out_dir, 'test.jsonl'))
# Write labels mapping
with open(os.path.join(args.out_dir, 'labels.json'), 'w', encoding='utf-8') as f:
json.dump({'label2id': label2id}, f, ensure_ascii=False, indent=2)
print(f"Prepared splits -> train:{len(train)} val:{len(val)} test:{len(test)}")
if __name__ == '__main__':
main()
|