Adding prompt message framework

app.py

@@ -8,6 +8,14 @@ pipe = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
 # Define the inference function
 def generate_text(prompt):
     start_time = time.time()
+    messages = [
+        {
+            "role": "system",
+            "content": "You are a friendly and helpful chatbot",
+        },
+        {"role": "user", "content": prompt},
+    ]
+    prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     results = pipe(prompt, max_length=100, num_return_sequences=1)
     end_time = time.time()
     response_time = end_time - start_time