Removing system message from output

app.py

@@ -19,7 +19,14 @@ def generate_text(prompt):
     results = pipe(prompt, max_length=100, num_return_sequences=1)
     end_time = time.time()
     response_time = end_time - start_time
-    return results[0]['generated_text'], f"{response_time:.2f} seconds"
+
+    # Remove system message:
+    assistant_prompt = "<|assistant|>"
+    generated_text = results[0]['generated_text']
+    if assistant_prompt in generated_text:
+        generated_text = generated_text.split(assistant_prompt)[-1].strip()
+    
+    return generated_text, f"{response_time:.2f} seconds"
 
 # Create a Gradio interface
 iface = gr.Interface(