import streamlit as st
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import numpy as np

# Load tokenizer and model from Hugging Face Hub
MODEL_NAME = "briangilbert/working"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)

# Define labels
id2label = {0: "NOT SCAM", 1: "SCAM"}

# Streamlit UI
st.title("💬 Fraud Detection in Text")
st.write("Enter a dialogue and check if it's a **SCAM** or **NOT SCAM**.")

# Text input
user_input = st.text_area("Enter a message:")

if st.button("Detect Fraud"):
    if user_input:
        # Tokenize input
        inputs = tokenizer(user_input, return_tensors="pt", truncation=True)

        # Get model prediction
        model.eval()
        with torch.no_grad():
            outputs = model(**inputs)
            predicted_class = torch.argmax(outputs.logits).item()

        # Display result
        st.success(f"🚨 Prediction: **{id2label[predicted_class]}**")
    else:
        st.warning("Please enter a dialogue.")