forked from Eshe/competition-vd
40 lines
1.6 KiB
Python
Executable File
40 lines
1.6 KiB
Python
Executable File
import sys
|
|
import os
|
|
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from src.data.data_loader import DataLoader
|
|
from src.data.preprocessor import Preprocessor
|
|
from src.data.feature_extractor import FeatureExtractor
|
|
import argparse
|
|
import json
|
|
|
|
def main(args):
|
|
data_loader = DataLoader(args.data_dir)
|
|
preprocessor = Preprocessor()
|
|
feature_extractor = FeatureExtractor()
|
|
|
|
for dataset in ['train', 'valid', 'test']:
|
|
print(f"Processing {dataset} dataset...")
|
|
data = data_loader.load_jsonl(f'{dataset}.jsonl')
|
|
processed_data = preprocessor.process_data(data)
|
|
features = feature_extractor.extract_features(processed_data)
|
|
|
|
output_dir = os.path.join(args.output_dir, dataset)
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
with open(os.path.join(output_dir, 'processed_data.json'), 'w') as f:
|
|
json.dump(processed_data, f)
|
|
|
|
np.save(os.path.join(output_dir, 'features.npy'), features['features'])
|
|
np.save(os.path.join(output_dir, 'labels.npy'), features['labels'])
|
|
np.save(os.path.join(output_dir, 'cvss_scores.npy'), features['cvss_scores'])
|
|
|
|
print("Preprocessing completed successfully.")
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Run preprocessing on vulnerability detection datasets")
|
|
parser.add_argument('--data_dir', type=str, default='data/raw', help='Directory containing raw data files')
|
|
parser.add_argument('--output_dir', type=str, default='data/processed', help='Directory to save processed data')
|
|
args = parser.parse_args()
|
|
|
|
main(args) |