Update app.py

app.py

@@ -7,7 +7,19 @@ def chat(message, history):
     if not message:
         return ""
     
-    response = pipe(message, max_length=100, pad_token_id=50256)
-    return response[0]["generated_text"]
+    prompt = f"User: {message}\nAI Queen:"
+    response = pipe(prompt, max_length=100, pad_token_id=50256)
+    text = response[0]["generated_text"]
+    
+    if "AI Queen:" in text:
+        text = text.split("AI Queen:")[-1].strip()
+    
+    return text
+
+with gr.Blocks(title="AI Queen") as demo:
+    gr.Markdown("# AI Queen")
+    gr.Markdown("Your intelligent and majestic assistant")
+    gr.ChatInterface(fn=chat)
+
+demo.launch()
 
-gr.ChatInterface(fn=chat).launch()