app.py

import streamlit as st
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import RobertaTokenizerFast, RobertaModel
import re
import os

# ── Page config ──────────────────────────────────────────────────────────────
st.set_page_config(
    page_title="Sarcasm Detector",
    page_icon="🙄",
    layout="centered"
)

# ── Model Definition (must match training) ───────────────────────────────────
class SarcasmDetector(nn.Module):
    def __init__(self, model_name="roberta-base", num_sarc=2, num_sent=3, dropout=0.15):
        super().__init__()
        self.roberta = RobertaModel.from_pretrained(model_name)
        H = self.roberta.config.hidden_size  # 768

        self.sarcasm_head = nn.Sequential(
            nn.Dropout(dropout), nn.Linear(H, 256),
            nn.GELU(), nn.Dropout(dropout), nn.Linear(256, num_sarc))

        self.sentiment_head = nn.Sequential(
            nn.Dropout(dropout), nn.Linear(H, 256),
            nn.GELU(), nn.Dropout(dropout), nn.Linear(256, num_sent))

        self.confidence_gate = nn.Linear(H, 1)

    def forward(self, input_ids, attention_mask):
        out = self.roberta(input_ids=input_ids, attention_mask=attention_mask)
        cls = out.last_hidden_state[:, 0, :]
        return (self.sarcasm_head(cls), self.sentiment_head(cls),
                torch.sigmoid(self.confidence_gate(cls)), None)

# ── Constants ─────────────────────────────────────────────────────────────────
SARCASM_ID2LABEL   = {0: "not_sarcastic", 1: "sarcastic"}
SENTIMENT_ID2LABEL = {0: "negative", 1: "neutral", 2: "positive"}
MAX_LEN            = 128
SARCASM_THRESH     = 0.45      # matches notebook CFG.SARCASM_THRESH
DEVICE             = torch.device("cpu")

# ── Load model (cached) ───────────────────────────────────────────────────────
@st.cache_resource
def load_model():
    tokenizer = RobertaTokenizerFast.from_pretrained("roberta-base")
    model     = SarcasmDetector("roberta-base").to(DEVICE)

    ckpt_path = "best_model.pt"
    if os.path.exists(ckpt_path):
        try:
            state_dict = torch.load(ckpt_path, map_location=DEVICE, weights_only=False)
            model.load_state_dict(state_dict, strict=True)
            model.eval()
            return tokenizer, model, True, "✅ Model loaded successfully!"
        except Exception as e:
            return tokenizer, model, False, f"❌ Load error: {e}"
    return tokenizer, model, False, f"❌ File not found: {ckpt_path} | Files: {os.listdir('.')}"

def clean_text(text):
    return re.sub(r"\s+", " ", str(text)).strip()

@torch.no_grad()
def predict(text, tokenizer, model):
    text = clean_text(text)
    enc  = tokenizer(text, max_length=MAX_LEN, padding="max_length",
                     truncation=True, return_tensors="pt")
    sar_l, sen_l, conf, _ = model(
        enc["input_ids"].to(DEVICE),
        enc["attention_mask"].to(DEVICE)
    )
    sar_p  = F.softmax(sar_l, dim=1)[0].cpu().numpy()
    sen_p  = F.softmax(sen_l, dim=1)[0].cpu().numpy()
    conf_v = conf[0].item()

    # ── Conversational sarcasm heuristic boost ────────────────────────────
    CONVERSATIONAL_TRIGGERS = [
        "best day ever", "oh great", "just what i needed", "wow thanks",
        "oh wonderful", "fantastic news", "oh perfect", "love it when",
        "because that's totally", "yeah right", "sure it is", "oh sure",
        "oh absolutely", "what could go wrong", "totally fine", "great idea",
        "brilliant idea", "oh brilliant", "oh how lovely", "how lovely",
        "oh joy", "oh lovely", "oh nice", "oh how nice", "oh how great",
        "best thing ever", "love mondays", "love traffic", "love waiting",
        "so fun", "so exciting", "oh can't wait", "my favourite",
        "my favorite", "best part", "best news",
    ]
    text_lower = text.lower()
    triggered  = any(t in text_lower for t in CONVERSATIONAL_TRIGGERS)
    if triggered and sar_p[1] < 0.60:
        boost    = 0.55 + (0.10 * sar_p[1])
        sar_p[1] = max(sar_p[1], boost)
        sar_p[0] = 1.0 - sar_p[1]
    # ─────────────────────────────────────────────────────────────────────

    sar_id = 1 if sar_p[1] >= SARCASM_THRESH else 0
    sen_id = int(sen_p.argmax())

    # Sarcasm-aware sentiment gate
    if sar_id == 1 and sar_p[1] > 0.65:
        final_sentiment = "implicit_negative ⚠️"
    else:
        final_sentiment = SENTIMENT_ID2LABEL[sen_id]

    return {
        "sarcasm_label":   SARCASM_ID2LABEL[sar_id],
        "sarcasm_prob":    float(sar_p[1]),
        "not_sarc_prob":   float(sar_p[0]),
        "sentiment":       final_sentiment,
        "confidence":      float(conf_v),
        "sen_probs":       sen_p.tolist(),
    }

