add automatic model retraining from user corrections

sourcecode/backend/categoriser.py

@@ -2,23 +2,153 @@
 from sentence_transformers import SentenceTransformer
 import pickle
 import os
+import pandas as pd
+import psycopg2
+from sklearn.linear_model import LogisticRegression
+from sklearn.model_selection import train_test_split
+
 
 _clf = None
 _st_model = None
 
+_MODEL_PATH = os.path.join(os.path.dirname(__file__), 'cipher_categoriser_v3.pkl')
+_TRAINING_DATA_PATH = os.path.join(os.path.dirname(__file__), 'cipher_training_data_combined.csv')
+
+_last_retrain_result = "not_yet"
+_last_retrain_time = None
+_corrections_at_last_retrain = 0
+_accuracy_before = None
+_accuracy_after = None
+
 def get_model():
     global _clf, _st_model
     if _clf is None:
         try:
-            model_path = os.path.join(os.path.dirname(__file__), 'cipher_categoriser_v3.pkl')
-            with open(model_path, 'rb') as f:
+            with open(_MODEL_PATH, 'rb') as f:
                 _clf = pickle.load(f)
             _st_model = SentenceTransformer('all-MiniLM-L6-v2')
             print("ML model loaded successfully!")
         except Exception as e:
             print(f"Warning: could not load ML model: {e}")
     return _clf, _st_model
 
