feat: add metrics

model/src/bin/main.rs

@@ -4,6 +4,7 @@ use clap::{Parser, Subcommand};
 use model::Model;
 
 pub const MODEL_PATH: &str = "assets/model.json";
+pub const METRICS_PATH: &str = "assets/metrics.json";
 pub const TRAIN_DATASET_PATH: &str = "assets/train.csv";
 pub const TEST_DATASET_PATH: &str = "assets/test.csv";
 
@@ -56,6 +57,7 @@ fn predict_command(text: &str) -> anyhow::Result<()> {
 fn train_command() -> anyhow::Result<()> {
     Model::train(
         MODEL_PATH.into(),
+        METRICS_PATH.into(),
         TRAIN_DATASET_PATH.into(),
         TEST_DATASET_PATH.into(),
     )?;

model/src/lib.rs

@@ -1,3 +1,4 @@
 pub mod model;
 
 pub use model::Model;
+pub use model::metrics::Metrics;

model/src/model.rs

@@ -6,11 +6,13 @@ use ndarray::Array1;
 use serde::{Deserialize, Serialize};
 
 use crate::model::logistic_regression::LogisticRegression;
+use crate::model::metrics::Metrics;
 use crate::model::tfidf::TfIdf;
 
 #[cfg(feature = "training")]
 mod gpu;
 pub mod logistic_regression;
+pub mod metrics;
 pub mod tfidf;
 
 /// A dataset record
@@ -109,6 +111,7 @@ impl Model {
 
     pub fn train(
         model_path: PathBuf,
+        metrics_path: PathBuf,
         train_dataset: PathBuf,
         test_dataset: PathBuf,
     ) -> anyhow::Result<()> {
@@ -130,19 +133,18 @@ impl Model {
         println!("Training...");
         let model = Self::fit_logreg(&x_train, &y_train, vectorizer.vocab_size())?;
 
-        let predictions = model.predict(&x_test);
-        let correct = predictions
-            .iter()
-            .zip(y_test.iter())
-            .filter(|(p, y)| **p == (**y >= 0.5))
-            .count();
-        let accuracy = correct as f64 / y_test.len() as f64;
+        let predictions: Vec<bool> = model.predict(&x_test).iter().copied().collect();
+        let labels: Vec<f64> = y_test.iter().copied().collect();
+        let metrics = Metrics::compute(&predictions, &labels);
 
         println!(
-            "\n Test accuracy: {:.2}% ({}/{})",
-            accuracy * 100.0,
-            correct,
-            y_test.len()
+            "\nTest metrics: accuracy={:.2}% precision={:.2}% recall={:.2}% f1={:.2}% ({}/{})",
+            metrics.accuracy * 100.0,
+            metrics.precision * 100.0,
+            metrics.recall * 100.0,
+            metrics.f1 * 100.0,
+            metrics.correct,
+            metrics.total
         );
 
         // Try it on a couple of brand-new sentences.
@@ -165,6 +167,9 @@ impl Model {
         model.save(model_path.clone())?;
         println!("\nSaved trained model to {}", model_path.to_string_lossy());
 
+        metrics.save(metrics_path.clone())?;
+        println!("Saved metrics to {}", metrics_path.to_string_lossy());
+
         Ok(())
     }
 }

model/src/model/metrics.rs

@@ -0,0 +1,111 @@
+use std::fs::File;
+use std::path::PathBuf;
+
+use anyhow::{Context, Result};
+use serde::{Deserialize, Serialize};
+
+/// Test-set evaluation metrics for a trained [`crate::Model`].
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Metrics {
+    pub accuracy: f64,
+    pub precision: f64,
+    pub recall: f64,
+    pub f1: f64,
+    pub correct: usize,
+    pub total: usize,
+}
+
+impl Metrics {
+    /// Computes accuracy/precision/recall/F1 from predicted vs. actual labels.
+    /// "positive" is the `true`/`>= 0.5` class.
+    pub fn compute(predictions: &[bool], labels: &[f64]) -> Self {
+        let mut true_positives = 0usize;
+        let mut false_positives = 0usize;
+        let mut true_negatives = 0usize;
+        let mut false_negatives = 0usize;
+
+        for (&predicted, &label) in predictions.iter().zip(labels.iter()) {
+            let actual = label >= 0.5;
+            match (predicted, actual) {
+                (true, true) => true_positives += 1,
+                (true, false) => false_positives += 1,
+                (false, false) => true_negatives += 1,
+                (false, true) => false_negatives += 1,
+            }
+        }
+
+        let total = predictions.len();
+        let correct = true_positives + true_negatives;
+        let accuracy = correct as f64 / total as f64;
+
+        let precision = if true_positives + false_positives == 0 {
+            0.0
+        } else {
+            true_positives as f64 / (true_positives + false_positives) as f64
+        };
+
+        let recall = if true_positives + false_negatives == 0 {
+            0.0
+        } else {
+            true_positives as f64 / (true_positives + false_negatives) as f64
+        };
+
+        let f1 = if precision + recall == 0.0 {
+            0.0
+        } else {
+            2.0 * precision * recall / (precision + recall)
+        };
+
+        Self {
+            accuracy,
+            precision,
+            recall,
+            f1,
+            correct,
+            total,
+        }
+    }
+
+    pub fn save(&self, path: PathBuf) -> Result<()> {
+        let file = File::create(path)?;
+        serde_json::to_writer_pretty(file, self).context("Failed to save metrics")
+    }
+
+    pub fn load(path: PathBuf) -> Result<Self> {
+        let file = File::open(path)?;
+        serde_json::from_reader(file).context("Failed to load metrics")
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::Metrics;
+
+    #[test]
+    fn computes_confusion_matrix_metrics() {
+        // 2 TP, 1 FP, 1 TN, 1 FN
+        let predictions = [true, true, true, false, false];
+        let labels = [1.0, 1.0, 0.0, 0.0, 1.0];
+
+        let metrics = Metrics::compute(&predictions, &labels);
+
+        assert_eq!(metrics.total, 5);
+        assert_eq!(metrics.correct, 3);
+        assert!((metrics.accuracy - 0.6).abs() < 1e-12);
+        assert!((metrics.precision - 2.0 / 3.0).abs() < 1e-12);
+        assert!((metrics.recall - 2.0 / 3.0).abs() < 1e-12);
+        assert!((metrics.f1 - 2.0 / 3.0).abs() < 1e-12);
+    }
+
+    #[test]
+    fn handles_no_positive_predictions() {
+        let predictions = [false, false];
+        let labels = [1.0, 0.0];
+
+        let metrics = Metrics::compute(&predictions, &labels);
+
+        assert_eq!(metrics.precision, 0.0);
+        assert_eq!(metrics.recall, 0.0);
+        assert_eq!(metrics.f1, 0.0);
+    }
+}

server/src/main.rs

@@ -8,9 +8,15 @@ use axum::routing::{get, post};
 use axum::{Form, Router};
 use serde::Deserialize;
 
-use model::Model;
+use model::{Metrics, Model};
 
 pub const MODEL_PATH: &str = "assets/model.json";
+pub const METRICS_PATH: &str = "assets/metrics.json";
+
+struct AppState {
+    model: Model,
+    metrics: Option<MetricsView>,
+}
 
 #[derive(Clone)]
 struct PredictionView {
@@ -18,11 +24,31 @@ struct PredictionView {
     prob: String,
 }
 
+#[derive(Clone)]
+struct MetricsView {
+    accuracy: String,
+    precision: String,
+    recall: String,
+    f1: String,
+}
+
+impl From<Metrics> for MetricsView {
+    fn from(metrics: Metrics) -> Self {
+        Self {
+            accuracy: format!("{:.2}", metrics.accuracy * 100.0),
+            precision: format!("{:.2}", metrics.precision * 100.0),
+            recall: format!("{:.2}", metrics.recall * 100.0),
+            f1: format!("{:.2}", metrics.f1 * 100.0),
+        }
+    }
+}
+
 #[derive(Template)]
 #[template(path = "index.html")]
 struct IndexTemplate {
     text: String,
     result: Option<PredictionView>,
+    metrics: Option<MetricsView>,
 }
 
 #[derive(Deserialize)]
@@ -43,15 +69,19 @@ impl IntoResponse for IndexTemplate {
     }
 }
 
-async fn index() -> IndexTemplate {
+async fn index(State(state): State<Arc<AppState>>) -> IndexTemplate {
     IndexTemplate {
         text: String::new(),
         result: None,
+        metrics: state.metrics.clone(),
     }
 }
 
-async fn predict(State(model): State<Arc<Model>>, Form(form): Form<PredictForm>) -> IndexTemplate {
-    let (label, prob) = model.predict(&form.text);
+async fn predict(
+    State(state): State<Arc<AppState>>,
+    Form(form): Form<PredictForm>,
+) -> IndexTemplate {
+    let (label, prob) = state.model.predict(&form.text);
 
     IndexTemplate {
         text: form.text,
@@ -59,6 +89,7 @@ async fn predict(State(model): State<Arc<Model>>, Form(form): Form<PredictForm>)
             label,
             prob: format!("{prob:.2}"),
         }),
+        metrics: state.metrics.clone(),
     }
 }
 
@@ -66,7 +97,8 @@ async fn predict(State(model): State<Arc<Model>>, Form(form): Form<PredictForm>)
 async fn main() -> anyhow::Result<()> {
     let model = Model::load(MODEL_PATH.into())
         .context("couldn't load {MODEL_PATH} ({e}) - run `cargo run --release -- train` first")?;
-    let state = Arc::new(model);
+    let metrics = Metrics::load(METRICS_PATH.into()).ok().map(MetricsView::from);
+    let state = Arc::new(AppState { model, metrics });
 
     let app = Router::new()
         .route("/", get(index))

server/templates/index.html

@@ -45,6 +45,18 @@
   .result.positive { background: #dcfce7; color: #14532d; }
   .result.negative { background: #fee2e2; color: #7f1d1d; }
   .prob { font-variant-numeric: tabular-nums; opacity: 0.8; }
+  .metrics {
+    margin-top: 2rem;
+    padding-top: 1rem;
+    border-top: 1px solid #ddd;
+    display: grid;
+    grid-template-columns: repeat(4, 1fr);
+    gap: 0.5rem;
+    text-align: center;
+  }
+  .metrics .stat { display: flex; flex-direction: column; }
+  .metrics .value { font-size: 1.1rem; font-variant-numeric: tabular-nums; }
+  .metrics .label { font-size: 0.75rem; color: #555; text-transform: uppercase; }
 </style>
 </head>
 <body>
@@ -60,5 +72,25 @@
     <span class="prob">(p={{ result.prob }})</span>
   </div>
   {% endif %}
+  {% if let Some(metrics) = metrics %}
+  <div class="metrics">
+    <div class="stat">
+      <span class="value">{{ metrics.accuracy }}%</span>
+      <span class="label">Accuracy</span>
+    </div>
+    <div class="stat">
+      <span class="value">{{ metrics.precision }}%</span>
+      <span class="label">Precision</span>
+    </div>
+    <div class="stat">
+      <span class="value">{{ metrics.recall }}%</span>
+      <span class="label">Recall</span>
+    </div>
+    <div class="stat">
+      <span class="value">{{ metrics.f1 }}%</span>
+      <span class="label">F1</span>
+    </div>
+  </div>
+  {% endif %}
 </body>
 </html>