Update app.py

app.py

@@ -18,27 +18,16 @@ class OrcaChatBot:
     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 = []
+        self.default_system_message = system_message
 
-    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:
-            if message.strip():
-                prompt += f"<|im_start|>{role}\n{message}<|im_end|>\n"
-#               if role == "assistant":
-#                    prompt += f"<|im_end|>\n"
-        prompt += "<|im_start|> assistant\n"
+    def format_prompt(self, user_message, system_message):
+        if system_message is None:
+            system_message = self.default_system_message
+        prompt = f"<|im_start|>assistant\n{self.system_message}<|im_end|>\n<|im_start|>\nuser\n{user_message}<|im_end|>\nassistant\n"
         return prompt
 
-    def predict(self, user_message, temperature=0.4, max_new_tokens=70, top_p=0.99, repetition_penalty=1.9):
-        self.update_conversation_history(user_message, "")
-        prompt = self.format_prompt()
+    def predict(self, user_message, system_message=None, temperature=0.4, max_new_tokens=70, top_p=0.99, repetition_penalty=1.9):
+        prompt = self.format_prompt(user_message, system_message)
         inputs = self.tokenizer(prompt, return_tensors='pt', add_special_tokens=False)
         input_ids = inputs["input_ids"].to(self.model.device)
 
@@ -48,19 +37,17 @@ class OrcaChatBot:
             temperature=temperature,
             top_p=top_p,
             repetition_penalty=repetition_penalty,
-#           pad_token_id=self.tokenizer.eos_token_id,
             do_sample=True
-    )
+        )
 
         response = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
-        self.update_conversation_history("", response)
         return response
-        
-Orca_bot = OrcaChatBot(model, tokenizer)
 
 def gradio_predict(user_message, system_message, max_new_tokens, temperature, top_p, repetition_penalty):
-    full_message = f"{system_message}\n{user_message}" if system_message else user_message
-    return Orca_bot.predict(full_message, temperature, max_new_tokens, top_p, repetition_penalty)
+    response = Orca_bot.predict(user_message, system_message, temperature, max_new_tokens, top_p, repetition_penalty)
+    return response
+    
+Orca_bot = OrcaChatBot(model, tokenizer)
 
 iface = gr.Interface(
     fn=gradio_predict,