Update app.py

app.py

@@ -1,16 +1,14 @@
-from fastapi import FastAPI, HTTPException
 import tensorflow as tf
-import numpy as np
 from tensorflow.keras.layers import Dense
 from tensorflow.keras.preprocessing.sequence import pad_sequences
+import numpy as np
 import pickle
-from stable_baselines3 import PPO
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from typing import Optional
+from stable_baselines3 import PPO
 
-print("Starting app.py...")
-
-# Define the custom HierarchicalPrediction layer
+# Define custom layer
 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,55 +51,65 @@ class HierarchicalPrediction(tf.keras.layers.Layer):
         config['cat_to_subcat_tensor'] = tf.constant(config['cat_to_subcat_tensor'], dtype=tf.int32)
         return cls(**config)
 
-print("Defined HierarchicalPrediction layer...")
-
-# Register custom layer
 tf.keras.utils.get_custom_objects()['HierarchicalPrediction'] = HierarchicalPrediction
 
-print("Loading BiLSTM model...")
+# 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 model: {e}")
+    raise RuntimeError(f"Failed to load BiLSTM model: {e}")
 
-print("Loading tokenizer...")
-with open('tokenizer.pkl', 'rb') as f:
-    tokenizer = pickle.load(f)
+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}")
 
-print("Loading category label encoder...")
-with open('le_category.pkl', 'rb') as f:
-    le_category = pickle.load(f)
+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}")
 
-print("Loading subcategory label encoder...")
-with open('le_subcategory.pkl', 'rb') as f:
-    le_subcategory = pickle.load(f)
+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}")
 
-print("Loading PPO model...")
 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}")
 
-# Define request model
+# Define models
 class TransactionRequest(BaseModel):
     description: str
     use_rl: Optional[bool] = False
 
-# Define response model
 class PredictionResponse(BaseModel):
     category: str
     subcategory: str
     category_confidence: float
     subcategory_confidence: float
 
-# Initialize FastAPI app
-app = FastAPI(title="Transaction Classifier API", description="API for classifying banking transactions.")
-print("FastAPI app initialized...")
-
 # Prediction functions
 def predict_category_subcategory(description):
     seq = tokenizer.texts_to_sequences([description])
@@ -131,7 +139,7 @@ def predict_with_rl(description):
     subcat_conf = float(pred[1][0][subcat_idx] * 100)
     return cat_pred, subcat_pred, cat_conf, subcat_conf
 
-# Define API routes
+# Define routes
 @app.get("/")
 async def root():
     return {"message": "Welcome to the Transaction Classifier API. Use POST /predict to classify transactions."}
@@ -156,4 +164,4 @@ async def predict(request: TransactionRequest):
 async def health_check():
     return {"status": "healthy"}
 
-print("API routes defined...")
+print("API fully loaded and ready.")