Update app.py

app.py

@@ -1,49 +1,81 @@
 from huggingface_hub import InferenceClient
 import gradio as gr
+import logging
+import sys
 
-client = InferenceClient(
-    "mistralai/Mixtral-8x7B-Instruct-v0.1"
-)
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-def format_prompt(message, history):
-    prompt = "<s>"
-    for user_prompt, bot_response in history:
-        prompt += f"[INST] {user_prompt} [/INST]"
-        prompt += f" {bot_response}</s> "
-    prompt += f"[INST] {message} [/INST]"
-    return prompt
-
-def generate(
-    prompt, history, system_prompt, temperature=0.9, max_new_tokens=512, top_p=0.95, repetition_penalty=1.0,
-):
-    temperature = float(temperature)
-    if temperature < 1e-2:
-        temperature = 1e-2
-    
-    top_p = float(top_p)
-    generate_kwargs = dict(
-        temperature=temperature,
-        max_new_tokens=max_new_tokens,
-        top_p=top_p,
-        repetition_penalty=repetition_penalty,
-        do_sample=True,
-        seed=42,
+try:
+    client = InferenceClient(
+        "mistralai/Mixtral-8x7B-Instruct-v0.1"
     )
+except Exception as e:
+    logger.error(f"Failed to initialize Hugging Face client: {str(e)}")
+    sys.exit(1)
 
-    # Check if the history is empty. If so, prepend the system prompt.
-    if not history:
-        formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", history)
-    else:
-        formatted_prompt = format_prompt(prompt, history)
+def format_prompt(message, history):
+    try:
+        prompt = "<s>"
+        if history:
+            for user_prompt, bot_response in history:
+                prompt += f"[INST] {user_prompt} [/INST]"
+                prompt += f" {bot_response}</s> "
+        prompt += f"[INST] {message} [/INST]"
+        logger.info(f"Formatted prompt: {prompt}")
+        return prompt
+    except Exception as e:
+        logger.error(f"Error in format_prompt: {str(e)}")
+        return None
+
+def generate(message, chat_history, system_prompt, temperature=0.9, max_new_tokens=512, top_p=0.95):
+    try:
+        logger.info(f"Received message: {message}")
+        logger.info(f"System prompt: {system_prompt}")
+        
+        # Format the full prompt
+        if not chat_history:
+            full_message = f"{system_prompt}\n\nUser: {message}"
+        else:
+            full_message = message
+            
+        formatted_prompt = format_prompt(full_message, chat_history)
+        
+        if not formatted_prompt:
+            return "I encountered an error formatting your message. Please try again."
+            
+        # Generation parameters
+        generate_kwargs = dict(
+            temperature=float(temperature),
+            max_new_tokens=int(max_new_tokens),
+            top_p=float(top_p),
+            do_sample=True,
+            seed=42,
+        )
+        
+        logger.info("Starting generation with parameters: %s", generate_kwargs)
+        
+        # Generate response
+        response_stream = client.text_generation(
+            formatted_prompt,
+            **generate_kwargs,
+            stream=True,
+            details=True,
+            return_full_text=False
+        )
+        
+        partial_message = ""
+        for response in response_stream:
+            if response.token.text:
+                partial_message += response.token.text
+                yield partial_message
 
-    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
-    output = ""
-    for response in stream:
-        output += response.token.text
-        yield output
-    return output
+    except Exception as e:
+        logger.error(f"Error in generate function: {str(e)}")
+        yield f"I encountered an error: {str(e)}"
 
-# Define a default system prompt for mental health support
+# Define the default system prompt
 DEFAULT_SYSTEM_PROMPT = """You are a supportive AI assistant trained to provide emotional support and general guidance. 
 Remember to: 
 1. Show empathy and understanding
@@ -52,52 +84,78 @@ Remember to:
 4. Encourage professional help when appropriate
 5. Maintain boundaries and ethical guidelines"""
 
-additional_inputs = [
-    gr.Textbox(
-        label="System Prompt",
-        value=DEFAULT_SYSTEM_PROMPT,
-        lines=3,
-        interactive=True,
-    ),
-    gr.Slider(
-        label="Temperature",
-        minimum=0.1,
-        maximum=1.0,
-        value=0.9,
-        step=0.1
-    ),
-    gr.Slider(
-        label="Max New Tokens",
-        minimum=64,
-        maximum=1024,
-        value=512,
-        step=64
+# Define the interface
+with gr.Blocks() as demo:
+    chatbot = gr.Chatbot(height=500)
+    msg = gr.Textbox(label="Message", placeholder="Type your message here...")
+    
+    with gr.Accordion("Advanced Options", open=False):
+        system_prompt = gr.Textbox(
+            value=DEFAULT_SYSTEM_PROMPT,
+            label="System Prompt",
+            lines=3
+        )
+        temperature = gr.Slider(
+            minimum=0.1,
+            maximum=1.0,
+            value=0.9,
+            step=0.1,
+            label="Temperature"
+        )
+        max_new_tokens = gr.Slider(
+            minimum=64,
+            maximum=1024,
+            value=512,
+            step=64,
+            label="Max Tokens"
+        )
+
+    clear = gr.Button("Clear")
+
+    def user(user_message, history):
+        return "", history + [[user_message, None]]
+
+    def bot(history, system_prompt, temperature, max_new_tokens):
+        if not history:
+            return history
+            
+        user_message = history[-1][0]
+        history[-1][1] = ""
+        
+        for chunk in generate(
+            user_message, 
+            history[:-1], 
+            system_prompt, 
+            temperature, 
+            max_new_tokens
+        ):
+            history[-1][1] = chunk
+            yield history
+
+    msg.submit(
+        user,
+        [msg, chatbot],
+        [msg, chatbot],
+        queue=False
+    ).then(
+        bot,
+        [chatbot, system_prompt, temperature, max_new_tokens],
+        chatbot
     )
-]
 
-examples = [
-    ["Patient is feeling stressed due to work and has trouble sleeping.", DEFAULT_SYSTEM_PROMPT, 0.9, 512, 0.95, 1.0],
-    ["Client is dealing with relationship issues and is seeking advice on communication strategies.", DEFAULT_SYSTEM_PROMPT, 0.9, 512, 0.95, 1.0],
-    ["Individual has recently experienced a loss and is having difficulty coping with grief.", DEFAULT_SYSTEM_PROMPT, 0.9, 512, 0.95, 1.0],
-]
+    clear.click(lambda: None, None, chatbot, queue=False)
 
-# Create the interface
-demo = gr.ChatInterface(
-    fn=generate,
-    chatbot=gr.Chatbot(
-        show_label=False, 
-        show_share_button=False, 
-        show_copy_button=True, 
-        layout="panel",
-        type="messages"  # Updated to use the new message format
-    ),
-    additional_inputs=additional_inputs,
-    title="PsyAssist - ADVANCING MENTAL HEALTH SUPPORT WITH AI-DRIVEN INTERACTION",
-    description="""This is an AI-powered mental health support chatbot. While it can provide emotional support and general guidance, 
-    it is not a replacement for professional mental health services. In case of emergency, please contact your local mental health crisis hotline.""",
-    examples=examples,
-    concurrency_limit=20,
-)
+    gr.Markdown("""
+    # PsyAssist - ADVANCING MENTAL HEALTH SUPPORT WITH AI-DRIVEN INTERACTION
+    
+    **Important Notice**: This is an AI-powered mental health support chatbot. While it can provide emotional support 
+    and general guidance, it is not a replacement for professional mental health services. In case of emergency, 
+    please contact your local mental health crisis hotline.
+    """)
 
 if __name__ == "__main__":
-    demo.launch(show_api=False)
+    try:
+        demo.queue().launch(show_api=False)
+    except Exception as e:
+        logger.error(f"Failed to launch Gradio interface: {str(e)}")
+        sys.exit(1)