Update app.py

app.py

@@ -18,16 +18,17 @@ 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.default_system_message = system_message
+        self.system_message = system_message
 
-    def format_prompt(self, user_message, system_message):
-        if system_message is None:
-            system_message = self.default_system_message
+    def set_system_message(self, new_system_message):
+        self.system_message = new_system_message
+
+    def format_prompt(self, user_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, 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)
+    def predict(self, user_message, temperature=0.4, max_new_tokens=70, top_p=0.99, repetition_penalty=1.9):
+        prompt = self.format_prompt(user_message)
         inputs = self.tokenizer(prompt, return_tensors='pt', add_special_tokens=False)
         input_ids = inputs["input_ids"].to(self.model.device)
 
@@ -44,7 +45,8 @@ class OrcaChatBot:
         return response
 
 def gradio_predict(user_message, system_message, max_new_tokens, temperature, top_p, repetition_penalty):
-    response = Orca_bot.predict(user_message, system_message, temperature, max_new_tokens, top_p, repetition_penalty)
+    Orca_bot.set_system_message(system_message)
+    response = Orca_bot.predict(user_message, temperature, max_new_tokens, top_p, repetition_penalty)
     return response
     
 Orca_bot = OrcaChatBot(model, tokenizer)