diff --git a/bayes.go b/bayes.go index 4cbc5ee..19b8ad1 100644 --- a/bayes.go +++ b/bayes.go @@ -5,14 +5,14 @@ import ( ) var ( - smoother = 1 // laplace + smoother = 1 // laplace defaultMinClassSize = 5 ) type Classifier struct { - Tokenizer *tokenizer `json:"-"` - Matrix *sparseMatrix `json:"matrix"` - MinClassSize int + Tokenizer *tokenizer `json:"-"` + Matrix *sparseMatrix `json:"matrix"` + MinClassSize int `json:"minClassSize"` } // Create a new multibayes classifier. @@ -24,8 +24,8 @@ func NewClassifier() *Classifier { sparse := newSparseMatrix() return &Classifier{ - Tokenizer: tokenize, - Matrix: sparse, + Tokenizer: tokenize, + Matrix: sparse, MinClassSize: defaultMinClassSize, } } diff --git a/encoding.go b/encoding.go index 99fec8b..3c18fab 100644 --- a/encoding.go +++ b/encoding.go @@ -6,11 +6,12 @@ import ( ) type jsonableClassifier struct { - Matrix *sparseMatrix `json:"matrix"` + Matrix *sparseMatrix `json:"matrix"` + MinClassSize int `json:"minClassSize"` } func (c *Classifier) MarshalJSON() ([]byte, error) { - return json.Marshal(&jsonableClassifier{c.Matrix}) + return json.Marshal(&jsonableClassifier{c.Matrix, c.MinClassSize}) } func (c *Classifier) UnmarshalJSON(buf []byte) error { @@ -23,6 +24,7 @@ func (c *Classifier) UnmarshalJSON(buf []byte) error { *c = *NewClassifier() c.Matrix = j.Matrix + c.MinClassSize = j.MinClassSize return nil }