fluxrec: personal Miniflux article recommender

1 files changed, 430 insertions, 0 deletions

diff --git a/cmds/train.go b/cmds/train.go
new file mode 100644
index 0000000..234ef06
--- /dev/null
+++ b/cmds/train.go
@@ -0,0 +1,430 @@
+// Train command: labels.jsonl → model.json + report.json.
+//
+// Time-split validation (train on older labels, evaluate on the newest
+// window — what serve actually faces). Headline metric is precision@15
+// (rank-based: of the 15 highest-scored validation entries, what fraction
+// are positives). Youden's J is computed and reported for information only;
+// it is NOT used as a cutoff — serve ranks by raw score.
+//
+// Reproducible: a fixed seed makes the downsampling deterministic, so the
+// same labels produce the same weights/vocab. IDF is fit on the full
+// (unbalanced) train set to preserve document frequencies, then the
+// majority class is downsampled to 1:1 before LR fit (matches the vendored
+// scholscan approach).
+package cmds
+
+import (
+	"encoding/json"
+	"flag"
+	"fmt"
+	"io"
+	"math"
+	"math/rand"
+	"os"
+	"sort"
+	"time"
+
+	"fluxrec/core"
+)
+
+// trainSeed fixes the downsampling RNG: same labels + same seed → same model.
+const trainSeed int64 = 20240301
+
+type TrainCommand struct {
+	Labels    string
+	ModelOut  string
+	ReportOut string
+	ValFrac   float64
+	Lambda    float64
+	MinDF     int
+	MaxDF     float64
+	NgramMax  int
+}
+
+func (c *TrainCommand) Name() string { return "train" }
+
+func (c *TrainCommand) Init(args []string) error {
+	fs := flag.NewFlagSet(c.Name(), flag.ContinueOnError)
+	fs.Usage = func() {
+		fmt.Fprint(fs.Output(), `Usage: fluxrec train LABELS_JSONL [flags]
+
+Train a TF-IDF + logistic regression model from labels.jsonl. Writes the
+model to --model-out and a validation report (precision@15 etc.) to
+--report-out. Time-split: the newest --val-frac of labels (by captured_at)
+are held out for validation; the rest train the model.
+`)
+		fs.PrintDefaults()
+	}
+	fs.StringVar(&c.ModelOut, "model-out", "model.json", "Path to write the trained model JSON")
+	fs.StringVar(&c.ReportOut, "report-out", "report.json", "Path to write the validation report JSON")
+	fs.Float64Var(&c.ValFrac, "val-frac", 0.2, "Fraction of newest labels held out for validation")
+	fs.Float64Var(&c.Lambda, "lambda", 0.001, "L2 regularization strength")
+	fs.IntVar(&c.MinDF, "min-df", 2, "Minimum document frequency (absolute)")
+	fs.Float64Var(&c.MaxDF, "max-df", 0.8, "Maximum document frequency (ratio)")
+	fs.IntVar(&c.NgramMax, "ngram-max", 2, "Max n-gram size (1=unigrams, 2=+bigrams)")
+
+	for _, a := range args {
+		if a == "--help" || a == "-h" {
+			fs.Usage()
+			return flag.ErrHelp
+		}
+	}
+	if len(args) == 0 {
+		return fmt.Errorf("labels.jsonl path argument is required")
+	}
+	c.Labels = args[0]
+	if err := fs.Parse(args[1:]); err != nil {
+		return err
+	}
+	if c.ValFrac <= 0 || c.ValFrac >= 1 {
+		return fmt.Errorf("--val-frac must be between 0 and 1, got %v", c.ValFrac)
+	}
+	return nil
+}
+
+func (c *TrainCommand) Run(stdin io.Reader, stdout io.Writer) error {
+	rows, err := core.ReadLabels(c.Labels)
+	if err != nil {
+		return err
+	}
+	if len(rows) == 0 {
+		return fmt.Errorf("no labels in %s", c.Labels)
+	}
+
+	trainRows, valRows := core.SplitByTime(rows, c.ValFrac)
+	if len(trainRows) == 0 || len(valRows) == 0 {
+		return fmt.Errorf("time split produced an empty side (train=%d val=%d); adjust --val-frac or accumulate more labels",
+			len(trainRows), len(valRows))
+	}
+	valPos := 0
+	for _, r := range valRows {
+		if r.Label == core.LabelPos {
+			valPos++
+		}
+	}
+	if valPos < 5 {
+		return fmt.Errorf("only %d positives in validation (need >=5 for signal, not noise); lower --val-frac or accumulate more labels",
+			valPos)
+	}
+
+	model, report, err := c.trainAndEvaluate(trainRows, valRows, time.Now().UTC())
+	if err != nil {
+		return err
+	}
+
+	if err := writeJSON(c.ModelOut, model); err != nil {
+		return err
+	}
+	if err := writeJSON(c.ReportOut, report); err != nil {
+		return err
+	}
+
+	fmt.Fprintf(stdout, "train: pos=%d neg=%d | val: pos=%d neg=%d | precision@15=%.3f youden_j=%.3f | vocab=%d → %s\n",
+		report.TrainPos, report.TrainNeg, report.ValPos, report.ValNeg,
+		report.PrecisionAt15, report.YoudensJ, report.VocabSize, c.ModelOut)
+	return nil
+}
+
+func writeJSON(path string, v any) error {
+	body, err := json.MarshalIndent(v, "", "  ")
+	if err != nil {
+		return err
+	}
+	if err := os.WriteFile(path, body, 0o644); err != nil {
+		return fmt.Errorf("writing %s: %w", path, err)
+	}
+	return nil
+}
+
+// trainAndEvaluate is the pure training core: given split rows and a clock
+// value, returns the model envelope and validation report. Pure w.r.t. the
+// inputs (no I/O, no wall clock) so two calls with the same rows + now are
+// byte-identical — the reproducibility guarantee.
+func (c *TrainCommand) trainAndEvaluate(trainRows, valRows []core.LabelRow, now time.Time) (*core.ModelEnvelope, *trainReport, error) {
+	trainDocs := make([]string, len(trainRows))
+	trainLabels := make([]float64, len(trainRows))
+	trainPos, trainNeg := 0, 0
+	for i, r := range trainRows {
+		trainDocs[i] = core.ScoringText(r.Title, r.FeedTitle)
+		if r.Label == core.LabelPos {
+			trainLabels[i] = 1.0
+			trainPos++
+		} else {
+			trainLabels[i] = 0.0
+			trainNeg++
+		}
+	}
+	valDocs := make([]string, len(valRows))
+	valLabels := make([]float64, len(valRows))
+	valPos, valNeg := 0, 0
+	for i, r := range valRows {
+		valDocs[i] = core.ScoringText(r.Title, r.FeedTitle)
+		if r.Label == core.LabelPos {
+			valLabels[i] = 1.0
+			valPos++
+		} else {
+			valLabels[i] = 0.0
+			valNeg++
+		}
+	}
+
+	// Fit IDF on the full (unbalanced) train set to preserve document
+	// frequencies, then downsample the majority to 1:1 before LR fit.
+	vectorizer := &core.TFIDFVectorizer{
+		NgramMin: 1, NgramMax: c.NgramMax,
+		MinDF: c.MinDF, MaxDF: c.MaxDF, VocabCap: 50000,
+		Vocabulary: make(map[string]float64),
+	}
+	vectorizer.Fit(trainDocs)
+
+	balDocs, balLabels := downsampleBalance(trainDocs, trainLabels, trainSeed)
+	trainVectors := vectorizer.Transform(balDocs)
+	valVectors := vectorizer.Transform(valDocs)
+
+	lr := &core.LogisticRegression{
+		LearningRate: 0.5, Lambda: c.Lambda, Iterations: 500, Tolerance: 0.000001,
+	}
+	lr.Validate()
+	weights, err := lr.Fit(trainVectors, balLabels, map[float64]float64{1.0: 1.0, 0.0: 1.0})
+	if err != nil {
+		return nil, nil, fmt.Errorf("training logistic regression: %w", err)
+	}
+
+	precisionAt15, youdenJ, posStats, negStats := evaluate(valVectors, valLabels, weights)
+
+	minCap, maxCap := trainRows[0].CapturedAt, trainRows[0].CapturedAt
+	for _, r := range trainRows {
+		if r.CapturedAt.Before(minCap) {
+			minCap = r.CapturedAt
+		}
+		if r.CapturedAt.After(maxCap) {
+			maxCap = r.CapturedAt
+		}
+	}
+
+	vecParams := map[string]any{
+		"ngram_min": vectorizer.NgramMin, "ngram_max": vectorizer.NgramMax,
+		"min_df": vectorizer.MinDF, "max_df": vectorizer.MaxDF, "vocab_cap": vectorizer.VocabCap,
+	}
+	modelParams := map[string]any{
+		"learning_rate": lr.LearningRate, "lambda": lr.Lambda,
+		"iterations": lr.Iterations, "tolerance": lr.Tolerance,
+	}
+
+	model := &core.ModelEnvelope{
+		Algorithm: "tfidf-go", Impl: "go", Version: "1",
+		CreatedAt: now,
+		Meta: map[string]any{
+			"scorer":            "tfidf",
+			"seed":              trainSeed,
+			"train_size":        len(trainRows),
+			"val_size":          len(valRows),
+			"train_pos":         trainPos,
+			"train_neg":         trainNeg,
+			"train_window_min":  minCap.UTC().Format(time.RFC3339),
+			"train_window_max":  maxCap.UTC().Format(time.RFC3339),
+			"val_frac":          c.ValFrac,
+			"precision_at_15":   precisionAt15,
+			"youdens_j":         youdenJ,
+			"vectorizer_params": vecParams,
+			"model_params":      modelParams,
+		},
+		Vectorizer:   vectorizer.Vocabulary,
+		OrderedVocab: vectorizer.OrderedVocab,
+		Weights:      weights,
+	}
+
+	report := &trainReport{
+		Scorer: "tfidf", Seed: trainSeed,
+		TrainSize: len(trainRows), ValSize: len(valRows),
+		TrainPos: trainPos, TrainNeg: trainNeg,
+		ValPos: valPos, ValNeg: valNeg,
+		TrainWindowMin: minCap.UTC(), TrainWindowMax: maxCap.UTC(),
+		ValFrac: c.ValFrac, VocabSize: len(vectorizer.OrderedVocab),
+		PrecisionAt15: precisionAt15, YoudensJ: youdenJ,
+		PositiveScores: posStats, NegativeScores: negStats,
+		VectorizerParams: vecParams, ModelParams: modelParams,
+	}
+	return model, report, nil
+}
+
+// downsampleBalance downsamples the majority class to a 1:1 ratio with the
+// minority, using a seeded RNG for reproducibility. If already balanced,
+// returns the inputs unchanged.
+func downsampleBalance(docs []string, labels []float64, seed int64) ([]string, []float64) {
+	var posDocs, negDocs []string
+	var posLabels, negLabels []float64
+	for i, label := range labels {
+		if label == 1.0 {
+			posDocs = append(posDocs, docs[i])
+			posLabels = append(posLabels, label)
+		} else {
+			negDocs = append(negDocs, docs[i])
+			negLabels = append(negLabels, label)
+		}
+	}
+	if len(posDocs) == len(negDocs) {
+		return docs, labels
+	}
+	var majorityDocs, minorityDocs []string
+	var majorityLabels, minorityLabels []float64
+	if len(negDocs) > len(posDocs) {
+		majorityDocs, minorityDocs = negDocs, posDocs
+		majorityLabels, minorityLabels = negLabels, posLabels
+	} else {
+		majorityDocs, minorityDocs = posDocs, negDocs
+		majorityLabels, minorityLabels = posLabels, negLabels
+	}
+	minoritySize := len(minorityDocs)
+	rng := rand.New(rand.NewSource(seed))
+	indices := make([]int, len(majorityDocs))
+	for i := range indices {
+		indices[i] = i
+	}
+	rng.Shuffle(len(indices), func(i, j int) {
+		indices[i], indices[j] = indices[j], indices[i]
+	})
+	outDocs := make([]string, 0, minoritySize*2)
+	outLabels := make([]float64, 0, minoritySize*2)
+	outDocs = append(outDocs, minorityDocs...)
+	outLabels = append(outLabels, minorityLabels...)
+	for i := 0; i < minoritySize; i++ {
+		outDocs = append(outDocs, majorityDocs[indices[i]])
+		outLabels = append(outLabels, majorityLabels[indices[i]])
+	}
+	return outDocs, outLabels
+}
+
+// evaluate scores the validation set and returns precision@15 (headline),
+// Youden's J (information only), and per-class score stats.
+func evaluate(valVectors [][]float64, valLabels []float64, weights []float64) (precisionAt15, youdenJ float64, posStats, negStats map[string]float64) {
+	type scored struct {
+		score float64
+		label float64
+		idx   int
+	}
+	ss := make([]scored, len(valVectors))
+	for i, v := range valVectors {
+		s, err := core.PredictScore(v, weights)
+		if err != nil {
+			s = 0
+		}
+		ss[i] = scored{s, valLabels[i], i}
+	}
+
+	// precision@15: rank by score desc (ties by idx), take top 15, count positives.
+	ranked := make([]scored, len(ss))
+	copy(ranked, ss)
+	sort.SliceStable(ranked, func(i, j int) bool {
+		if ranked[i].score == ranked[j].score {
+			return ranked[i].idx < ranked[j].idx
+		}
+		return ranked[i].score > ranked[j].score
+	})
+	k := 15
+	if k > len(ranked) {
+		k = len(ranked)
+	}
+	topPos := 0
+	for i := 0; i < k; i++ {
+		if ranked[i].label == 1.0 {
+			topPos++
+		}
+	}
+	if k > 0 {
+		precisionAt15 = float64(topPos) / float64(k)
+	}
+
+	// Youden's J (information only): best sensitivity + specificity - 1 over
+	// a threshold sweep. Not used as a cutoff anywhere.
+	bestJ := -1.0
+	for t := 5; t <= 95; t++ {
+		threshold := float64(t) / 100.0
+		var tp, fp, tn, fn int
+		for _, s := range ss {
+			pred := s.score > threshold
+			switch {
+			case pred && s.label == 1.0:
+				tp++
+			case pred && s.label == 0.0:
+				fp++
+			case !pred && s.label == 1.0:
+				fn++
+			default:
+				tn++
+			}
+		}
+		sens := 0.0
+		if tp+fn > 0 {
+			sens = float64(tp) / float64(tp+fn)
+		}
+		spec := 0.0
+		if tn+fp > 0 {
+			spec = float64(tn) / float64(tn+fp)
+		}
+		if j := sens + spec - 1.0; j > bestJ {
+			bestJ = j
+		}
+	}
+	youdenJ = bestJ
+
+	var pos, neg []float64
+	for _, s := range ss {
+		if s.label == 1.0 {
+			pos = append(pos, s.score)
+		} else {
+			neg = append(neg, s.score)
+		}
+	}
+	posStats = scoreStats(pos)
+	negStats = scoreStats(neg)
+	return
+}
+
+// scoreStats returns min/max/mean/std for a slice of scores.
+func scoreStats(scores []float64) map[string]float64 {
+	if len(scores) == 0 {
+		return map[string]float64{"min": 0, "max": 0, "mean": 0, "std": 0}
+	}
+	min, max, sum := scores[0], scores[0], 0.0
+	for _, s := range scores {
+		if s < min {
+			min = s
+		}
+		if s > max {
+			max = s
+		}
+		sum += s
+	}
+	mean := sum / float64(len(scores))
+	variance := 0.0
+	for _, s := range scores {
+		variance += (s - mean) * (s - mean)
+	}
+	variance /= float64(len(scores))
+	return map[string]float64{
+		"min": min, "max": max, "mean": mean, "std": math.Sqrt(variance),
+	}
+}
+
+// trainReport is the validation report written to --report-out.
+type trainReport struct {
+	Scorer           string             `json:"scorer"`
+	Seed             int64              `json:"seed"`
+	TrainSize        int                `json:"train_size"`
+	ValSize          int                `json:"val_size"`
+	TrainPos         int                `json:"train_pos"`
+	TrainNeg         int                `json:"train_neg"`
+	ValPos           int                `json:"val_pos"`
+	ValNeg           int                `json:"val_neg"`
+	TrainWindowMin   time.Time          `json:"train_window_min"`
+	TrainWindowMax   time.Time          `json:"train_window_max"`
+	ValFrac          float64            `json:"val_frac"`
+	VocabSize        int                `json:"vocab_size"`
+	PrecisionAt15    float64            `json:"precision_at_15"`
+	YoudensJ         float64            `json:"youdens_j"`
+	PositiveScores   map[string]float64 `json:"positive_scores"`
+	NegativeScores   map[string]float64 `json:"negative_scores"`
+	VectorizerParams map[string]any     `json:"vectorizer_params"`
+	ModelParams      map[string]any     `json:"model_params"`
+}