forked from Eshe/competition-vd
44 lines
1.4 KiB
Python
Executable File
44 lines
1.4 KiB
Python
Executable File
import sys
|
|
import os
|
|
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from src.model.llm_inference import initialize_model, generate_predictions
|
|
from src.data.data_loader import load_dataset, save_jsonl
|
|
from src.utils.preprocessing import prepare_messages, extract_cwe_id
|
|
from src.analysis.metrics import calculate_metrics, plot_confusion_matrix, plot_metrics
|
|
|
|
def main():
|
|
tokenizer, llm, sampling_params = initialize_model()
|
|
|
|
test_data = load_dataset('data/test.jsonl')
|
|
|
|
messages_list = [prepare_messages(item) for item in test_data]
|
|
|
|
predictions = generate_predictions(tokenizer, llm, sampling_params, messages_list)
|
|
|
|
predicted_cwes = [extract_cwe_id(pred) for pred in predictions]
|
|
true_cwes = [item['cwe_id'] for item in test_data]
|
|
|
|
metrics = calculate_metrics(true_cwes, predicted_cwes)
|
|
print("评估指标:", metrics)
|
|
|
|
unique_cwes = list(set(true_cwes + predicted_cwes))
|
|
plot_confusion_matrix(true_cwes, predicted_cwes, unique_cwes)
|
|
|
|
plot_metrics(metrics)
|
|
|
|
results = [
|
|
{
|
|
'function': item['function'],
|
|
'cve_description': item['cve_description'],
|
|
'true_cwe': item['cwe_id'],
|
|
'predicted_cwe': pred_cwe,
|
|
'model_output': pred
|
|
}
|
|
for item, pred_cwe, pred in zip(test_data, predicted_cwes, predictions)
|
|
]
|
|
save_jsonl(results, 'results/output.json')
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |