Update app.py

app.py

@@ -1,76 +1,75 @@
 import gradio as gr
 import torch
 from transformers import AlbertTokenizer, AlbertForSequenceClassification, BertTokenizer, BertForSequenceClassification
-import numpy as np
+import torch.nn.functional as F
 
-# Load ALBERT model and tokenizer from Hugging Face Hub
-albert_model = AlbertForSequenceClassification.from_pretrained("Deepaksai1/albert-fraud-detector-v2")
+# Load models
+albert_model = AlbertForSequenceClassification.from_pretrained("Deepaksai1/albert-fraud-detector-v2").eval()
 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-v2")
+finbert_model = BertForSequenceClassification.from_pretrained("Deepaksai1/finbert-fraud-detector-v2").eval()
 finbert_tokenizer = BertTokenizer.from_pretrained("Deepaksai1/finbert-fraud-detector-v2")
-finbert_model.eval()
 
-# ALBERT prediction
+# Inference function
+def predict_model(step, tx_type, amount, old_org, new_org, old_dest, new_dest, model_name):
+    text = f"Step: {step}, Type: {tx_type}, Amount: {amount}, " \
+           f"OldBalOrig: {old_org}, NewBalOrig: {new_org}, " \
+           f"OldBalDest: {old_dest}, NewBalDest: {new_dest}"
+    
+    tokenizer = albert_tokenizer if model_name == "ALBERT" else finbert_tokenizer
+    model = albert_model if model_name == "ALBERT" else finbert_model
 
-def predict_with_albert(text):
-    inputs = albert_tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=128)
+    inputs = 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)
-        pred_class = torch.argmax(probs).item()
-        pred_prob = probs[0][1].item()
-    return ("Fraud" if pred_class == 1 else "Not Fraud"), float(pred_prob)
+        outputs = model(**inputs)
+        probs = F.softmax(outputs.logits, dim=1)
+        pred = torch.argmax(probs).item()
+        fraud_score = probs[0][1].item()
 
-# FinBERT prediction
+    return "Fraud" if pred == 1 else "Not Fraud", round(fraud_score, 4)
 
-def predict_with_finbert(text):
-    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)
-        pred_class = torch.argmax(probs).item()
-        pred_prob = probs[0][1].item()
-    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)
-    else:
-        return "Unknown Model", 0.0
-
-# Updated examples: 3 fraud + 3 non-fraud, using training-format features
+# Example values
 examples = [
-    ["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"]
+    [151, "CASH_OUT", 1633227.0, 1633227.0, 0.0, 2865353.22, 4498580.23, "ALBERT"],
+    [353, "CASH_OUT", 174566.53, 174566.53, 0.0, 1191715.74, 1366282.27, "FinBERT"],
+    [357, "TRANSFER", 484493.06, 484493.06, 0.0, 0.0, 0.0, "ALBERT"],
+    [43, "CASH_OUT", 81571.63, 0.0, 0.0, 176194.2, 257765.83, "FinBERT"],
+    [307, "DEBIT", 247.82, 11544.0, 11296.18, 3550535.53, 3550783.36, "ALBERT"],
+    [350, "DEBIT", 4330.57, 3766.0, 0.0, 239435.41, 243765.98, "FinBERT"]
 ]
 
 # Gradio Interface
-gui = gr.Interface(
-    fn=predict,
-    inputs=[
-        gr.Textbox(label="Enter Transaction Description"),
-        gr.Dropdown(choices=["ALBERT", "FinBERT"], label="Select Model", value="ALBERT")
-    ],
-    outputs=[
-        gr.Label(label="Prediction"),
-        gr.Number(label="Fraud Probability")
-    ],
-    examples=examples,
-    title="💸 Fraud Detection Assistant",
-    description="Analyze transaction text using ALBERT or FinBERT models. Format: Step, Type, Amount, OldBalOrig, NewBalOrig, OldBalDest, NewBalDest."
-)
+with gr.Blocks() as demo:
+    gr.Markdown("## 🔎 Fraud Detection with ALBERT and FinBERT")
+    
+    with gr.Row():
+        step = gr.Number(label="Step", value=1)
+        tx_type = gr.Dropdown(choices=["CASH_OUT", "TRANSFER", "PAYMENT", "DEBIT", "CASH_IN"], label="Transaction Type")
+        amount = gr.Number(label="Amount", value=0.0)
+    
+    with gr.Row():
+        old_org = gr.Number(label="Old Balance Orig", value=0.0)
+        new_org = gr.Number(label="New Balance Orig", value=0.0)
+    
+    with gr.Row():
+        old_dest = gr.Number(label="Old Balance Dest", value=0.0)
+        new_dest = gr.Number(label="New Balance Dest", value=0.0)
+
+    model_selector = gr.Dropdown(choices=["ALBERT", "FinBERT"], value="ALBERT", label="Select Model")
+    
+    with gr.Row():
+        predict_btn = gr.Button("Predict")
+        pred_label = gr.Label(label="Prediction")
+        prob_score = gr.Number(label="Fraud Probability")
+
+    # Bind function
+    predict_btn.click(fn=predict_model,
+                      inputs=[step, tx_type, amount, old_org, new_org, old_dest, new_dest, model_selector],
+                      outputs=[pred_label, prob_score])
+
+    gr.Examples(examples=examples,
+                inputs=[step, tx_type, amount, old_org, new_org, old_dest, new_dest, model_selector])
 
-# Launch the app
+# Launch app
 if __name__ == "__main__":
-    gui.launch()
+    demo.launch()