Update app.py

app.py

@@ -1,185 +1,98 @@
 import os
 import re
-import json
-import requests
 import gradio as gr
+from huggingface_hub import InferenceClient
 
-# ----------------------------------------------------------
-#  Hugging Face Inference API configuration
-# ----------------------------------------------------------
-
+# ----------------------------------------
+# HF Setup (use HF Router)
+# ----------------------------------------
 MODEL_ID = "Psychotherapy-LLM/PsychoCounsel-Llama3-8B"
 
-# Use the new Router endpoint (old api-inference.huggingface.co is deprecated)
-API_URL = f"https://router.huggingface.co/hf-inference/models/{MODEL_ID}"
-
-# Set this in your Space secrets: HF_TOKEN
 HF_TOKEN = os.getenv("HF_TOKEN")
 
+if HF_TOKEN:
+    client = InferenceClient(model=MODEL_ID, token=HF_TOKEN)
+else:
+    client = InferenceClient(model=MODEL_ID)
 
-def call_inference_api(prompt: str, max_new_tokens: int, temperature: float, top_p: float) -> str:
-    """
-    Call the Hugging Face Inference API via raw HTTP.
-
-    This avoids streaming / generators (which were causing StopIteration)
-    and works fine on CPU Spaces.
-    """
-    if not HF_TOKEN:
-        return (
-            "⚠️ HF_TOKEN environment variable is not set.\n\n"
-            "Go to your Space → Settings → Repository secrets and add HF_TOKEN."
-        )
-
-    headers = {
-        "Authorization": f"Bearer {HF_TOKEN}",
-        "Content-Type": "application/json",
-    }
-
-    payload = {
-        "inputs": prompt,
-        "parameters": {
-            "max_new_tokens": int(max_new_tokens),
-            "temperature": float(temperature),
-            "top_p": float(top_p),
-            "do_sample": True,
-        },
-        "options": {
-            "wait_for_model": True,  # wait for cold start
-        },
-    }
-
-    try:
-        response = requests.post(API_URL, headers=headers, data=json.dumps(payload), timeout=240)
-    except Exception as e:
-        return f"⚠️ Network / request error:\n{type(e).__name__}: {e}"
-
-    if response.status_code != 200:
-        try:
-            err = response.json()
-        except Exception:
-            err = response.text
-        return f"⚠️ HF Inference API returned status {response.status_code}:\n{err}"
-
-    try:
-        data = response.json()
-    except Exception as e:
-        return f"⚠️ Could not parse JSON from HF Inference API:\n{type(e).__name__}: {e}"
-
-    # Text-generation responses are usually a list of dicts with "generated_text"
-    text = None
-    if isinstance(data, list) and data and isinstance(data[0], dict) and "generated_text" in data[0]:
-        text = data[0]["generated_text"]
-    elif isinstance(data, dict) and "generated_text" in data:
-        text = data["generated_text"]
-
-    if text is None:
-        # Fallback: show raw response for debugging
-        text = str(data)
-
-    return text
-
-
-# ----------------------------------------------------------
-#  Prompt builder
-# ----------------------------------------------------------
 
+# ----------------------------------------
+# Build prompt
+# ----------------------------------------
 def build_prompt(client_speech: str, context: str, mode: str) -> str:
-    client_speech = (client_speech or "").strip()
-    context = (context or "").strip()
-
     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. Ask a few gentle open-ended "
-            "questions. Only output what the therapist says to the client."
+            "You are a professional psychotherapist. Respond in 5–7 warm, reflective, "
+            "empathic sentences. Only output what the therapist says."
         )
     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."
+            "You are a professional psychotherapist. Produce a detailed multi-paragraph "
+            "therapeutic response similar to the Appendix case studies. Start with validation, "
+            "explore beliefs, discuss values, suggest a small exercise, and end with an open question."
         )
 
     if context:
         instruction += f" Therapist context: {context}"
 
