Update app.py

app.py

@@ -1,30 +1,31 @@
 import gradio as gr
 import torch
 from transformers import AlbertTokenizer, AlbertForSequenceClassification, BertTokenizer, BertForSequenceClassification
+from catboost import CatBoostClassifier
+from huggingface_hub import hf_hub_download
 import numpy as np
-import joblib
-import os
 
-# Load ALBERT model
+# Load ALBERT model and tokenizer from Hugging Face Hub
 albert_model = AlbertForSequenceClassification.from_pretrained("Deepaksai1/albert-fraud-detector")
 albert_tokenizer = AlbertTokenizer.from_pretrained("Deepaksai1/albert-fraud-detector")
 albert_model.eval()
 
-# Load FinBERT model
+# Load FinBERT model and tokenizer from Hugging Face Hub
 finbert_model = BertForSequenceClassification.from_pretrained("Deepaksai1/finbert-fraud-detector")
 finbert_tokenizer = BertTokenizer.from_pretrained("Deepaksai1/finbert-fraud-detector")
 finbert_model.eval()
 
-# Load CatBoost model
-from catboost import CatBoostClassifier
-catboost_model_path = "catboost_fraud_model.cbm"
+# Download and load CatBoost model from Hugging Face Hub
+catboost_model_path = hf_hub_download(repo_id="Deepaksai1/catboost-fraud-detector", filename="catboost_fraud_model.cbm")
 catboost_model = CatBoostClassifier()
 catboost_model.load_model(catboost_model_path)
 
-# CatBoost prediction (expects structured features, here we simulate with dummy value)
+# CatBoost prediction (simple numeric extraction)
 def predict_with_catboost(text):
-    # Simulate with simple heuristic
-    amount = float([s for s in text.split(',') if 'Amount' in s][0].split(':')[1].strip())
+    try:
+        amount = float([s for s in text.split(',') if 'Amount' in s][0].split(':')[1].strip())
+    except:
+        return "Invalid input", 0.0
     prediction = catboost_model.predict([[amount]])[0]
     proba = catboost_model.predict_proba([[amount]])[0][1]
     return ("Fraud" if prediction == 1 else "Not Fraud"), float(proba)
@@ -49,7 +50,7 @@ def predict_with_finbert(text):
         pred_prob = probs[0][1].item()
     return ("Fraud" if pred_class == 1 else "Not Fraud"), float(pred_prob)
 
-# Main prediction selector
+# Model selector
 def predict(text, model_name):
     if model_name == "ALBERT":
         return predict_with_albert(text)
@@ -60,7 +61,7 @@ def predict(text, model_name):
     else:
         return "Unknown Model", 0.0
 
-# Example transactions
+# Examples from dataset
 examples = [
     "Step: 305, Type: CASH_OUT, Amount: 2321633.57, Origin Balance: 2321633.57, Dest Balance: 0.0",
     "Step: 6, Type: CASH_OUT, Amount: 13704.0, Origin Balance: 13704.0, Dest Balance: 3382.84",
@@ -69,7 +70,7 @@ examples = [
     "Step: 372, Type: CASH_IN, Amount: 187503.32, Origin Balance: 76827.0, Dest Balance: 0.0"
 ]
 
-# Gradio interface
+# Gradio Interface
 gui = gr.Interface(
     fn=predict,
     inputs=[
@@ -81,10 +82,10 @@ gui = gr.Interface(
         gr.Number(label="Fraud Probability")
     ],
     examples=[[ex, "ALBERT"] for ex in examples],
-    title="💸 Fraud Detection Assistant (ALBERT, FinBERT, CatBoost)",
-    description="Analyze transaction text for fraud using your choice of model."
+    title="💸 Fraud Detection Assistant",
+    description="Analyze transaction text using ALBERT, FinBERT, or CatBoost models."
 )
 
-# Launch
+# Launch the app
 if __name__ == "__main__":
     gui.launch()