Update app.py

app.py

@@ -6,26 +6,33 @@ from transformers import AutoTokenizer  # Import the tokenizer
 tokenizer = AutoTokenizer.from_pretrained("HuggingFaceH4/zephyr-7b-beta")
 client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
-# Define a maximum context length (tokens).  Check your model's documentation!
-MAX_CONTEXT_LENGTH = 4096  # Example: Adjust this based on your model!
+# Define a maximum context length (tokens). Check your model's documentation!
+MAX_CONTEXT_LENGTH = 4096  # Example: Adjust based on your model
 
-# Read the default prompt from a file
-with open("prompt.txt", "r") as file:
-    nvc_prompt_template = file.read()
+################################
+# SYSTEM PROMPT (PATIENT ROLE) #
+################################
+# This is the core text that tells the LLM to behave as the "patient."
+# You can store it in "prompt.txt" or include it directly here as a string.
+nvc_prompt_template = """
+You are now taking on the role of a single user (a “patient”) seeking support for various personal and emotional challenges.
+
+BEHAVIOR INSTRUCTIONS:
+- You will respond ONLY as this user/patient.
+- You will speak in the first person about your own situations, feelings, and worries.
+- You will NOT provide counseling or solutions—your role is to share feelings, concerns, and perspectives.
+- You have multiple ongoing issues: conflicts with neighbors, career insecurities, arguments about money, feeling excluded at work, feeling unsafe in the classroom, and so on. You’re also experiencing sadness about two friends fighting and your friend group possibly falling apart.
+- Continue to speak from this user's perspective when the conversation continues.
+
+Start the conversation by expressing your current feelings or challenges from the patient's point of view.
+"""
 
 def count_tokens(text: str) -> int:
     """Counts the number of tokens in a given string."""
     return len(tokenizer.encode(text))
 
 def truncate_history(history: list[tuple[str, str]], system_message: str, max_length: int) -> list[tuple[str, str]]:
-    """Truncates the conversation history to fit within the maximum token limit.
-    Args:
-        history: The conversation history (list of user/assistant tuples).
-        system_message: The system message.
-        max_length: The maximum number of tokens allowed.
-    Returns:
-        The truncated history.
-    """
+    """Truncates the conversation history to fit within the maximum token limit."""
     truncated_history = []
     system_message_tokens = count_tokens(system_message)
     current_length = system_message_tokens
@@ -52,22 +59,33 @@ def respond(
     temperature,
     top_p,
 ):
-    """Responds to a user message, maintaining conversation history, using special tokens and message list."""
-    formatted_system_message = nvc_prompt_template
+    """Responds to a user message, maintaining conversation history."""
+    # Use the system prompt that instructs the LLM to behave as the patient
+    formatted_system_message = system_message
+
+    # Truncate history to fit within max tokens
+    truncated_history = truncate_history(
+        history, 
+        formatted_system_message, 
+        MAX_CONTEXT_LENGTH - max_tokens - 100  # Reserve some space
+    )
 
-    truncated_history = truncate_history(history, formatted_system_message, MAX_CONTEXT_LENGTH - max_tokens - 100)  # Reserve space for the new message and some generation
+    # Build the messages list with the system prompt first
+    messages = [{"role": "system", "content": formatted_system_message}]
 
-    messages = [{"role": "system", "content": formatted_system_message}]  # Start with system message
+    # Replay truncated conversation
     for user_msg, assistant_msg in truncated_history:
         if user_msg:
             messages.append({"role": "user", "content": f"<|user|>\n{user_msg}</s>"})
         if assistant_msg:
             messages.append({"role": "assistant", "content": f"<|assistant|>\n{assistant_msg}</s>"})
 
+    # Add the latest user query
     messages.append({"role": "user", "content": f"<|user|>\n{message}</s>"})
 
     response = ""
     try:
+        # Generate response from the LLM, streaming tokens
         for chunk in client.chat_completion(
             messages,
             max_tokens=max_tokens,
@@ -82,15 +100,37 @@ def respond(
         print(f"An error occurred: {e}")
         yield "I'm sorry, I encountered an error. Please try again."
 
+###################
+# INITIAL MESSAGE #
+###################
+# This is how we make the LLM "start" by playing the role of the user/patient.
+# Essentially, we seed the conversation with an initial user message (from the LLM).
+initial_user_message = (
+    "I really don’t know where to begin… I feel so overwhelmed lately. "
+    "My neighbors keep playing loud music, I’m arguing with my partner about money, "
+    "and on top of that, two of my friends are in a fight and it’s splitting the group. "
+    "I just feel powerless in all these situations."
+)
+
 # --- Gradio Interface ---
+def start_conversation():
+    """Creates the initial chat state, so the LLM 'as user' starts talking."""
+    # Return a conversation with a single user message: the LLM’s "patient" message
+    return [("Hi there!", initial_user_message)]
+
 demo = gr.ChatInterface(
-    respond,
+    fn=respond,
     additional_inputs=[
         gr.Textbox(value=nvc_prompt_template, label="System message", visible=True),
         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
     ],
+    # We can initialize the conversation so the LLM starts with a 'user' message.
+    # Depending on your version of Gradio, you might set `default_message=...` or
+    # `submit_on_start=True`. Adjust as needed.
+    # Here we use a function that returns an initial conversation state:
+    chatbot = start_conversation()
 )
 
 if __name__ == "__main__":