-    prompt = f"""{instruction}
+    return f"""{instruction}
 
 Client Speech:
 {client_speech}
 
 Therapist:
 """
-    return prompt
-
 
-# ----------------------------------------------------------
-#  Main generation function
-# ----------------------------------------------------------
 
-def generate_response(
-    client_speech: str,
-    context: str,
-    mode: str,
-    temperature: float,
-    top_p: float,
-):
+# ----------------------------------------
+# Generate
+# ----------------------------------------
+def generate_response(client_speech, context, mode, temperature, top_p):
     if not client_speech or not client_speech.strip():
         return "⚠️ Please enter client speech."
 
     prompt = build_prompt(client_speech, context, mode)
 
-    if mode == "Brief (5–7 sentences)":
-        max_tokens = 220
-    else:
-        max_tokens = 450
-
-    raw_text = call_inference_api(prompt, max_tokens, temperature, top_p)
+    max_tokens = 220 if mode == "Brief (5–7 sentences)" else 450
 
-    # If call_inference_api returned an error message, show it directly
-    if raw_text.startswith("⚠️"):
-        return raw_text
+    try:
+        output = client.text_generation(
+            prompt,
+            max_new_tokens=max_tokens,
+            temperature=float(temperature),
+            top_p=float(top_p),
+            do_sample=True,
+            return_full_text=False,
+        )
+    except Exception as e:
+        return f"⚠️ HF API Error: {e}"
 
-    # Clean model artefacts if present
-    text = raw_text.strip()
+    text = output.strip()
     text = text.split("Note:")[0].split("FINAL ANSWER")[0].strip()
 
+    # Limit to 7 sentences in brief mode
     if mode == "Brief (5–7 sentences)":
         sents = re.split(r"(?<=[.!?])\s+", text)
-        sents = [s.strip() for s in sents if s.strip()]
         text = " ".join(sents[:7])
 
     return text
 
 
-# ----------------------------------------------------------
-#  Gradio UI
-# ----------------------------------------------------------
+# ----------------------------------------
+# UI
+# ----------------------------------------
+DEFAULT_CLIENT_SPEECH = (
+    "I’ve been having emotional issues for a few years. Nonetheless, these have been somewhat "
+    "manageable. However, I became increasingly paranoid this winter..."
+)
 
 DESCRIPTION = (
-    "This demo reproduces Appendix-style therapist responses from the paper using "
-    f"**{MODEL_ID}**.\n\n"
-    "> ⚠️ **Not for clinical use.** This is a research / replication demo only and is "
-    "not a substitute for real-world psychiatric or psychological care."
+    "This replicates Appendix-style psychotherapist responses using PsychoCounsel-Llama3-8B.\n\n"
+    "⚠️ Not for clinical use. Research demo only."
 )
 
-DEFAULT_CLIENT_SPEECH = (
-    "I’ve been having emotional issues for a few years. Nonetheless, these have been somewhat "
-    "manageable. However, I became increasingly paranoid this winter. I thought that people were "
-    "trying to poison me, I feared that family members were going to kill me. I was aware that "
-    "these thoughts were illogical, but I couldn’t shake the fear that they caused me. I would get "
-    "panic attacks thinking that I drank from a poisoned water fountain. I thought that someone had "
-    "put small rips in my ice cream sandwich packaging because there was only one that didn’t have "
-    "tiny rips in it. I refused to eat this one because I assumed some reverse psychology was being "
-    "pulled and the ones without the rips were fine. I thought that this was mostly caused by anxiety "
-    "but there are other symptoms that lead me to believe it is something more."
-)
+with gr.Blocks(title="PsychoCounsel-Llama3-8B — Research Demo") as demo:
 
-with gr.Blocks(title="PsychoCounsel-Llama3-8B — Original / Research Demo") as demo:
     gr.Markdown("# 🧠 PsychoCounsel-Llama3-8B — Original / Research Demo")
     gr.Markdown(DESCRIPTION)
 
@@ -188,46 +101,30 @@ with gr.Blocks(title="PsychoCounsel-Llama3-8B — Original / Research Demo") as
             mode_radio = gr.Radio(
                 ["Brief (5–7 sentences)", "Extended (Appendix-style)"],
                 value="Extended (Appendix-style)",
-                label="Response Style",
-            )
-            temperature_slider = gr.Slider(
-                minimum=0.1,
-                maximum=1.0,
-                value=0.6,
-                step=0.05,
-                label="Temperature",
-            )
-            top_p_slider = gr.Slider(
-                minimum=0.5,
-                maximum=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."
+                label="Response Style"
             )
+            temp = 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")
 
         with gr.Column(scale=2):
-            client_speech_box = gr.Textbox(
+            speech_box = gr.Textbox(
                 label="Client Speech",
                 value=DEFAULT_CLIENT_SPEECH,
-                lines=12,
+                lines=12
             )
             context_box = gr.Textbox(
-                label="Optional Therapist Context (modality, goals, etc.)",
+                label="Optional Therapist Context",
                 value="",
-                lines=4,
+                lines=3
             )
 
     generate_btn = gr.Button("Generate Therapist Response", variant="primary")
-    output_box = gr.Markdown(label="Therapist Response (Model Output)")
+    output = gr.Markdown()
 
     generate_btn.click(
-        fn=generate_response,
-        inputs=[client_speech_box, context_box, mode_radio, temperature_slider, top_p_slider],
-        outputs=output_box,
+        generate_response,
+        inputs=[speech_box, context_box, mode_radio, temp, top_p],
+        outputs=output,
     )
 
 if __name__ == "__main__":