Update app.py

app.py

@@ -1,14 +1,24 @@
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import Optional
 import tensorflow as tf
 from tensorflow.keras.layers import Dense
 from tensorflow.keras.preprocessing.sequence import pad_sequences
 import numpy as np
 import pickle
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
-from typing import Optional
 from stable_baselines3 import PPO
 
-# Define custom layer
+app = FastAPI(title="Transaction Classifier API", description="API for classifying banking transactions.")
+print("FastAPI app initialized...")
+
+# Lazy-loaded resources
+model = None
+tokenizer = None
+le_category = None
+le_subcategory = None
+ppo_model = None
+max_len = 20
+
 class HierarchicalPrediction(tf.keras.layers.Layer):
     def __init__(self, num_subcategories, cat_to_subcat_tensor, max_subcats_per_cat, **kwargs):
         super(HierarchicalPrediction, self).__init__(**kwargs)
@@ -53,53 +63,37 @@ class HierarchicalPrediction(tf.keras.layers.Layer):
 
 tf.keras.utils.get_custom_objects()['HierarchicalPrediction'] = HierarchicalPrediction
 
-# Initialize FastAPI app at the top
-app = FastAPI(title="Transaction Classifier API", description="API for classifying banking transactions.")
-
-# Load models and resources
-try:
-    print("Loading BiLSTM model...")
-    model = tf.keras.models.load_model('model.h5', custom_objects={'HierarchicalPrediction': HierarchicalPrediction})
-    print("BiLSTM model loaded.")
-except Exception as e:
-    raise RuntimeError(f"Failed to load BiLSTM model: {e}")
-
-try:
-    print("Loading tokenizer...")
-    with open('tokenizer.pkl', 'rb') as f:
-        tokenizer = pickle.load(f)
-    print("Tokenizer loaded.")
-except Exception as e:
-    raise RuntimeError(f"Failed to load tokenizer: {e}")
-
-try:
-    print("Loading category label encoder...")
-    with open('le_category.pkl', 'rb') as f:
-        le_category = pickle.load(f)
-    print("Category label encoder loaded.")
-except Exception as e:
-    raise RuntimeError(f"Failed to load category label encoder: {e}")
-
-try:
-    print("Loading subcategory label encoder...")
-    with open('le_subcategory.pkl', 'rb') as f:
-        le_subcategory = pickle.load(f)
-    print("Subcategory label encoder loaded.")
-except Exception as e:
-    raise RuntimeError(f"Failed to load subcategory label encoder: {e}")
-
-max_len = 20
-num_subcategories = len(le_subcategory.classes_)
-print(f"max_len: {max_len}, num_subcategories: {num_subcategories}")
-
-try:
-    print("Loading PPO model...")
-    ppo_model = PPO.load('ppo_finetuned_model')
-    print("PPO model loaded.")
-except Exception as e:
-    raise RuntimeError(f"Failed to load PPO model: {e}")
+def load_resources():
+    global model, tokenizer, le_category, le_subcategory
+    if model is None:
+        print("Loading BiLSTM model...")
+        model = tf.keras.models.load_model('model.h5', custom_objects={'HierarchicalPrediction': HierarchicalPrediction})
+        print("BiLSTM model loaded.")
+    if tokenizer is None:
+        print("Loading tokenizer...")
+        with open('tokenizer.pkl', 'rb') as f:
+            tokenizer = pickle.load(f)
+        print("Tokenizer loaded.")
+    if le_category is None:
+        print("Loading category label encoder...")
+        with open('le_category.pkl', 'rb') as f:
+            le_category = pickle.load(f)
+        print("Category label encoder loaded.")
+    if le_subcategory is None:
+        print("Loading subcategory label encoder...")
+        with open('le_subcategory.pkl', 'rb') as f:
+            le_subcategory = pickle.load(f)
+        print("Subcategory label encoder loaded.")
+    return model, tokenizer, le_category, le_subcategory
+
+def load_ppo_model():
+    global ppo_model
+    if ppo_model is None:
+        print("Loading PPO model...")
+        ppo_model = PPO.load('ppo_finetuned_model')
+        print("PPO model loaded.")
+    return ppo_model
 
