Update app.py

app.py

@@ -1,10 +1,10 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
 
-# Initialize the client with your model
+# Initialize the client with your model from Hugging Face Hub
 client = InferenceClient("Arnic/gemma2-2b-it-Pubmed20k-TPU")
 
-# Define response function
+# Define the function to handle chat responses
 def respond(
     message,
     history: list[tuple[str, str]],
@@ -13,23 +13,20 @@ def respond(
     temperature,
     top_p,
 ):
+    # System message to set the chatbot's tone
     system_message = (
         "You are a good listener. You advise relaxation exercises, suggest avoiding negative thoughts, "
         "and guide through steps to manage stress. Let's discuss what's on your mind, "
         "or ask me for a quick relaxation exercise."
     )
 
-    # Format history and system message as prompt text
-    chat_history = ""
+    # Format prompt with system message, chat history, and user message
+    prompt = system_message + "\n\n"
     for user_msg, bot_reply in history:
-        if user_msg:
-            chat_history += f"User: {user_msg}\n"
-        if bot_reply:
-            chat_history += f"Assistant: {bot_reply}\n"
-    
-    prompt = f"{system_message}\n\n{chat_history}User: {message}\nAssistant:"
+        prompt += f"User: {user_msg}\nAssistant: {bot_reply}\n"
+    prompt += f"User: {message}\nAssistant:"
 
-    # Generate response using the InferenceClient text generation method
+    # Call the text generation API
     response = client.text_generation(
         prompt=prompt,
         max_new_tokens=max_tokens,
@@ -37,11 +34,11 @@ def respond(
         top_p=top_p
     )
 
-    # Extract and yield the text response
-    generated_text = response["generated_text"].replace(prompt, "").strip()
+    # Extract the response text and yield it as output
+    generated_text = response.get("generated_text", "").replace(prompt, "").strip()
     yield generated_text
 
-# Set up Gradio interface
+# Gradio UI setup
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[