import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import numpy as np
from sklearn.preprocessing import LabelEncoder

MODEL_ID = "ModSpecialization/distilbert-base-uncased-fraud-classifer"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)

label_encoder = LabelEncoder()
label_encoder.classes_ = np.array(["No", "Yes"])

def classify_fraud(customer_id, terminal_id, amount, timestamp):
    text = f"{customer_id} | {terminal_id} | {amount} | {timestamp}"
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    outputs = model(**inputs)
    pred = np.argmax(outputs.logits.detach().numpy(), axis=1)
    print(pred)
    return f"Prediction: {label_encoder.inverse_transform(pred)[0]}"

demo = gr.Interface(
    fn=classify_fraud,
    inputs=[
        gr.Textbox(label="Customer ID"),
        gr.Textbox(label="Terminal ID"),
        gr.Number(label="Transaction Amount ($)"),
        gr.Textbox(label="Transaction Time (YYYY-MM-DD HH:MM:SS)")
    ],
    outputs="text",
    title="Credit Card Fraud Classifier",
    description="Enter transaction details to predict if it's fraudulent"
)

demo.launch()