-
Notifications
You must be signed in to change notification settings - Fork 52
/
nb.py
88 lines (71 loc) · 2.55 KB
/
nb.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
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
"""
A collection of Naive Bayes type classifier models.
.. autoclass:: revscoring.scorer_models.GaussianNB
:members:
:member-order:
.. autoclass:: revscoring.scorer_models.MultinomialNB
:members:
:member-order:
.. autoclass:: revscoring.scorer_models.BernoulliNB
:members:
:member-order:
"""
import logging
from sklearn import naive_bayes
from .sklearn_classifier import ScikitLearnClassifier
logger = logging.getLogger("revscoring.scorers.nb")
class NB(ScikitLearnClassifier):
def __init__(self, features, *, version=None, nb=None,
sklearn_class=None, **kwargs):
if nb is None:
nb = sklearn_class(**kwargs)
super().__init__(features, classifier_model=nb, version=version)
NBModel = NB
"Alias for backwards compatibility"
class GaussianNB(NBModel):
"""
Implements a Gaussian Naive Bayes model.
:Params:
features : `list` ( :class:`revscoring.Feature` )
The features that the model will be trained on
version : str
A version string representing the version of the model
`**kwargs`
Passed to :class:`sklearn.naive_bayes.GaussianNB`
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, sklearn_class=naive_bayes.GaussianNB, **kwargs)
GaussianNBModel = GaussianNB
"Alias for backwards compatibility"
class MultinomialNB(NBModel):
"""
Implements a Multinomial Naive Bayes model.
:Params:
features : `list` ( :class:`revscoring.Feature` )
The features that the model will be trained on
version : str
A version string representing the version of the model
`**kwargs`
Passed to :class:`sklearn.naive_bayes.MultinomialNB`
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, sklearn_class=naive_bayes.MultinomialNB,
**kwargs)
MultinomialNBModel = MultinomialNB
"Alias for backwards compatibility"
class BernoulliNB(NBModel):
"""
Implements a Bernoulli Naive Bayes model.
:Params:
features : `list` ( :class:`revscoring.Feature` )
The features that the model will be trained on
version : str
A version string representing the version of the model
`**kwargs`
Passed to :class:`sklearn.naive_bayes.BernoulliNB`
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, sklearn_class=naive_bayes.BernoulliNB,
**kwargs)
BernoulliNBModel = BernoulliNB
"Alias for backwards compatibility"