topic.go

// Package topic is a single-topic content classifier over the gist feature
// engine: it wraps a binary grove model plus the featurizer that produced its
// inputs, so one artifact answers "is this <topic>?" (e.g. gambling yes/no). A
// multi-topic categorizer is a set of these, one per topic (one-vs-rest).
package topic

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"os"

	"github.com/netstar-labs/gist"
	"github.com/netstar-labs/grove"
)

// Classifier scores text for one topic. Build one with [New] after training a
// binary grove model on gist feature vectors, or load a saved one with [Load]/
// [Open]. Score returns P(topic); Is applies the threshold.
type Classifier struct {
	Name      string
	Threshold float64
	feat      gist.Featurizer
	model     *grove.Model
}

// New wraps a trained binary grove model and the featurizer used to build its
// training vectors. The model must be binary (grove.Binary): P(topic) is P(class
// = 1). threshold is the P(topic) at or above which [Classifier.Is] returns true.
func New(name string, threshold float64, feat gist.Featurizer, m *grove.Model) (*Classifier, error) {
	if m == nil {
		return nil, fmt.Errorf("topic: nil model")
	}
	if m.NumClass != 1 {
		return nil, fmt.Errorf("topic: model NumClass=%d; a topic classifier needs a binary grove model (NumClass=1)", m.NumClass)
	}
	if m.NumFeature != feat.FeatureCount() {
		return nil, fmt.Errorf("topic: model expects %d features but featurizer produces %d", m.NumFeature, feat.FeatureCount())
	}
	return &Classifier{Name: name, Threshold: threshold, feat: feat, model: m}, nil
}

// Score returns P(topic) for text in [0,1].
func (c *Classifier) Score(text string) float64 {
	return c.model.Predict(c.feat.Vectorize(text))[0]
}

// Is reports whether text belongs to the topic (Score >= Threshold).
func (c *Classifier) Is(text string) bool { return c.Score(text) >= c.Threshold }

// ScoreIs returns P(topic) and whether it meets the threshold in a single pass.
// Callers that need both must use this rather than Score + Is, which would
// vectorize and run the tree ensemble twice for the same text.
func (c *Classifier) ScoreIs(text string) (float64, bool) {
	s := c.Score(text)
	return s, s >= c.Threshold
}

// envelope is the on-disk form: the pipeline identity + featurizer shape + the
// embedded grove model. Storing the pipeline lets Load reject an incompatible
// model up front instead of silently scoring on the wrong features.
type envelope struct {
	Pipeline  string          `json:"pipeline"`
	Name      string          `json:"name"`
	Threshold float64         `json:"threshold"`
	Dims      int             `json:"dims"`
	CharN     int             `json:"char_n"`
	WordN     int             `json:"word_n"`
	Model     json.RawMessage `json:"model"`
}

// Save writes the classifier (metadata + grove model) as one JSON document.
func (c *Classifier) Save(w io.Writer) error {
	var mb bytes.Buffer
	if err := c.model.Save(&mb); err != nil {
		return err
	}
	return json.NewEncoder(w).Encode(envelope{
		Pipeline:  gist.PipelineVersion,
		Name:      c.Name,
		Threshold: c.Threshold,
		Dims:      c.feat.Dims,
		CharN:     c.feat.CharN,
		WordN:     c.feat.WordN,
		Model:     mb.Bytes(),
	})
}

// Load reads a classifier written by [Classifier.Save], rejecting a model built
// under a different [gist.PipelineVersion] (the features would be incomparable).
func Load(r io.Reader) (*Classifier, error) {
	var e envelope
	if err := json.NewDecoder(r).Decode(&e); err != nil {
		return nil, err
	}
	if e.Pipeline != gist.PipelineVersion {
		return nil, fmt.Errorf("topic: model pipeline %q != %q — retrain under the current pipeline", e.Pipeline, gist.PipelineVersion)
	}
	m, err := grove.Load(bytes.NewReader(e.Model))
	if err != nil {
		return nil, fmt.Errorf("topic: load model: %w", err)
	}
	return New(e.Name, e.Threshold, gist.Featurizer{Dims: e.Dims, CharN: e.CharN, WordN: e.WordN}, m)
}

// Open loads a classifier from a file path.
func Open(path string) (*Classifier, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return Load(f)
}