#!/bin/bash
# 数据下载与清洗脚本
# 适用于 Qwen3.5-0.8B-HOS 训练流程

set -e

echo "============================================"
echo "  准备训练数据集"
echo "============================================"

# 数据集目录
DATA_DIR="datasets"
mkdir -p "$DATA_DIR"

# 从 HuggingFace 下载数据集
echo "[1/4] 下载数据集..."
python -c "
from datasets import load_dataset
from huggingface_hub import hf_hub_download
import os

# 下载 infosec-dataset-training 数据集
print('下载 cybersecurity_hq...')
ds = load_dataset('lxcxjxhx/infosec-dataset-training', 'cybersecurity_hq')
ds['train'].to_json(os.path.join('$DATA_DIR', 'cybersecurity_hq.jsonl'), orient='records', lines=True, force_ascii=False)

print('下载 cybersecurity_sharegpt...')
ds = load_dataset('lxcxjxhx/infosec-dataset-training', 'cybersecurity_sharegpt')
ds['train'].to_json(os.path.join('$DATA_DIR', 'cybersecurity_sharegpt.jsonl'), orient='records', lines=True, force_ascii=False)

print('下载 cnvd_vulnerability...')
ds = load_dataset('lxcxjxhx/infosec-dataset-training', 'cnvd_vulnerability')
ds['train'].to_json(os.path.join('$DATA_DIR', 'cnvd_vulnerability.jsonl'), orient='records', lines=True, force_ascii=False)

print('下载 merged_cybersec...')
ds = load_dataset('lxcxjxhx/infosec-dataset-training', 'merged_cybersec')
ds['train'].to_json(os.path.join('$DATA_DIR', 'merged_cybersec.jsonl'), orient='records', lines=True, force_ascii=False)

print('数据集下载完成!')
"

# 数据清洗
echo "[2/4] 数据清洗..."
python scripts/clean_data.py

# 验证数据集
echo "[3/4] 验证数据集..."
python -c "
import json
import os

data_dir = '$DATA_DIR'
files = ['cybersecurity_hq.jsonl', 'cybersecurity_sharegpt.jsonl', 'cnvd_vulnerability.jsonl', 'merged_cybersec.jsonl']

for f in files:
    path = os.path.join(data_dir, f)
    if os.path.exists(path):
        with open(path, 'r', encoding='utf-8') as file:
            count = sum(1 for _ in file)
        print(f'{f}: {count} 条记录')
    else:
        print(f'{f}: 文件不存在!')
"

# 复制 dataset_info.json
echo "[4/4] 配置数据集信息..."
cp datasets/dataset_info.json "$DATA_DIR/dataset_info.json" 2>/dev/null || echo "dataset_info.json 已存在"

echo "============================================"
echo "  数据集准备完成!"
echo "============================================"