import streamlit as st
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from peft import PeftModel, PeftConfig
import torch

# Set page config
st.set_page_config(
    page_title="SMS Spam Detection",
    page_icon="🚀",
    layout="wide"
)

# Title and description
st.title("🚀 SMS Spam Detection")
st.markdown("""
This app uses a fine-tuned TinyLlama 1.1B model to detect spam messages.
Enter a message below to check if it's spam or not.
""")

# Cache the model loading
@st.cache_resource
def load_model():
    adapter_model_name = "deathVader-afk/tinyllama-sms-spam"
    
    # Load base model and tokenizer
    tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
    # Load model with PEFT adapter
    model = AutoModelForSequenceClassification.from_pretrained(
        adapter_model_name,
        num_labels=2,
        device_map="auto"
    )
    model.eval()    
    return tokenizer, model

# Load model with spinner
with st.spinner("Loading model..."):
    tokenizer, model = load_model()

st.success("Model loaded successfully!")

# Create two columns
col1, col2 = st.columns([2, 1])

with col1:
    # Text input
    message = st.text_area(
        "Enter your message:",
        height=150,
        placeholder="Type or paste your message here..."
    )
    
    # Predict button
    if st.button("🔍 Check Message", type="primary"):
        if message:
            with st.spinner("Analyzing..."):
                # Tokenize input
                inputs = tokenizer(message, return_tensors="pt", truncation=True, max_length=512)
                
                # Get prediction
                with torch.no_grad():
                    outputs = model(**inputs)
                    logits = outputs.logits
                    prediction = torch.argmax(logits, dim=1).item()
                    probabilities = torch.softmax(logits, dim=1)[0]
                
                # Display results
                st.markdown("---")
                st.subheader("📊 Results")
                
                if prediction == 1:
                    st.error("⚠️ This message is likely SPAM")
                    confidence = probabilities[1].item() * 100
                else:
                    st.success("✅ This message appears to be HAM (Not Spam)")
                    confidence = probabilities[0].item() * 100
                
                st.metric("Confidence", f"{confidence:.2f}%")
                
                # Show probability distribution
                st.markdown("### Probability Distribution")
                prob_col1, prob_col2 = st.columns(2)
                with prob_col1:
                    st.metric("HAM Probability", f"{probabilities[0].item() * 100:.2f}%")
                with prob_col2:
                    st.metric("SPAM Probability", f"{probabilities[1].item() * 100:.2f}%")
        else:
            st.warning("Please enter a message to analyze.")

with col2:
    st.markdown("### 📝 Examples")
    st.markdown("""
    **Try these examples:**
    
    🟢 **HAM:**
    - "Hey, want to grab coffee tomorrow?"
    - "Meeting at 3pm in conference room"
    - "Thanks for your help yesterday!"
    
    🔴 **SPAM:**
    - "URGENT! You've won $1000! Click here now!"
    - "FREE iPhone! Limited time offer!"
    - "Congratulations! Claim your prize now!"
    """)

# Footer
st.markdown("---")
st.markdown("""
<div style='text-align: center; color: gray;'>
    <p>Powered by TinyLlama 1.1B | Fine-tuned for SMS Spam Detection</p>
    <p>Model: <a href='https://huggingface.co/deathVader-afk/tinyllama-sms-spam' target='_blank'>deathVader-afk/tinyllama-sms-spam</a></p>
</div>
""", unsafe_allow_html=True)