init

src/streamlit_app.py

@@ -1,40 +1,146 @@
-import altair as alt
-import numpy as np
-import pandas as pd
 import streamlit as st
+import torch
+import numpy as np
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+# ── Config ────────────────────────────────────────────────────────────────────
+HF_MODEL_ID = "ishaan1402/cbt-thought-pattern-classifier"
+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.",
+}
+
+# ── Model loading (cached so it only runs once per session) ───────────────────
+@st.cache_resource
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_ID)
+    model     = AutoModelForSequenceClassification.from_pretrained(HF_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)},
+    }
+
+# ── UI ────────────────────────────────────────────────────────────────────────
+st.set_page_config(page_title="CBT Thought Classifier", page_icon="🧠", layout="centered")
+
+st.title("🧠 CBT Thought Pattern Classifier")
+st.caption(
+    "Enter a thought below. The model will identify which unhelpful cognitive pattern "
+    "it exhibits, if any, based on Cognitive Behavioral Therapy (CBT) research."
+)
+
+st.divider()
+
+persona = st.text_input(
+    "Persona (optional)",
+    placeholder="e.g. I am a college student. I love playing 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
+)
+
+classify_btn = st.button("Classify", type="primary", use_container_width=True)
+
+if classify_btn:
+    if not thought.strip():
+        st.warning("Please enter a thought to classify.")
+    else:
+        with st.spinner("Loading model...") if "tokenizer" not in st.session_state else st.spinner("Classifying..."):
+            tokenizer, model = load_model()
+            result = classify(thought, persona, tokenizer, model)
+
+        st.divider()
+
+        pattern     = result["predicted_pattern"]
+        confidence  = result["confidence"]
+        unhelpful   = result["is_unhelpful"]
+
+        # ── Result banner ──
+        if unhelpful:
+            st.error(f"**Unhelpful pattern detected:** {pattern}")
+        else:
+            st.success("**No unhelpful pattern detected** — this thought looks okay.")
+
+        # ── Pattern description ──
+        if pattern in PATTERN_DESCRIPTIONS:
+            st.info(f"📖 **{pattern}:** {PATTERN_DESCRIPTIONS[pattern]}")
+
+        # ── Confidence ──
+        st.metric("Model confidence", f"{confidence:.1%}")
+
+        # ── Full distribution ──
+        with st.expander("See full probability distribution"):
+            dist = result["distribution"]
+            sorted_dist = sorted(dist.items(), key=lambda x: x[1], reverse=True)
+            for label, prob in sorted_dist:
+                bar_pct = int(prob * 100)
+                # Highlight the predicted class
+                label_display = f"**{label}**" if label == pattern else label
+                col1, col2 = st.columns([3, 1])
+                with col1:
+                    st.markdown(label_display)
+                    st.progress(bar_pct)
+                with col2:
+                    st.markdown(f"`{prob:.3f}`")
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+st.divider()
+st.caption("Model trained on the [PATTERNREFRAME dataset](https://github.com/facebookresearch/ParlAI/tree/main/projects/reframe_thoughts) · Built with RoBERTa-large")