src/streamlit_app.py

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

# ── Config ────────────────────────────────────────────────────────────────────
HF_CLASSIFIER_MODEL_ID = "ishaan1402/cbt-thought-pattern-classifier"
HF_T2T_MODEL_ID        = "ishaan1402/cbt-thought-pattern-t2t"

THRESHOLD  = 0.15
MAX_LENGTH = 256

PATTERN_CLASSES = [
    "Catastrophizing",
    "Discounting the positive",
    "Labeling and mislabeling",
    "Mental filtering",
    "Jumping to conclusions: mind reading",
    "Jumping to conclusions: Fortune-telling",
    "Overgeneralization",
    "Personalization",
    "Black-and-white or polarized thinking / All or nothing thinking",
    "Should statements",
    "None"
]

PATTERN_DESCRIPTIONS = {
    "Catastrophizing":              "Giving greater weight to the worst possible outcome.",
    "Discounting the positive":     "Rejecting positive experiences by insisting they don't count.",
    "Labeling and mislabeling":     "Attributing actions to character rather than situation.",
    "Mental filtering":             "Dwelling only on the negative details of a situation.",
    "Jumping to conclusions: mind reading": "Inferring negative thoughts from someone's behaviour.",
    "Jumping to conclusions: Fortune-telling": "Predicting negative outcomes of events.",
    "Overgeneralization":           "Making faulty generalisations from insufficient evidence.",
    "Personalization":              "Assigning disproportionate personal blame to oneself.",
    "Black-and-white or polarized thinking / All or nothing thinking":
                                    "Viewing things as either all good or all bad with no middle ground.",
    "Should statements":            "Demanding particular behaviours regardless of realistic circumstances.",
    "None":                         "No unhelpful thought pattern detected.",
}

STRATEGY_DEFINITIONS = {
    "Growth Mindset":   "Reframe a challenging event as an opportunity to grow instead of dwelling on the setbacks.",
    "Impermanence":     "Say that bad things don't last forever, will get better soon, and that others have experienced similar struggles.",
    "Neutralizing":     "Challenge the negative or catastrophic possibilities and reframe it with a neutral possibility.",
    "Optimism":         "Focus and be thankful for the positive aspects of the current situation.",
    "Self-Affirmation": "Say that the character can overcome the challenging event because of their strengths or values.",
}

# Shameless hardcoding, mapping thought patterns to their recommended CBT strategy
PATTERN_DEFAULT_STRATEGY = {
    "Catastrophizing":              "Neutralizing",
    "Discounting the positive":     "Optimism",
    "Labeling and mislabeling":     "Self-Affirmation",
    "Mental filtering":             "Optimism",
    "Jumping to conclusions: mind reading": "Neutralizing",
    "Jumping to conclusions: Fortune-telling": "Neutralizing",
    "Overgeneralization":           "Impermanence",
    "Personalization":              "Self-Affirmation",
    "Black-and-white or polarized thinking / All or nothing thinking": "Neutralizing",
    "Should statements":            "Growth Mindset",
    "None":                         "Optimism",
}

# ── Model loading ─────────────────────────────────────────────────────────────
@st.cache_resource
def load_classifier():
    tokenizer = AutoTokenizer.from_pretrained(HF_CLASSIFIER_MODEL_ID)
    model     = AutoModelForSequenceClassification.from_pretrained(HF_CLASSIFIER_MODEL_ID)
    model.eval()
    return tokenizer, model

@st.cache_resource
def load_reframer():
    tokenizer = AutoTokenizer.from_pretrained(HF_T2T_MODEL_ID)
    model     = AutoModelForSeq2SeqLM.from_pretrained(HF_T2T_MODEL_ID)
    model.eval()
    return tokenizer, model

# ── Inference ─────────────────────────────────────────────────────────────────
def classify(thought: str, persona: str, tokenizer, model) -> dict:
    input_text = f"Persona: {persona} | Thought: {thought}" if persona.strip() \
                 else f"Persona:  | Thought: {thought}"

    inputs = tokenizer(
        input_text,
        return_tensors="pt",
        max_length=MAX_LENGTH,
        truncation=True,
        padding="max_length"
    )
    with torch.no_grad():
        logits = model(**inputs).logits

    probs      = torch.softmax(logits, dim=-1).squeeze().numpy()
    top_idx    = int(np.argmax(probs))
    top_label  = PATTERN_CLASSES[top_idx]
    confidence = float(probs[top_idx])
    is_unhelpful = (top_label != "None") and (confidence >= THRESHOLD)

    return {
        "is_unhelpful":      is_unhelpful,
        "predicted_pattern": top_label if is_unhelpful else "None",
        "confidence":        confidence,
        "distribution":      {PATTERN_CLASSES[i]: float(p) for i, p in enumerate(probs)},
    }


