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

# ====================== APP CONFIG ======================
st.set_page_config(
    page_title="AI Text Detector",
    page_icon="🤖",
    layout="centered"
)

st.title("🧠 AI Text Detector (DeBERTa-v3-large)")
st.markdown("""
This tool detects whether the given text is **Human-written** or **AI-generated**  
using a fine-tuned `microsoft/deberta-v3-large` model.
""")

# ====================== LOAD MODEL ======================
@st.cache_resource
def load_model():
    model_name = "kishankachhadiya/ai-text-detector-v-45.5"  # Replace with your HF model repo
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(model_name)
    return tokenizer, model

tokenizer, model = load_model()

# ====================== TEXT INPUT ======================
user_text = st.text_area(
    "Enter text to analyze:",
    placeholder="Paste or write any text here...",
    height=200
)

if st.button("🔍 Analyze Text", type="primary"):
    if not user_text.strip():
        st.warning("⚠️ Please enter some text.")
    else:
        with st.spinner("Analyzing..."):
            inputs = tokenizer(user_text, return_tensors="pt", truncation=True, padding=True, max_length=512)
            with torch.no_grad():
                outputs = model(**inputs)
                probs = torch.nn.functional.softmax(outputs.logits, dim=-1)[0]
                confidence, prediction = torch.max(probs, dim=0)

            label = "🤖 AI-generated" if prediction.item() == 1 else "🧍 Human-written"
            confidence_percent = confidence.item() * 100

        st.success(f"**Prediction:** {label}")
        st.progress(confidence.item())
        st.write(f"**Confidence:** {confidence_percent:.2f}%")

        # Detailed Probabilities
        st.markdown("### 📊 Detailed Probabilities")
        st.write({
            "Human (0)": f"{probs[0].item() * 100:.2f}%",
            "AI (1)": f"{probs[1].item() * 100:.2f}%"
        })

# ====================== FOOTER ======================
st.markdown("---")
st.caption("Built with ❤️ using [Streamlit](https://streamlit.io) and [Hugging Face Transformers](https://huggingface.co/transformers).")