Update app.py

app.py

@@ -21,13 +21,28 @@ tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
 model = transformers.AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
 
 class OrcaChatBot:
+    # Code below from [microsoft/ari9dam](https://huggingface.co/spaces/ari9dam/Orca-2-13B)  
     def __init__(self, model, tokenizer, system_message="You are Orca, an AI language model created by Microsoft. You are a cautious assistant. You carefully follow instructions. You are helpful and harmless and you follow ethical guidelines and promote positive behavior."):
         self.model = model
         self.tokenizer = tokenizer
         self.system_message = system_message
+        self.conversation_history = []
+
+    def update_conversation_history(self, user_message, assistant_message):
+        self.conversation_history.append(("user", user_message))
+        self.conversation_history.append(("assistant", assistant_message))
+
+    
+    def format_prompt(self):
+        prompt = f"<|im_start|>assistant\n{self.system_message}<|im_end|>\n"
+        for role, message in self.conversation_history:
+            prompt += f"<|im_start|>{role}\n{message}<|im_end|>\n"
+        prompt += "<|im_start|> assistant\n"
+        return prompt
 
     def predict(self, user_message, temperature=0.4, max_new_tokens=70, top_p=0.99, repetition_penalty=1.9):
-        prompt = f"<|im_start|>system\n{self.system_message}<|im_end|>\n<|im_start|>user\n{user_message}<|im_end|>\n<|im_start|>assistant"
+        self.update_conversation_history(user_message, "")
+        prompt = self.format_prompt()
         inputs = self.tokenizer(prompt, return_tensors='pt', add_special_tokens=False)
         input_ids = inputs["input_ids"].to(self.model.device)
 
@@ -42,7 +57,7 @@ class OrcaChatBot:
         )
 
         response = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
-
+        self.update_conversation_history("", response)
         return response
 
 Orca_bot = OrcaChatBot(model, tokenizer)