example.py

import logging
from analyzer import analyzer

LABELS = [
    "FUNCTION",
    "VERSION",
    "SOURCECODE",
    "DRIVER",
    "STRUCT",
    "VULNERABILITY",
    "CAPABILITY"
]

DESCRIPTION = '''
Stack-based buffer overflow in the ExecuteSoapAction function in the SOAPAction
handler in the HTTP service in MiniUPnP MiniUPnPd 1.0 allows remote attackers
to execute arbitrary code via a long quoted method.
'''

log = logging.getLogger('CVE_Analyzer')
log.setLevel(logging.DEBUG)

print('1) SPLIT DATASET TRAIN/TEST')
trainset, testset = analyzer.get_train_and_test_sets("./dataset/dataset.json", 2/3)

print('2) TRAIN PHASE')
model = analyzer.train(trainset, LABELS, 40, 0.35)

print('3) TEST PERFORMANCES')
performances = analyzer.test(model, testset)
analyzer.pp_performances(
    performances['accuracy'],
    performances['precision'],
    performances['recall'],
    performances['f_measure'])

print('4) SAVE MODEL')
analyzer.save_model('./models/example_model', 'example_model', model)

print('5) GET MODEL')
analyzer.get_model('./models/example_model')

print('6) ANALYZE DESCRIPTION')
prediction = analyzer.get_prediction_for_description(model, DESCRIPTION)
print(prediction)