Update app.py

app.py

@@ -1,83 +1,119 @@
 import gradio as gr
 from transformers import pipeline
+import re
 
-"""
-Psychotherapy Session Summarizer (Open-Source Implementation)
-
-This is a Hugging Face Spaces-compatible Gradio app that:
-- Takes psychotherapy session transcripts (as text input).
-- Summarizes key themes, emotional tones, and patterns.
-- Optionally allows custom instructions or focus areas (e.g., "Focus on client's progress since last session").
-- Utilizes open-source models only.
+# Initialize pipelines
+# Summarization pipeline with FLAN-T5
+summarizer = pipeline("text2text-generation", model="google/flan-t5-small", tokenizer="google/flan-t5-small")
 
-Modular:
-- Summarization model can be swapped easily.
-- Sentiment analysis model can be changed if desired.
+# Sentiment analysis pipeline
+sentiment_analyzer = pipeline("sentiment-analysis", model="cardiffnlp/twitter-roberta-base-sentiment-latest")
 
-Scalable:
-- Can be extended to handle multiple transcripts and clustering by topic.
+# Automatic speech recognition pipeline for audio
+asr_pipeline = pipeline("automatic-speech-recognition", model="openai/whisper-small")
 
-Note: 
-All used models are open-source and freely available on Hugging Face:
-- Summarization model: "google/flan-t5-small" (FLAN-T5 is open-sourced by Google)
-- Sentiment model: "cardiffnlp/twitter-roberta-base-sentiment-latest" (open-source sentiment analysis)
+def convert_to_json(transcript_text):
+    """
+    Convert the transcript into a structured JSON format.
+    Attempts to identify speaker turns based on lines starting with 'Therapist:' or 'Client:'.
+    If no clear pattern is found, the entire transcript is considered one turn.
+    """
+    lines = transcript_text.strip().split("\n")
+    session_data = []
+    
+    # Regex patterns to identify lines with a speaker
+    therapist_pattern = re.compile(r"^\s*(Therapist|T):", re.IGNORECASE)
+    client_pattern = re.compile(r"^\s*(Client|C):", re.IGNORECASE)
+    
+    current_speaker = None
+    current_text = []
+    
+    for line in lines:
+        line = line.strip()
+        if therapist_pattern.match(line):
+            # If we have accumulated text from previous speaker, store it
+            if current_speaker and current_text:
+                session_data.append({"speaker": current_speaker, "text": " ".join(current_text).strip()})
+                current_text = []
+            
+            current_speaker = "Therapist"
+            # Remove the speaker prefix
+            text_part = therapist_pattern.sub("", line).strip()
+            current_text.append(text_part)
+        
+        elif client_pattern.match(line):
+            if current_speaker and current_text:
+                session_data.append({"speaker": current_speaker, "text": " ".join(current_text).strip()})
+                current_text = []
+            
+            current_speaker = "Client"
+            text_part = client_pattern.sub("", line).strip()
+            current_text.append(text_part)
+        
+        else:
+            # Just text, append to current speaker's segment if identified
+            if current_speaker is None:
+                # No speaker identified yet, assume unknown
+                current_speaker = "Unknown"
+            current_text.append(line)
 
-You can adapt these to different models if needed, as long as they are open-source.
+    # Append the last collected segment
+    if current_speaker and current_text:
+        session_data.append({"speaker": current_speaker, "text": " ".join(current_text).strip()})
 
-Instructions:
-1. Paste or upload a session transcript (preferably a few paragraphs of conversation).
-2. (Optional) Provide a custom instruction or focus question, e.g., "Highlight moments of cognitive reframing."
-3. Click "Summarize" to generate a concise summary with themes and emotional insights.
+    # If no speakers identified at all and just one big chunk, still return it as JSON
+    if not session_data:
+        session_data = [{"speaker": "Unknown", "text": transcript_text.strip()}]
 
-"""
+    # Create a final JSON structure
+    json_data = {"session": session_data}
+    return json_data
 
-# Initialize pipelines
-summarizer = pipeline("text2text-generation", model="google/flan-t5-small", tokenizer="google/flan-t5-small")
-# Note: "text2text-generation" pipeline for FLAN-T5 also works for summarization if we prompt it properly.
-sentiment_analyzer = pipeline("sentiment-analysis", model="cardiffnlp/twitter-roberta-base-sentiment-latest")
-
-def analyze_session(transcript, custom_instruction):
-    # Basic input validation
-    if not transcript.strip():
-        return "Please provide a transcript."
+def analyze_session(transcript, custom_instruction, audio):
+    # If audio is provided, we transcribe it and ignore the text transcript field
+    if audio is not None:
+        # Transcribe audio
+        asr_result = asr_pipeline(audio)
+        transcript_text = asr_result['text']
+    else:
+        # Use the provided transcript text
+        transcript_text = transcript
     
-    # Construct a prompt for the summarization model
-    # FLAN-T5 is instruction tuned, so we can give it instructions directly.
-    # For example:
-    prompt = "Summarize the following psychotherapy session transcript, focusing on key themes, emotional shifts, and patterns."
-    if custom_instruction.strip():
-        prompt += " Additionally, " + custom_instruction.strip()
-    prompt += "\n\nTranscript:\n" + transcript.strip()
+    if not transcript_text.strip():
+        return "Please provide a transcript or an audio file."
     
