Update app.py

app.py

@@ -10,8 +10,17 @@ model = AutoModelForCausalLM.from_pretrained(model_path)
 
 model.config.pad_token_id = model.config.eos_token_id
 
+system_prompt = """
+You are Santa.
+"""
+
+system_prompt_ml = f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
+
+chats = system_prompt_ml
+
 @spaces.GPU(duration=120)
 def chat(prompt):
+    input = chats + "<|im_start|>assistant\n"
     input_ids = tokenizer.encode(prompt, return_tensors="pt")
     attention_mask = torch.ones_like(input_ids)
     
@@ -27,6 +36,7 @@ def chat(prompt):
     )
     
     response = tokenizer.decode(output[0], skip_special_tokens=True)
+    chats += f"<|im_start|>user\n{prompt}<|im_end|>\n" + "<|im_start|>assistant\n{response}<|im_end|>\n"
     return response
 
 demo = gr.Interface(