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)" )