forked from mozilla/bugbug
-
Notifications
You must be signed in to change notification settings - Fork 0
/
run.py
67 lines (56 loc) · 2.6 KB
/
run.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
# -*- coding: utf-8 -*-
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this file,
# You can obtain one at http://mozilla.org/MPL/2.0/.
import argparse
import numpy as np
from bugbug import bugzilla
from bugbug import db
from bugbug import repository # noqa
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--lemmatization', help='Perform lemmatization (using spaCy)', action='store_true')
parser.add_argument('--train', help='Perform training', action='store_true')
parser.add_argument('--goal', help='Goal of the classifier', choices=['bug', 'regression', 'tracking', 'qaneeded', 'uplift', 'component', 'devdocneeded'], default='bug')
parser.add_argument('--classify', help='Perform evaluation', action='store_true')
args = parser.parse_args()
model_file_name = f'{args.goal}model'
if args.goal == 'bug':
from bugbug.models.bug import BugModel
model_class = BugModel
elif args.goal == 'regression':
from bugbug.models.regression import RegressionModel
model_class = RegressionModel
elif args.goal == 'tracking':
from bugbug.models.tracking import TrackingModel
model_class = TrackingModel
elif args.goal == 'qaneeded':
from bugbug.models.qaneeded import QANeededModel
model_class = QANeededModel
elif args.goal == 'uplift':
from bugbug.models.uplift import UpliftModel
model_class = UpliftModel
elif args.goal == 'component':
from bugbug.models.component import ComponentModel
model_class = ComponentModel
elif args.goal == 'devdocneeded':
from bugbug.models.devdocneeded import DevDocNeededModel
model_class = DevDocNeededModel
if args.train:
db.download()
model = model_class(args.lemmatization)
model.train()
else:
model = model_class.load(model_file_name)
if args.classify:
for bug in bugzilla.get_bugs():
print(f'https://bugzilla.mozilla.org/show_bug.cgi?id={ bug["id"] } - { bug["summary"]} ')
probas, importances = model.classify(bug, probabilities=True, importances=True)
feature_names = model.get_feature_names()
for i, (index, is_positive, contrib) in enumerate(importances[:20]):
print(f'{i + 1}. \'{feature_names[index]}\' ({"+" if is_positive else "-"}{contrib})')
if np.argmax(probas) == 1:
print(f'Positive! {probas}')
else:
print(f'Negative! {probas}')
input()