app.py

import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import numpy as np
import matplotlib.pyplot as plt
import io
import pandas as pd
from PIL import Image
import gc
import shap 

# 1. MODEL CONFIGURATION
MODELS = {
    "MuRIL": "SpamX/MuRIL_V2",
    "XLM-RoBERTa": "SpamX/XLM_RoBERTa_V2"
}

cache = {"model": None, "tokenizer": None, "current_id": None, "last_winner": "MuRIL"}

def purge_memory():
    global cache
    cache.update({"model": None, "tokenizer": None, "current_id": None})
    gc.collect()
    if torch.cuda.is_available(): torch.cuda.empty_cache()

def get_assets(model_name):
    global cache
    model_id = MODELS[model_name]
    if cache["current_id"] != model_id:
        purge_memory()
        cache["tokenizer"] = AutoTokenizer.from_pretrained(model_id)
        cache["model"] = AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True)
        cache["current_id"] = model_id
    return cache["tokenizer"], cache["model"]

def get_gating_prediction(text):
    tk_m, md_m = get_assets("MuRIL")
    with torch.no_grad():
        out_m = torch.softmax(md_m(**tk_m(text, return_tensors="pt")).logits, dim=1)[0]
        prob_m_spam = out_m[1].item()
    conf_m = abs(prob_m_spam - 0.5) * 2 

    tk_x, md_x = get_assets("XLM-RoBERTa")
    with torch.no_grad():
        out_x = torch.softmax(md_x(**tk_x(text, return_tensors="pt")).logits, dim=1)[0]
        prob_x_spam = out_x[1].item()
    conf_x = abs(prob_x_spam - 0.5) * 2

    if conf_m >= conf_x:
        final_prob_spam, winner = prob_m_spam, "MuRIL"
    else:
        final_prob_spam, winner = prob_x_spam, "XLM-RoBERTa"
    
    cache["last_winner"] = winner
    label = "SPAM" if final_prob_spam > 0.5 else "HAM"
    score = final_prob_spam if label == "SPAM" else 1 - final_prob_spam
    return label, score, winner

def generate_shap_plot(tokens, shap_values, model_name):
    data = []
    for t, v in zip(tokens, shap_values):
        clean_t = t.replace(' ', '').replace('▁', '')
        if clean_t not in ['[CLS]', '[SEP]', '[PAD]', '<s>', '</s>', '']:
            data.append((clean_t, v))
    
    if not data: return None
    labels, values = zip(*data)
    colors = ['#ff4444' if v > 0 else '#00ff88' for v in values]

    plt_height = max(5, len(labels) * 0.4)
    plt.figure(figsize=(12, plt_height), facecolor='#050505')
    ax = plt.axes()
    ax.set_facecolor('#050505')
    ax.barh(labels, values, color=colors)
    ax.tick_params(axis='both', colors='white', labelsize=10)
    for s in ax.spines.values(): s.set_color('#333')
    plt.title(f"SHAP ATTRIBUTION: {model_name}", color='white', pad=20, fontsize=14)
    plt.xlabel("Shapley Value (Marginal Contribution)", color='white')
    plt.gca().invert_yaxis()
    
    buf = io.BytesIO()
    plt.savefig(buf, format='png', bbox_inches='tight', facecolor='#050505')
    buf.seek(0); plt.close()
    return Image.open(buf)

def predict(text, model_choice):
    if not text.strip(): return {"NULL": 0}, "STATUS: INPUT_REQUIRED"
    try:
        if model_choice == "Confidence-based Gating":
            label, conf, winner = get_gating_prediction(text)
            info = f"Decision Authority: {winner}"
        else:
            tk, md = get_assets(model_choice)
            cache["last_winner"] = model_choice
            with torch.no_grad():
                out = torch.softmax(md(**tk(text, return_tensors="pt")).logits, dim=1)[0]
                prob_spam = out[1].item()
                label = "SPAM" if prob_spam > 0.5 else "HAM"
                conf = prob_spam if label == "SPAM" else 1 - prob_spam
            info = f"Model: {model_choice}"
        return {label: conf, ("HAM" if label == "SPAM" else "SPAM"): 1-conf}, f"RESULT: {label} | {info}"
    except Exception as e:
        purge_memory()
        return {"ERROR": 1}, f"Error: {str(e)}"

