forked from Eshe/competition-vd
47 lines
1.5 KiB
Python
Executable File
47 lines
1.5 KiB
Python
Executable File
import yaml
|
|
from vulnpatch.preprocessing.data_cleaner import DataCleaner
|
|
from vulnpatch.modeling.patch_classifier import PatchClassifier
|
|
from vulnpatch.utils.metrics import calculate_metrics, optimal_threshold
|
|
import pandas as pd
|
|
import numpy as np
|
|
from sklearn.model_selection import train_test_split
|
|
|
|
def load_config():
|
|
with open('config.yaml', 'r') as file:
|
|
return yaml.safe_load(file)
|
|
|
|
def main():
|
|
config = load_config()
|
|
|
|
data_cleaner = DataCleaner(config['train_data_path'])
|
|
prepared_data = data_cleaner.prepare_for_model()
|
|
|
|
train_data, val_data = train_test_split(prepared_data, test_size=0.2, random_state=42)
|
|
|
|
classifier = PatchClassifier()
|
|
|
|
train_results = classifier.classify_patches(train_data)
|
|
classifier.save_predictions(train_results, config['train_predictions_path'])
|
|
|
|
val_results = classifier.classify_patches(val_data)
|
|
classifier.save_predictions(val_results, config['val_predictions_path'])
|
|
|
|
threshold = optimal_threshold(val_results['label'], val_results['probability'])
|
|
|
|
val_metrics = calculate_metrics(
|
|
val_results['label'],
|
|
(val_results['probability'] >= threshold).astype(int),
|
|
val_results['probability']
|
|
)
|
|
|
|
print("验证集指标:")
|
|
for metric, value in val_metrics.items():
|
|
print(f"{metric}: {value:.4f}")
|
|
|
|
print(f"最佳阈值: {threshold:.4f}")
|
|
|
|
with open(config['model_output_path'], 'w') as f:
|
|
yaml.dump({'threshold': threshold}, f)
|
|
|
|
if __name__ == "__main__":
|
|
main() |