forked from Eshe/competition-vd
96 lines
4.3 KiB
Python
Executable File
96 lines
4.3 KiB
Python
Executable File
import sys
|
|
import os
|
|
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from src.models.linevul import LineVulModel
|
|
from src.models.xgboost_classifier import XGBoostVulClassifier, XGBoostVulRegressor
|
|
from src.data.data_loader import DataLoader
|
|
from src.evaluation.evaluator import Evaluator
|
|
from src.visualization.plot_utils import VisualizationUtils
|
|
import argparse
|
|
import torch
|
|
import numpy as np
|
|
import json
|
|
|
|
def load_linevul_model(model_path, num_labels):
|
|
model = LineVulModel(num_labels=num_labels)
|
|
model.load_state_dict(torch.load(model_path))
|
|
return model
|
|
|
|
def load_xgboost_models(classifier_path, regressor_path):
|
|
classifier = XGBoostVulClassifier()
|
|
classifier.model.load_model(classifier_path)
|
|
regressor = XGBoostVulRegressor()
|
|
regressor.model.load_model(regressor_path)
|
|
return classifier, regressor
|
|
|
|
def main(args):
|
|
data_loader = DataLoader(args.data_dir)
|
|
evaluator = Evaluator()
|
|
viz_utils = VisualizationUtils()
|
|
|
|
test_data = {
|
|
'features': np.load(os.path.join(args.data_dir, 'test', 'features.npy')),
|
|
'labels': np.load(os.path.join(args.data_dir, 'test', 'labels.npy')),
|
|
'cvss_scores': np.load(os.path.join(args.data_dir, 'test', 'cvss_scores.npy'))
|
|
}
|
|
|
|
if args.model == 'linevul':
|
|
model = load_linevul_model(args.model_path, num_labels=len(set(test_data['labels'])))
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
model.to(device)
|
|
model.eval()
|
|
|
|
test_dataset = torch.utils.data.TensorDataset(
|
|
torch.tensor(test_data['features']).float(),
|
|
torch.tensor(test_data['labels']).long()
|
|
)
|
|
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size)
|
|
|
|
all_preds = []
|
|
all_labels = []
|
|
with torch.no_grad():
|
|
for batch in test_loader:
|
|
inputs, labels = batch
|
|
inputs = inputs.to(device)
|
|
outputs = model(inputs)
|
|
_, preds = torch.max(outputs, 1)
|
|
all_preds.extend(preds.cpu().numpy())
|
|
all_labels.extend(labels.numpy())
|
|
|
|
results = evaluator.evaluate_classification(all_labels, all_preds)
|
|
elif args.model == 'xgboost':
|
|
classifier, regressor = load_xgboost_models(args.classifier_path, args.regressor_path)
|
|
class_preds = classifier.predict(test_data['features'])
|
|
reg_preds = regressor.predict(test_data['features'])
|
|
|
|
results = evaluator.evaluate_classification(test_data['labels'], class_preds)
|
|
results.update(evaluator.evaluate_regression(test_data['cvss_scores'], reg_preds))
|
|
else:
|
|
raise ValueError(f"Unsupported model: {args.model}")
|
|
|
|
evaluator.evaluate_and_print(results)
|
|
|
|
viz_utils.plot_confusion_matrix(test_data['labels'], all_preds if args.model == 'linevul' else class_preds,
|
|
labels=list(set(test_data['labels'])), title=f'Confusion Matrix - {args.model}')
|
|
|
|
if args.model == 'xgboost':
|
|
feature_importance = classifier.feature_importance()
|
|
viz_utils.plot_feature_importance(feature_importance, [f'Feature_{i}' for i in range(len(feature_importance))],
|
|
title=f'Feature Importance - {args.model}')
|
|
|
|
with open(os.path.join(args.output_dir, f'{args.model}_evaluation_results.json'), 'w') as f:
|
|
json.dump(results, f, indent=2)
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Run evaluation for vulnerability detection models")
|
|
parser.add_argument('--model', type=str, choices=['linevul', 'xgboost'], required=True, help='Model to evaluate')
|
|
parser.add_argument('--data_dir', type=str, default='data/processed', help='Directory containing processed data')
|
|
parser.add_argument('--model_path', type=str, help='Path to the trained LineVul model')
|
|
parser.add_argument('--classifier_path', type=str, help='Path to the trained XGBoost classifier')
|
|
parser.add_argument('--regressor_path', type=str, help='Path to the trained XGBoost regressor')
|
|
parser.add_argument('--output_dir', type=str, default='outputs/evaluation', help='Directory to save evaluation results')
|
|
parser.add_argument('--batch_size', type=int, default=32, help='Batch size for evaluation')
|
|
args = parser.parse_args()
|
|
|
|
main(args) |