+def _score(clf, embeddings, labels) -> float:
+    if len(labels) == 0:
+        return 0.0
+    preds = clf.predict(embeddings)
+    return sum(p == t for p, t in zip(preds, labels)) / len(labels) * 100
+
+def retrain_model(db_url: str):
+    global _clf, _last_retrain_result, _last_retrain_time, _corrections_at_last_retrain, _accuracy_before, _accuracy_after
+
+    old_clf, st_model = get_model()
+    if st_model is None:
+        print("Retrain skipped: sentence transformer not loaded.")
+        return
+
+    # ── 1. load corrections from training_queue ──
+    try:
+        conn = psycopg2.connect(db_url)
+        cursor = conn.cursor()
+        cursor.execute("SELECT description, correct_category FROM training_queue WHERE correct_category != ''")
+        rows = cursor.fetchall()
+        conn.close()
+
+        corrections = pd.DataFrame(rows, columns=['description', 'category'])
+        if corrections.empty:
+            print("Retrain skipped: no corrections in training_queue.")
+            return
+
+    except Exception as e:
+        print(f"Retrain failed (DB fetch): {e}")
+        return
+
+    # ── 2. load base CSV and split into train/test ──
+    try:
+        base_df = pd.read_csv(_TRAINING_DATA_PATH)
+    except Exception as e:
+        print(f"Retrain failed (CSV load): {e}")
+        return
+
+    base_train, base_test = train_test_split(base_df, test_size=0.2, random_state=42)
+
+    # ── 3. build training set: 80% base + all corrections ──
+    train_df = pd.concat([base_train, corrections], ignore_index=True).dropna(subset=['description', 'category'])
+
+    # ── 4. embed everything needed ──
+    train_X = st_model.encode(train_df['description'].tolist(), show_progress_bar=False)
+    train_y = train_df['category'].tolist()
+
+    test_old_X = st_model.encode(base_test['description'].tolist(), show_progress_bar=False)
+    test_old_y = base_test['category'].tolist()
+
+    correction_X = st_model.encode(corrections['description'].tolist(), show_progress_bar=False)
+    correction_y = corrections['category'].tolist()
+
+
+    # ── 5. score OLD model on both test sets ──
+    if old_clf is not None:
+        old_on_base = _score(old_clf, test_old_X, test_old_y)
+        old_on_corrections = _score(old_clf, correction_X, correction_y)
+        n_base = len(test_old_y)
+        n_corr = len(correction_y)
+        old_combined = (old_on_base * n_base + old_on_corrections * n_corr) / (n_base + n_corr)
+    else:
+        old_on_base = old_on_corrections = old_combined = None
+
+    # ── 6. train NEW model ──
+    print(f"Retraining on {len(train_df)} samples ({len(corrections)} corrections)...")
+    new_clf = LogisticRegression(max_iter=1000, class_weight='balanced')
+    new_clf.fit(train_X, train_y)
+
+    # ── 7. score NEW model on same two test sets ──
+    new_on_base = _score(new_clf, test_old_X, test_old_y)
+    new_on_corrections = _score(new_clf, correction_X, correction_y)
+    n_base = len(test_old_y)
+    n_corr = len(correction_y)
+    new_combined = (new_on_base * n_base + new_on_corrections * n_corr) / (n_base + n_corr)
+
+    # ── 8. record stats ──
+    from datetime import datetime, timezone
+    _last_retrain_time = datetime.now(timezone.utc).isoformat()
+    _corrections_at_last_retrain = len(corrections)
+    _accuracy_before = {
+        "old_data": round(old_on_base, 1) if old_on_base is not None else None,
+        "corrections": round(old_on_corrections, 1) if old_on_corrections is not None else None,
+        "combined": round(old_combined, 1) if old_combined is not None else None,
+    }
+    _accuracy_after = {
+        "old_data": round(new_on_base, 1),
+        "corrections": round(new_on_corrections, 1),
+        "combined": round(new_combined, 1),
+    }
+
+    # ── 9. swap or roll back ──
+    if old_combined is None or new_combined > old_combined:
+        with open(_MODEL_PATH, 'wb') as f:
+            pickle.dump(new_clf, f)
+        _clf = new_clf
+        _last_retrain_result = "improved"
+        old_str = f"{old_combined:.1f}%" if old_combined is not None else "no baseline"
+        print(f"Model improved: combined {old_str} → {new_combined:.1f}%. Hot-swapped.")
+    else:
+        _last_retrain_result = "rolled_back"
+        print(f"No improvement: combined {old_combined:.1f}% → {new_combined:.1f}%. Keeping old model.")
+
+def get_model_stats() -> dict:
+    return {
+        "last_retrain_result": _last_retrain_result,
+        "last_retrain_time": _last_retrain_time,
+        "corrections_at_last_retrain": _corrections_at_last_retrain,
+        "accuracy_before": _accuracy_before,
+        "accuracy_after": _accuracy_after,
+        "improvement": (
+            round(_accuracy_after["combined"] - _accuracy_before["combined"], 1)
+            if _accuracy_after and _accuracy_before and _accuracy_before["combined"] is not None
+            else None
+        ),
+    }
+
 # Rule-based transaction categoriser for Singapore merchants -- fallback if ML model is funky
 RULES = {
     # check these first — more specific categories before general ones

sourcecode/backend/main.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, File, UploadFile, Depends, HTTPException, Header
+from fastapi import FastAPI, File, UploadFile, Depends, HTTPException, Header, BackgroundTasks
 from fastapi.responses import StreamingResponse
 from typing import List
 from fastapi.middleware.cors import CORSMiddleware
@@ -9,7 +9,7 @@
 import io
 import json
 from pydantic import BaseModel, EmailStr
-from categoriser import categorise_transactions
+from categoriser import categorise_transactions, retrain_model
 from features import extract_features 
 from archetypes import assign_archetype, generate_insights
 from auth import hash_password, verify_password, create_token, decode_token
@@ -574,7 +574,7 @@ def delete_transaction(tx_id: str, user_id: str = Depends(get_current_user)):
         db.close()
 
 @app.put("/transactions/{tx_id}")
-def update_transaction(tx_id: str, request: dict, user_id: str = Depends(get_current_user)):
+def update_transaction(tx_id: str, request: dict, background_tasks: BackgroundTasks, user_id: str = Depends(get_current_user)):
     db = SessionLocal()
     try:
         # get original category before updating
@@ -624,6 +624,9 @@ def update_transaction(tx_id: str, request: dict, user_id: str = Depends(get_cur
                 "correct_category": new_category,
                 "original_category": original_category
             })
+            count = db.execute(text("SELECT COUNT(*) FROM training_queue")).scalar()
+            if count % 50 == 0:
+                background_tasks.add_task(retrain_model, os.getenv("DATABASE_URL"))
 
         db.commit()
         return {"success": True}
@@ -724,3 +727,15 @@ async def stream():
         yield f"data: {json.dumps({'type': 'done'})}\n\n"
 
     return StreamingResponse(stream(), media_type="text/event-stream")
+
+@app.get("/model-stats")
+def model_stats(user_id: str = Depends(get_current_user)):
+    from categoriser import get_model_stats
+    db = SessionLocal()
+    try:
+        total_corrections = db.execute(text("SELECT COUNT(*) FROM training_queue")).scalar()
+        stats = get_model_stats()
+        return {**stats, "total_corrections": total_corrections}
+    finally:
+        db.close()
+