File size: 1,722 Bytes
25d2492 13cb8f6 25d2492 13cb8f6 25d2492 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | # app.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Literal
import joblib
import numpy as np
app = FastAPI(title="Transaction Categorizer")
# Load models once at startup
shop_model = joblib.load("shop_classifier.pkl")
category_model = joblib.load("category_classifier.pkl")
def predict_or_na(pipeline, text: str, threshold: float = 0.3) -> str:
"""
Returns the top class if its probability ≥ threshold, else "N/A"
"""
probs = pipeline.predict_proba([text])[0]
top_idx = int(np.argmax(probs))
return str(pipeline.classes_[top_idx]) if probs[top_idx] >= threshold else "N/A"
class TransactionIn(BaseModel):
id: str
description: str
class TransactionOut(BaseModel):
id: str
description: str
shop: str
category: str
@app.post("/predict", response_model=List[TransactionOut])
def predict_transactions(
items: List[TransactionIn],
threshold: float = 0.3
):
"""
Predict shop and category for each transaction.
- **items**: list of `{ id, description }`
- **threshold**: optional override for probability cutoff
"""
results = []
for item in items:
desc_norm = item.description.lower().strip()
shop_pred = predict_or_na(shop_model, desc_norm, threshold)
category_pred = predict_or_na(category_model, desc_norm, threshold)
results.append(
TransactionOut(
id=item.id,
description=item.description,
shop=shop_pred,
category=category_pred
)
)
return results
# Optional healthcheck
@app.get("/health")
def health():
return {"status": "ok"} |