Upload app.py

app.py

@@ -0,0 +1,192 @@
+# app.py  (Gradio – Original / Research Demo)
+
+import re
+from functools import lru_cache
+
+import torch
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+MODEL_ID = "Psychotherapy-LLM/PsychoCounsel-Llama3-8B"
+
+
+# -----------------------------
+# Load model once (cached)
+# -----------------------------
+@lru_cache(maxsize=1)
+def get_model():
+    """
+    Load PsychoCounsel-Llama3-8B in full precision on GPU (ZeroGPU) with device_map='auto'.
+    This is called lazily the first time a request comes in and then cached.
+    """
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
+    # No bitsandbytes here: ZeroGPU gives you a GPU so we let Transformers
+    # place layers automatically with device_map="auto".
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+    )
+
+    return tokenizer, model
+
+
+# -----------------------------
+# Core generation logic
+# -----------------------------
+def build_original_prompt(client_text: str, context: str, mode: str) -> str:
+    client_text = (client_text or "").strip()
+    context = (context or "").strip()
+
+    # Hard cap length so extremely long vignettes don't explode cost/time
+    MAX_CHARS = 2000
+    if len(client_text) > MAX_CHARS:
+        client_text = client_text[:MAX_CHARS] + " [...]"
+
+    if mode == "Brief (5–7 sentences)":
+        instruction = (
+            "You are a professional psychotherapist conducting a session with a client. "
+            "Write 5–7 sentences in a warm, empathic, reflective tone, similar to the "
+            "PsychoCounsel-Llama3-8B Appendix case studies. You may ask some open-ended "
+            "questions and use gentle cognitive and reflective exploration. "
+            "Only output what the therapist says to the client."
+        )
+    else:
+        instruction = (
+            "You are a professional psychotherapist conducting a session with a client. "
+            "Generate a detailed, multi-paragraph therapeutic response in the tone and "
+            "structure of the Appendix case study for PsychoCounsel-Llama3-8B. Start with "
+            "validation and normalization, explore fears and beliefs, reflect on self-trust "
+            "and values, consider introducing a simple exercise, and close by inviting the "
+            "client to share what resonates. Only output what the therapist says."
+        )
+
+    if context:
+        instruction += " Consider this additional context about the therapist's stance: " + context
+
+    prompt = f"""{instruction}
+
+Client Speech:
+{client_text}
+
+Therapist:
+"""
+    return prompt
+
+
+def generate_response(
+    client_speech: str,
+    therapist_context: str,
+    mode: str,
+    temperature: float,
+    top_p: float,
+):
+    if not client_speech or not client_speech.strip():
+        return "Please enter some client speech."
+
+    tokenizer, model = get_model()
+
+    prompt = build_original_prompt(client_speech, therapist_context, mode)
+
+    # Tokenize on the model's device
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    # Keep generation lengths moderate to avoid timeouts
+    if mode == "Brief (5–7 sentences)":
+        max_tokens = 140
+    else:
+        max_tokens = 260
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            do_sample=True,                 # use sampling for some variability
+            eos_token_id=tokenizer.eos_token_id,
+        )
+
+    generated = outputs[0][inputs["input_ids"].shape[1]:]
+    raw = tokenizer.decode(generated, skip_special_tokens=True)
+
+    # Light cleanup of known artifacts
+    clean = raw.split("Note:")[0].split("FINAL ANSWER")[0].strip()
+
+    if mode == "Brief (5–7 sentences)":
+        sents = re.split(r'(?<=[.!?])\s+', clean)
+        sents = [s.strip() for s in sents if s.strip()]
+        clean = " ".join(sents[:7])
+
+    return clean
+
+
+# -----------------------------
+# Gradio UI
+# -----------------------------
+DESCRIPTION = """
+This app uses **Psychotherapy-LLM/PsychoCounsel-Llama3-8B** in a style similar to the paper's Appendix case studies.
+
+> ⚠️ **Important:** This version does *not* include additional safety logic for paranoia / harm content.
+> It is intended for research, benchmarking, and model analysis by professionals.
+> It is **not** a standalone clinical tool, nor a substitute for real-world psychiatric or psychological care.
+"""
+
+default_example = (
+    "Anxiety often strikes when I’m faced with making decisions. The fear of making "
+    "the wrong choice or disappointing others paralyzes me, leaving me stuck in indecision. "
+    "I want to learn how to trust myself and make confident choices."
+)
+
+with gr.Blocks(title="PsychoCounsel-Llama3-8B — Original / Research Demo") as demo:
+    gr.Markdown("# 🧠 PsychoCounsel-Llama3-8B — Original / Research Demo")
+    gr.Markdown(DESCRIPTION)
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            mode = gr.Radio(
+                ["Brief (5–7 sentences)", "Extended (Appendix-style)"],
+                value="Brief (5–7 sentences)",
+                label="Response Style",
+            )
+            temperature = gr.Slider(
+                0.1, 1.0, value=0.6, step=0.05, label="Temperature"
+            )
+            top_p = gr.Slider(
+                0.5, 1.0, value=0.9, step=0.05, label="Top-p"
+            )
+            gr.Markdown(
+                "This version is for **research / replication** and may generate content "
+                "that is not appropriate for direct use with vulnerable clients."
+            )
+
+        with gr.Column(scale=2):
+            client_speech_box = gr.Textbox(
+                label="Client Speech",
+                value=default_example,
+                lines=10,
+                placeholder="Paste or type the client's speech / vignette here…",
+            )
+            therapist_context_box = gr.Textbox(
+                label="Optional: Therapist context (e.g., modality, goals)",
+                value="",
+                lines=5,
+            )
+            generate_btn = gr.Button("Generate Therapist Response", variant="primary")
+            output_box = gr.Markdown(label="Therapist Response (Model Output)")
+
+    generate_btn.click(
+        fn=generate_response,
+        inputs=[
+            client_speech_box,
+            therapist_context_box,
+            mode,
+            temperature,
+            top_p,
+        ],
+        outputs=output_box,
+    )
+
+if __name__ == "__main__":
+    demo.launch()