def explain(text, model_choice):
    """Formal SHAP Audit: Synchronized with Gating Result."""
    if not text.strip(): return None, "### INPUT REQUIRED"
    
    # Force synchronization with the winner of the scan
    explainer_id = cache.get("last_winner", "MuRIL") if model_choice == "Confidence-based Gating" else model_choice
    tk, md = get_assets(explainer_id)
    
    def f(x):
        tv = tk(x.tolist(), padding=True, truncation=True, return_tensors="pt")
        with torch.no_grad(): outputs = md(**tv).logits
        return torch.softmax(outputs, dim=1).numpy()

    # Use a smaller perturbation sample for very short text like "Hello"
    explainer = shap.Explainer(f, tk)
    shap_results = explainer([text])
    
    prediction_scores = f(np.array([text]))[0]
    target_idx = np.argmax(prediction_scores)
    verdict = "SPAM" if target_idx == 1 else "HAM"
    v_color = "#ff4444" if verdict == "SPAM" else "#00ff88"
    
    # Get values and tokens
    current_values = shap_results.values[0][:, target_idx]
    raw_tokens = tk.convert_ids_to_tokens(tk(text)["input_ids"])
    expected_value = shap_results.base_values[0][target_idx]
    
    pos_impact, neg_impact, rows = 0, 0, []
    
    # Fixed Loop: Use zip to ensure we don't go out of bounds
    for t, v in zip(raw_tokens, current_values):
        # Clean special characters based on the specific model's vocabulary
        clean_t = t.replace(' ', '').replace('▁', '').replace('Ġ', '')
        
        # Skip padding and special markers
        if clean_t in ['[CLS]', '[SEP]', '<s>', '</s>', '<pad>', ''] or not clean_t.strip():
            continue
        
        if v > 0: 
            pos_impact += v
            c = v_color
            nature = verdict 
        else: 
            neg_impact += v
            c = "#ff4444" if verdict == "HAM" else "#00ff88"
            nature = "SPAM" if verdict == "HAM" else "HAM"
            
        rows.append(f"<tr><td style='padding:12px; border-bottom:1px solid #222;'>{clean_t}</td><td style='color:{c}; padding:12px; border-bottom:1px solid #222; font-weight:bold;'>{nature}</td><td style='color:{c}; padding:12px; border-bottom:1px solid #222;'>{v:+.4f}</td></tr>")

    # Final Convergence Calculation
    final_val = expected_value + pos_impact + neg_impact
    net_shift = pos_impact + neg_impact

    # If no tokens were meaningful (very common for "Hello"), provide a fallback
    if not rows:
        rows.append(f"<tr><td colspan='3' style='padding:20px; text-align:center; color:#888;'>Neutral content: No significant feature attribution found.</td></tr>")

    report_html = f"""
    <div style='font-family: sans-serif; color: white; background: #0a0a0a; padding: 25px; border-radius: 12px; border: 1px solid #333;'>
        <h2 style='color: {v_color}; text-align: center; border: 2px solid {v_color}; padding: 10px; border-radius: 8px; margin-bottom: 25px;'>
            SPAMX EXPLAINABILITY REPORT: {verdict}
        </h2>
        
        <div style='background: #111; padding: 20px; border-radius: 8px; border-left: 6px solid {v_color}; margin-bottom: 25px;'>
            <h3 style='margin-top: 0; color: #888; font-size: 0.85em;'>MATHEMATICAL AUDIT ({explainer_id})</h3>
            <p style='font-size: 1.1em; color: #ddd;'>
                Audit initialized at a <b>BASELINE</b> of <b>{expected_value:.4f}</b>.
            </p>

            <div style='background: #000; padding: 20px; border-radius: 8px; border: 1px solid #444; margin: 15px 0; font-family: monospace;'>
                <p style='color: #888; margin-bottom: 10px;'>CALCULATION LOG:</p>
                <p style='margin: 5px 0; font-size: 1.2em;'>&nbsp;&nbsp;Baseline: {expected_value:.4f}</p>
                <p style='margin: 5px 0; font-size: 1.2em; color: #00ff88;'>+ HAM IMPACT: { (pos_impact if verdict == "HAM" else neg_impact):+.4f}</p>
                <p style='margin: 5px 0; font-size: 1.2em; color: #ff4444;'>+ SPAM IMPACT: { (neg_impact if verdict == "HAM" else pos_impact):+.4f}</p>
                <hr style='border: 0.5px solid #333;'>
                <p style='margin: 10px 0 0 0; font-size: 1.3em; font-weight: bold; color: {v_color};'>
                    = FINAL CONVERGENCE: {final_val:.4f} ({final_val:.2%})
                </p>
            </div>

            <p style='margin-bottom: 0; font-size: 1.1em;'>
                <b>Conclusion:</b> A net shift of <b>{net_shift:+.4f}</b> confirmed the <b>{verdict}</b> verdict.
            </p>
        </div>

        <h3 style='color: #888; font-size: 0.85em; margin-bottom: 10px;'>TOKEN-LEVEL SHAPLEY VALUES</h3>
        <table style='width: 100%; border-collapse: collapse; font-size: 1em;'>
            <thead style='background: #1a1a1a; color: #888;'>
                <tr>
                    <th style='padding:12px; text-align:left; border-bottom: 2px solid #333;'>FEATURE (TOKEN)</th>
                    <th style='padding:12px; text-align:left; border-bottom: 2px solid #333;'>FEATURE INFLUENCE</th>
                    <th style='padding:12px; text-align:left; border-bottom: 2px solid #333;'>SHAPLEY VALUE</th>
                </tr>
            </thead>
            <tbody>
                {"".join(rows)}
            </tbody>
        </table>
    </div>
    """
    
    # Generate Plot
    plot_val = current_values if target_idx == 1 else -current_values
    plot_img = generate_shap_plot(raw_tokens, plot_val, explainer_id)
    return plot_img, report_html

