Update app.py

app.py

@@ -1,38 +1,22 @@
 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
 
 # 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 = AlbertForSequenceClassification.from_pretrained("Deepaksai1/albert-fraud-detector-v2")
+albert_tokenizer = AlbertTokenizer.from_pretrained("Deepaksai1/albert-fraud-detector-v2")
 albert_model.eval()
 
 # 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 = BertForSequenceClassification.from_pretrained("Deepaksai1/finbert-fraud-detector-v2")
+finbert_tokenizer = BertTokenizer.from_pretrained("Deepaksai1/finbert-fraud-detector-v2")
 finbert_model.eval()
 
-# 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 (simple numeric extraction)
-def predict_with_catboost(text):
-    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)
-
 # ALBERT prediction
+
 def predict_with_albert(text):
-    inputs = albert_tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=64)
+    inputs = albert_tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=128)
     with torch.no_grad():
         outputs = albert_model(**inputs)
         probs = torch.nn.functional.softmax(outputs.logits, dim=1)
@@ -41,8 +25,9 @@ def predict_with_albert(text):
     return ("Fraud" if pred_class == 1 else "Not Fraud"), float(pred_prob)
 
 # FinBERT prediction
+
 def predict_with_finbert(text):
-    inputs = finbert_tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=64)
+    inputs = finbert_tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=128)
     with torch.no_grad():
         outputs = finbert_model(**inputs)
         probs = torch.nn.functional.softmax(outputs.logits, dim=1)
@@ -51,23 +36,23 @@ def predict_with_finbert(text):
     return ("Fraud" if pred_class == 1 else "Not Fraud"), float(pred_prob)
 
 # Model selector
+
 def predict(text, model_name):
     if model_name == "ALBERT":
         return predict_with_albert(text)
     elif model_name == "FinBERT":
         return predict_with_finbert(text)
-    elif model_name == "CatBoost":
-        return predict_with_catboost(text)
     else:
         return "Unknown Model", 0.0
 
-# Examples from dataset
+# Updated examples: 3 fraud + 3 non-fraud, using training-format features
 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",
-    "Step: 285, Type: TRANSFER, Amount: 125487.45, Origin Balance: 0.0, Dest Balance: 524556.64",
-    "Step: 352, Type: PAYMENT, Amount: 41263.42, Origin Balance: 0.0, Dest Balance: 0.0",
-    "Step: 372, Type: CASH_IN, Amount: 187503.32, Origin Balance: 76827.0, Dest Balance: 0.0"
+    ["Step: 151, Type: CASH_OUT, Amount: 1633227.0, OldBalOrig: 1633227.0, NewBalOrig: 0.0, OldBalDest: 2865353.22, NewBalDest: 4498580.23", "ALBERT"],
+    ["Step: 353, Type: CASH_OUT, Amount: 174566.53, OldBalOrig: 174566.53, NewBalOrig: 0.0, OldBalDest: 1191715.74, NewBalDest: 1366282.27", "FinBERT"],
+    ["Step: 357, Type: TRANSFER, Amount: 484493.06, OldBalOrig: 484493.06, NewBalOrig: 0.0, OldBalDest: 0.0, NewBalDest: 0.0", "ALBERT"],
+    ["Step: 43, Type: CASH_OUT, Amount: 81571.63, OldBalOrig: 0.0, NewBalOrig: 0.0, OldBalDest: 176194.2, NewBalDest: 257765.83", "FinBERT"],
+    ["Step: 307, Type: DEBIT, Amount: 247.82, OldBalOrig: 11544.0, NewBalOrig: 11296.18, OldBalDest: 3550535.53, NewBalDest: 3550783.36", "ALBERT"],
+    ["Step: 350, Type: DEBIT, Amount: 4330.57, OldBalOrig: 3766.0, NewBalOrig: 0.0, OldBalDest: 239435.41, NewBalDest: 243765.98", "FinBERT"]
 ]
 
 # Gradio Interface
@@ -75,15 +60,15 @@ gui = gr.Interface(
     fn=predict,
     inputs=[
         gr.Textbox(label="Enter Transaction Description"),
-        gr.Dropdown(choices=["ALBERT", "FinBERT", "CatBoost"], label="Select Model", value="ALBERT")
+        gr.Dropdown(choices=["ALBERT", "FinBERT"], label="Select Model", value="ALBERT")
     ],
     outputs=[
         gr.Label(label="Prediction"),
         gr.Number(label="Fraud Probability")
     ],
-    examples=[[ex, "ALBERT"] for ex in examples],
+    examples=examples,
     title="💸 Fraud Detection Assistant",
-    description="Analyze transaction text using ALBERT, FinBERT, or CatBoost models."
+    description="Analyze transaction text using ALBERT or FinBERT models. Format: Step, Type, Amount, OldBalOrig, NewBalOrig, OldBalDest, NewBalDest."
 )
 
 # Launch the app