Create app.py

app.py

@@ -0,0 +1,151 @@
+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 pickle
+from stable_baselines3 import PPO
+from pydantic import BaseModel
+from typing import Optional
+
+# Define the custom HierarchicalPrediction 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)
+        self.num_subcategories = num_subcategories
+        self.cat_to_subcat_tensor = cat_to_subcat_tensor
+        self.max_subcats_per_cat = max_subcats_per_cat
+        self.subcategory_dense = Dense(num_subcategories, activation=None)
+
+    def build(self, input_shape):
+        super(HierarchicalPrediction, self).build(input_shape)
+
+    def call(self, inputs):
+        lstm_output, category_probs = inputs
+        subcat_logits = self.subcategory_dense(lstm_output)
+        batch_size = tf.shape(category_probs)[0]
+        predicted_categories = tf.argmax(category_probs, axis=1)
+        valid_subcat_indices = tf.gather(self.cat_to_subcat_tensor, predicted_categories)
+        batch_indices = tf.range(batch_size)
+        batch_indices_expanded = tf.tile(batch_indices[:, tf.newaxis], [1, self.max_subcats_per_cat])
+        update_indices = tf.stack([batch_indices_expanded, valid_subcat_indices], axis=-1)
+        update_indices = tf.reshape(update_indices, [-1, 2])
+        valid_mask = tf.not_equal(valid_subcat_indices, -1)
+        valid_indices = tf.boolean_mask(update_indices, tf.reshape(valid_mask, [-1]))
+        updates = tf.ones(tf.shape(valid_indices)[0], dtype=tf.float32)
+        mask = tf.scatter_nd(valid_indices, updates, [batch_size, self.num_subcategories])
+        masked_logits = subcat_logits * mask + (1 - mask) * tf.float32.min
+        return tf.nn.softmax(masked_logits)
+
+    def get_config(self):
+        config = super(HierarchicalPrediction, self).get_config()
+        config.update({
+            'num_subcategories': self.num_subcategories,
+            'cat_to_subcat_tensor': self.cat_to_subcat_tensor.numpy(),
+            'max_subcats_per_cat': self.max_subcats_per_cat
+        })
+        return config
+
+    @classmethod
+    def from_config(cls, config):
+        config['cat_to_subcat_tensor'] = tf.constant(config['cat_to_subcat_tensor'], dtype=tf.int32)
+        return cls(**config)
+
+# Register custom layer
+tf.keras.utils.get_custom_objects()['HierarchicalPrediction'] = HierarchicalPrediction
+
+# Load model and preprocessing objects
+try:
+    model = tf.keras.models.load_model('model.h5', custom_objects={'HierarchicalPrediction': HierarchicalPrediction})
+except Exception as e:
+    raise RuntimeError(f"Failed to load model: {e}")
+
+with open('tokenizer.pkl', 'rb') as f:
+    tokenizer = pickle.load(f)
+with open('le_category.pkl', 'rb') as f:
+    le_category = pickle.load(f)
+with open('le_subcategory.pkl', 'rb') as f:
+    le_subcategory = pickle.load(f)
+
+max_len = 20
+num_subcategories = len(le_subcategory.classes_)
+
+# Load PPO model (optional)
+try:
+    ppo_model = PPO.load('ppo_finetuned_model')
+    rl_available = True
+except FileNotFoundError:
+    ppo_model = None
+    rl_available = False
+
+# Define request model
+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.")
+
+# 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):
+    if not rl_available:
+        raise HTTPException(status_code=400, detail="RL model not available")
+    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 API routes
+@app.get("/")
+async def root():
+    return {"message": "Welcome to the Transaction Classifier API. Use POST /predict to classify transactions."}
+
+@app.post("/predict", response_model=PredictionResponse)
+async def predict(request: TransactionRequest):
+    try:
+        if request.use_rl:
+            cat_pred, subcat_pred, cat_conf, subcat_conf = predict_with_rl(request.description)
+        else:
+            cat_pred, subcat_pred, cat_conf, subcat_conf = predict_category_subcategory(request.description)
+        return {
+            "category": cat_pred,
+            "subcategory": subcat_pred,
+            "category_confidence": cat_conf,
+            "subcategory_confidence": subcat_conf
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Prediction error: {str(e)}")
+
+# Optional: Health check endpoint
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy"}