# ── UI ────────────────────────────────────────────────────────────────────────
st.markdown("""
<h1 style='text-align:center; font-size:2.5rem;'>🙄 Sarcasm Detector</h1>
<p style='text-align:center; color:gray; font-size:1rem;'>
Multi-Task RoBERTa · F1-Macro 0.977 · AUC 0.997 · Sentiment-Incongruence Framework
</p>
""", unsafe_allow_html=True)

st.markdown("---")

# Load model
tokenizer, model, model_loaded, load_msg = load_model()

st.caption(load_msg)

if not model_loaded:
    st.error("⚠️ Model weights not loaded correctly. Check logs above.")
    st.stop()

# Input
st.markdown("### 📝 Enter Text")
user_input = st.text_area(
    label="",
    placeholder='e.g. "Study finds that watching TV all day is actually good for your health"',
    height=120
)

# Example buttons
st.markdown("**Try an example:**")
col1, col2, col3 = st.columns(3)
with col1:
    if st.button("😏 Headline Sarcasm"):
        user_input = "Study Confirms That Working 80 Hours A Week Is The Key To A Fulfilling Life"
with col2:
    if st.button("😏 Casual Sarcasm"):
        user_input = "Oh great, my flight got delayed again. Best day ever."
with col3:
    if st.button("😊 Not Sarcastic"):
        user_input = "Dog reunited with owner after being missing for three days."

# Predict
if st.button("🔍 Analyse", type="primary", use_container_width=True):
    if not user_input.strip():
        st.warning("Please enter some text first!")
    else:
        with st.spinner("Analysing..."):
            result = predict(user_input, tokenizer, model)

        st.markdown("---")
        st.markdown("### 📊 Results")

        # Sarcasm result
        is_sarc = result["sarcasm_label"] == "sarcastic"
        sarc_color  = "#e74c3c" if is_sarc else "#2ecc71"
        sarc_emoji  = "😏 SARCASTIC" if is_sarc else "😊 NOT SARCASTIC"
        sarc_pct    = result["sarcasm_prob"] * 100

        st.markdown(f"""
        <div style='background:{sarc_color}22; border-left:5px solid {sarc_color};
                    padding:16px; border-radius:8px; margin-bottom:12px;'>
            <h2 style='color:{sarc_color}; margin:0;'>{sarc_emoji}</h2>
            <p style='margin:4px 0 0; font-size:1rem;'>
                Sarcasm probability: <strong>{sarc_pct:.1f}%</strong> &nbsp;|&nbsp;
                Confidence: <strong>{result['confidence']*100:.1f}%</strong>
            </p>
        </div>
        """, unsafe_allow_html=True)

        # Sentiment result
        sent = result["sentiment"]
        sent_color = {"positive":"#2ecc71","negative":"#e74c3c","neutral":"#f39c12"}.get(
            sent.split()[0], "#9b59b6")
        sent_emoji = {"positive":"😊","negative":"😔","neutral":"😐"}.get(
            sent.split()[0], "⚠️")

        st.markdown(f"""
        <div style='background:{sent_color}22; border-left:5px solid {sent_color};
                    padding:12px; border-radius:8px; margin-bottom:12px;'>
            <p style='margin:0; font-size:1rem;'>
                {sent_emoji} <strong>Sentiment:</strong> {sent.upper()}
            </p>
        </div>
        """, unsafe_allow_html=True)

        # Probability bars
        st.markdown("#### Sarcasm Probability")
        c1, c2 = st.columns(2)
        with c1:
            st.metric("😏 Sarcastic",    f"{result['sarcasm_prob']*100:.1f}%")
        with c2:
            st.metric("😊 Not Sarcastic", f"{result['not_sarc_prob']*100:.1f}%")

        st.progress(result["sarcasm_prob"])

        # Sentiment breakdown
        st.markdown("#### Sentiment Breakdown")
        sent_labels = ["Negative", "Neutral", "Positive"]
        sent_colors_list = ["#e74c3c", "#f39c12", "#2ecc71"]
        cols = st.columns(3)
        for i, (col, label, color) in enumerate(zip(cols, sent_labels, sent_colors_list)):
            with col:
                st.metric(label, f"{result['sen_probs'][i]*100:.1f}%")

st.markdown("---")

# Model info
with st.expander("ℹ️ About this model"):
    st.markdown("""
    **Architecture:** Multi-Task RoBERTa-base
    - Primary task: Sarcasm detection (binary, Focal Loss γ=2)
    - Auxiliary task: Sentiment analysis (3-class)
    - Novel contribution: Sentiment-Incongruence Auto-Labeler

    **Performance (Twitter Headlines dataset):**
    | Metric | Score |
    |--------|-------|
    | Accuracy | 97.74% |
    | F1-Macro | **0.9773** |
    | AUC-ROC | **0.9967** |

    **Training:** 4 epochs · RoBERTa-base · Focal Loss · WeightedRandomSampler

    **Built by:** Bharath Kesav R · Amrita Vishwa Vidyapeetham
    """)