# UI CONFIGURATION
css = """
body, .gradio-container { background-color: #050505 !important; color: white !important; font-family: 'Courier New', monospace; }
#title { text-align: center; color: #ff0000; letter-spacing: 5px; font-weight: 900; }
.gr-button-primary { background: #aa0000 !important; border: none !important; }
"""

with gr.Blocks(css=css) as demo:
    gr.Markdown("# SPAMX", elem_id="title")
    with gr.Row():
        with gr.Column(scale=2):
            input_text = gr.Textbox(label="INPUT", placeholder="Enter comment...", lines=4)
            model_select = gr.Dropdown(["MuRIL", "XLM-RoBERTa", "Confidence-based Gating"], value="Confidence-based Gating", label="CORE SELECTION")
            with gr.Row():
                btn_run = gr.Button("SCAN", variant="primary")
                btn_exp = gr.Button("EXPLAIN", variant="secondary")
        with gr.Column(scale=1):
            output_metric = gr.Label(label="PROBABILITY")
            output_verdict = gr.Markdown("STATUS: READY", elem_id="output-box")
    
    with gr.Tabs():
        with gr.TabItem("1. NARRATIVE"): 
            explain_html = gr.HTML("<div style='color: #666;'>Analysis required.</div>")
        with gr.TabItem("2. VISUAL"): 
            explain_plot = gr.Image(label="SHAP Attribution Map")

    gr.Markdown("---")
    gr.Examples(
        examples=[
            ["BROOO!! SUBSCRIBE CHEYYATHAVAR SUBSCRIBE CHEYYU!!"], 
            ["Sathyam parayalo video super! Pakshe check the link in my profile for cash."],
            ["Amazing production! Does anyone follow @Cr_ypto_Ma_ster strategy?🚀"],
            ["Don't click that link guys, it's a scam! Stay safe!"],
            ["Thakarppan video aliya! Waiting for next part!"]
        ],
        inputs=input_text,
        label="EXAMPLES"
    )

    btn_run.click(predict, [input_text, model_select], [output_metric, output_verdict])
    btn_exp.click(explain, [input_text, model_select], [explain_plot, explain_html])

if __name__ == "__main__":
    demo.launch()