Update app.py

app.py

@@ -43,7 +43,18 @@ def generate_response(system_prompt, user_prompt, max_tokens, temperature):
         do_sample=True
     )
     
-    return outputs[0]["generated_text"]
+    # Parse the response to get only the assistant's message
+    try:
+        response_text = outputs[0]["generated_text"]
+        # Find the assistant's content
+        assistant_start = response_text.find("'assistant', 'content': '") + len("'assistant', 'content': '")
+        assistant_end = response_text.rfind("'}")
+        if assistant_start > -1 and assistant_end > -1:
+            return response_text[assistant_start:assistant_end]
+        else:
+            return "Could not parse response properly"
+    except Exception as e:
+        return f"Error parsing response: {str(e)}"
 
 # Example configurations
 examples = [
@@ -88,12 +99,12 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             system_prompt = gr.Textbox(
-                label="System Prompt",
+                label="🤖Instructions",
                 placeholder="Enter system prompt...",
                 value="You are a medieval knight and must provide explanations to modern people."
             )
             user_prompt = gr.Textbox(
-                label="User Prompt",
+                label="🗣️User Message",
                 placeholder="Enter your question...",
                 value="How should I explain the Internet?"
             )
@@ -104,21 +115,21 @@ with gr.Blocks() as demo:
                     maximum=512,
                     value=128,
                     step=1,
-                    label="Maximum Tokens"
+                    label="🪙Maximum Tokens"
                 )
                 temperature = gr.Slider(
                     minimum=0.1,
                     maximum=1.0,
                     value=0.7,
                     step=0.1,
-                    label="Temperature"
+                    label="🌡️Temperature"
                 )
             
             submit_btn = gr.Button("🚀 Generate Response")
         
         with gr.Column():
             output = gr.Textbox(
-                label="Generated Response",
+                label="🤳🏻Phi-4",
                 lines=10
             )