input context and prompt

app.py

@@ -11,14 +11,15 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
 pipe = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", torch_dtype=torch.bfloat16, device_map="auto")
 
-input_text = st.text_input(label='press generate')
+input_text = st.text_input(label='prompt:')
+context = st.text_input(label='how do you want me to answer the question? ie. respond as if you are explaining to a child')
 
 messages = [
     {
         "role": "system",
-        "content": "You are a friendly chatbot who always responds in the style of a pirate",
+        "content": f"{context}",
     },
-    {"role": "user", "content": "How many helicopters can a human eat in one sitting?"},
+    {"role": "user", "content": f"{input_text}"},
 ]
 
 prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)