From 96ab2fcb2ff442698465389d75390afa91629165 Mon Sep 17 00:00:00 2001 From: Sam Scholten Date: Sun, 30 Aug 2026 12:50:24 +1000 Subject: fluxrec: personal Miniflux article recommender --- cmds/train.go | 430 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 430 insertions(+) create mode 100644 cmds/train.go (limited to 'cmds/train.go') 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"` +} -- cgit v1.2.3