-# Define models
 class TransactionRequest(BaseModel):
     description: str
     use_rl: Optional[bool] = False
@@ -110,36 +104,6 @@ class PredictionResponse(BaseModel):
     category_confidence: float
     subcategory_confidence: float
 
-# Prediction functions
-def predict_category_subcategory(description):
-    seq = tokenizer.texts_to_sequences([description])
-    pad = pad_sequences(seq, maxlen=max_len, padding='post')
-    pred = model.predict(pad, verbose=0)
-    cat_probs = pred[0][0]
-    subcat_probs = pred[1][0]
-    cat_idx = np.argmax(cat_probs)
-    subcat_idx = np.argmax(subcat_probs)
-    cat_pred = le_category.inverse_transform([cat_idx])[0]
-    subcat_pred = le_subcategory.inverse_transform([subcat_idx])[0]
-    cat_conf = float(cat_probs[cat_idx] * 100)
-    subcat_conf = float(subcat_probs[subcat_idx] * 100)
-    return cat_pred, subcat_pred, cat_conf, subcat_conf
-
-def predict_with_rl(description):
-    seq = tokenizer.texts_to_sequences([description])
-    pad = pad_sequences(seq, maxlen=max_len, padding='post')
-    pred = model.predict(pad, verbose=0)
-    obs = pad[0]
-    action, _ = ppo_model.predict(obs)
-    cat_idx = action // num_subcategories
-    subcat_idx = action % num_subcategories
-    cat_pred = le_category.inverse_transform([cat_idx])[0]
-    subcat_pred = le_subcategory.inverse_transform([subcat_idx])[0]
-    cat_conf = float(pred[0][0][cat_idx] * 100)
-    subcat_conf = float(pred[1][0][subcat_idx] * 100)
-    return cat_pred, subcat_pred, cat_conf, subcat_conf
-
-# Define routes
 @app.get("/")
 async def root():
     return {"message": "Welcome to the Transaction Classifier API. Use POST /predict to classify transactions."}
@@ -147,10 +111,30 @@ async def root():
 @app.post("/predict", response_model=PredictionResponse)
 async def predict(request: TransactionRequest):
     try:
+        model, tokenizer, le_category, le_subcategory = load_resources()
+        num_subcategories = len(le_subcategory.classes_)
+        
+        seq = tokenizer.texts_to_sequences([description])
+        pad = pad_sequences(seq, maxlen=max_len, padding='post')
+        pred = model.predict(pad, verbose=0)
+        
         if request.use_rl:
-            cat_pred, subcat_pred, cat_conf, subcat_conf = predict_with_rl(request.description)
+            ppo = load_ppo_model()
+            obs = pad[0]
+            action, _ = ppo.predict(obs)
+            cat_idx = action // num_subcategories
+            subcat_idx = action % num_subcategories
         else:
-            cat_pred, subcat_pred, cat_conf, subcat_conf = predict_category_subcategory(request.description)
+            cat_probs = pred[0][0]
+            subcat_probs = pred[1][0]
+            cat_idx = np.argmax(cat_probs)
+            subcat_idx = np.argmax(subcat_probs)
+        
+        cat_pred = le_category.inverse_transform([cat_idx])[0]
+        subcat_pred = le_subcategory.inverse_transform([subcat_idx])[0]
+        cat_conf = float(pred[0][0][cat_idx] * 100)
+        subcat_conf = float(pred[1][0][subcat_idx] * 100)
+        
         return {
             "category": cat_pred,
             "subcategory": subcat_pred,
@@ -164,4 +148,4 @@ async def predict(request: TransactionRequest):
 async def health_check():
     return {"status": "healthy"}
 
-print("API fully loaded and ready.")
+print("API ready with lazy-loaded resources.")