From d9b778eb9e370fcdb4cec10924712842c542392f Mon Sep 17 00:00:00 2001 From: Tevyn Tan Date: Thu, 9 Jul 2026 19:46:59 +0800 Subject: [PATCH] add automatic model retraining from user corrections --- sourcecode/backend/categoriser.py | 134 +++++++++++++++++++++++++++++- sourcecode/backend/main.py | 21 ++++- 2 files changed, 150 insertions(+), 5 deletions(-) diff --git a/sourcecode/backend/categoriser.py b/sourcecode/backend/categoriser.py index cb8d519..7bca4eb 100644 --- a/sourcecode/backend/categoriser.py +++ b/sourcecode/backend/categoriser.py @@ -2,16 +2,29 @@ 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!") @@ -19,6 +32,123 @@ def get_model(): 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 diff --git a/sourcecode/backend/main.py b/sourcecode/backend/main.py index 2769328..98adf4e 100644 --- a/sourcecode/backend/main.py +++ b/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() +