Support multiclass models in the generate_sheet.py script (#896)

Fixes #822
This commit is contained in:
SanthoshBala18 2019-09-02 18:14:54 +05:30 коммит произвёл Marco
Родитель ffec2e35b5
Коммит 391121970a
1 изменённых файлов: 10 добавлений и 1 удалений

Просмотреть файл

@ -5,6 +5,8 @@ import csv
import os
from datetime import datetime, timedelta
import numpy as np
from bugbug import bugzilla
from bugbug.models import get_model_class
@ -31,10 +33,17 @@ def generate_sheet(model_name, token, days, threshold):
for bug in bugs.values():
p = model.classify(bug, probabilities=True)
probability = p[0]
if len(probability) > 2:
index = np.argmax(probability)
prediction = model.class_names[index]
else:
prediction = "y" if probability[1] >= threshold else "n"
rows.append(
[
f'https://bugzilla.mozilla.org/show_bug.cgi?id={bug["id"]}',
"y" if p[0][1] >= threshold else "n",
prediction,
"",
bug["summary"],
]