def reframe(thought: str, pattern: str, strategy: str, tokenizer, model) -> str:
    strategy_def = STRATEGY_DEFINITIONS.get(strategy, "")
    input_text = (
        f"reframe thought: {pattern} | "
        f"strategy: {strategy} | "
        f"definition: {strategy_def} | "
        f"{thought}"
    )
    inputs = tokenizer(
        input_text,
        return_tensors="pt",
        max_length=MAX_LENGTH,
        truncation=True,
        padding="max_length"
    )
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=128,
            num_beams=4,
            early_stopping=True
        )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# ── UI ────────────────────────────────────────────────────────────────────────
st.set_page_config(page_title="CBT Thought Reframer", page_icon="🧠", layout="wide")

st.title("CBT Thought Reframer")
st.caption(
    "Enter a thought to detect its unhelpful cognitive pattern and generate an optimistic way to reframe your sentence. "
)

st.divider()

col_input, col_spacer = st.columns([2, 1])
with col_input:
    persona = st.text_input(
        "Persona (optional)",
        placeholder="e.g. I am a college student. I love playing the guitar.",
        help="Adding context about the person improves classification accuracy."
    )
    thought = st.text_area(
        "Thought",
        placeholder="e.g. I failed this test, I'm going to fail my entire degree.",
        height=120
    )
    run_btn = st.button("Analyze", type="primary", use_container_width=True)

if run_btn:
    if not thought.strip():
        st.warning("Please enter a thought to analyze.")
    else:
        st.divider()

        cls_col, ref_col = st.columns(2)

        # ── Classification ──
        with cls_col:
            st.subheader("Pattern Classification")
            with st.spinner("Classifying..."):
                cls_tok, cls_model = load_classifier()
                result = classify(thought, persona, cls_tok, cls_model)

            pattern    = result["predicted_pattern"]
            confidence = result["confidence"]
            unhelpful  = result["is_unhelpful"]

            if unhelpful:
                st.error(f"**{pattern}**")
            else:
                st.success("**No unhelpful pattern detected**")

            if pattern in PATTERN_DESCRIPTIONS:
                st.caption(PATTERN_DESCRIPTIONS[pattern])

            st.metric("Confidence", f"{confidence:.1%}")

            with st.expander("Full distribution"):
                sorted_dist = sorted(result["distribution"].items(), key=lambda x: x[1], reverse=True)
                for label, prob in sorted_dist:
                    col1, col2 = st.columns([3, 1])
                    label_display = f"**{label}**" if label == pattern else label
                    with col1:
                        st.markdown(label_display)
                        st.progress(int(prob * 100))
                    with col2:
                        st.markdown(f"`{prob:.3f}`")

        # ── Using classifier output to reframe ──
        with ref_col:
            st.subheader("Positive Reframe")

            # fall back to "None" for the reframer
            reframe_pattern = pattern if unhelpful else "None"
            strategy = PATTERN_DEFAULT_STRATEGY.get(reframe_pattern, "Neutralizing")

            with st.spinner("Reframing..."):
                ref_tok, ref_model = load_reframer()
                reframed = reframe(thought, reframe_pattern, strategy, ref_tok, ref_model)

            st.success(reframed)
            st.caption(
                f"Pattern fed to reframer: **{reframe_pattern}** · "
                f"Strategy: **{strategy}** — {STRATEGY_DEFINITIONS[strategy]}"
            )

        st.divider()
        st.caption(
            "Classifier: RoBERTa-large fine-tuned on PATTERNREFRAME · "
            "Reframer: T5-base fine-tuned on PATTERNREFRAME · "
            "Dataset: [facebookresearch/ParlAI](https://github.com/facebookresearch/ParlAI/tree/main/projects/reframe_thoughts)"
        )