-    # Generate summary
+    # Convert transcript to JSON
+    json_data = convert_to_json(transcript_text)
+
+    # Prepare the prompt for summarization
+    prompt = (
+        "You are a helpful assistant that summarizes psychotherapy sessions. "
+        "The session is provided in JSON format with speaker turns. "
+        "Summarize the key themes, emotional shifts, and patterns from this session. "
+    )
+    if custom_instruction.strip():
+        prompt += f" Additionally, {custom_instruction.strip()}"
+    prompt += "\n\nJSON data:\n" + str(json_data)
+
+    # Summarize using the LLM
     summary_output = summarizer(prompt, max_length=200, do_sample=False)
     summary = summary_output[0]['generated_text'].strip()
     
-    # Sentiment analysis: We'll run it on the entire transcript to gauge the overall emotional tone.
-    sentiment_results = sentiment_analyzer(transcript)
-    # The sentiment model returns something like: [{'label': 'positive', 'score': ...}]
-    # We'll aggregate the results (though it's a single input) and just pick the top.
+    # Sentiment analysis of the entire transcript
+    sentiment_results = sentiment_analyzer(transcript_text)
     main_sentiment = sentiment_results[0]['label']
     
-    # Construct a more informative result
-    # You could elaborate this logic: detect recurring concerns by simple keyword frequency analysis, etc.
-    # For a simple first iteration, just provide summary and sentiment.
-    
-    # Optional: Identify recurring concerns (simple keyword extraction)
-    # We'll do a naive keyword frequency approach just as a demonstration:
-    words = transcript.lower().split()
-    # Common therapy-related words (just a naive approach, could be replaced by a proper keyword extraction model)
-    # This is a placeholder for demonstration
+    # Simple keyword-based recurring concerns
+    words = transcript_text.lower().split()
     keywords_of_interest = ["anxiety", "depression", "relationship", "stress", "fear", "goals", "progress", "cognitive", "behavior"]
     recurring_concerns = [word for word in words if word in keywords_of_interest]
-    recurring_concerns = list(set(recurring_concerns))  # unique
+    recurring_concerns = list(set(recurring_concerns))
     if not recurring_concerns:
         recurring_concerns_str = "No specific recurring concerns identified from the predefined list."
     else:
         recurring_concerns_str = "Recurring concerns include: " + ", ".join(recurring_concerns)
     
-    # Recommended follow-up topics (just a heuristic based on summary)
-    # If certain keywords appear in summary, we can suggest follow-up:
+    # Suggest follow-up topics based on summary
     follow_up_suggestions = []
     if "progress" in summary.lower():
         follow_up_suggestions.append("Explore client's perception of progress in more detail.")
@@ -85,10 +121,8 @@ def analyze_session(transcript, custom_instruction):
         follow_up_suggestions.append("Discuss client's relationship dynamics further.")
     if not follow_up_suggestions:
         follow_up_suggestions.append("Consider following up on the emotional themes identified in the summary.")
-    
     follow_up_suggestions_str = " ".join(follow_up_suggestions)
     
-    # Combine results into a final output
     final_output = f"**Summary of Session:**\n{summary}\n\n**Overall Sentiment:** {main_sentiment}\n\n**{recurring_concerns_str}**\n\n**Suggested Follow-Up Topics:** {follow_up_suggestions_str}"
     
     return final_output
@@ -96,19 +130,24 @@ def analyze_session(transcript, custom_instruction):
 # Build Gradio UI
 description = """# Psychotherapy Session Summarizer
 
-Upload or paste your psychotherapy session transcript and optionally provide a custom instruction (e.g., "Focus on anxiety and coping strategies.").  
-Click 'Summarize' to generate a concise summary of key themes, emotional tones, recurring concerns, and suggested follow-up topics.
+This tool summarizes psychotherapy session transcripts (text or audio) into key themes, emotional shifts, and patterns.
+
+**How to Use:**
+- You may upload an audio file of the session or paste the text transcript.
+- Optionally provide a custom focus or instruction (e.g., "Focus on how the client talks about their anxiety.").
+- Click 'Summarize' to generate a summary along with identified concerns and suggested follow-ups.
 """
 
 with gr.Blocks() as demo:
     gr.Markdown(description)
     with gr.Row():
-        transcript_input = gr.Textbox(label="Session Transcript", lines=10, placeholder="Paste the session transcript here...")
-    custom_instruction_input = gr.Textbox(label="Custom Instruction (Optional)", placeholder="e.g., Focus on how the client describes their feelings about progress.")
+        transcript_input = gr.Textbox(label="Session Transcript (Text)", lines=10, placeholder="Paste the session transcript here...")
+        audio_input = gr.Audio(source="upload", type="file", label="Session Audio (Optional)")
+    custom_instruction_input = gr.Textbox(label="Custom Instruction (Optional)", placeholder="e.g., Focus on anxiety and coping strategies.")
     summarize_button = gr.Button("Summarize")
     output_box = gr.Markdown()
     
-    summarize_button.click(fn=analyze_session, inputs=[transcript_input, custom_instruction_input], outputs=output_box)
+    summarize_button.click(fn=analyze_session, inputs=[transcript_input, custom_instruction_input, audio_input], outputs=output_box)
 
 if __name__ == "__main__":
     demo.launch()