Update app.py

app.py

@@ -30,14 +30,10 @@ class OrcaChatBot:
         self.system_message = system_message
 
     def predict(self, user_message, temperature=0.4, max_new_tokens=70, top_p=0.99, repetition_penalty=1.9):
-        # Prepare the prompt
-        prompt = f"system\n{self.system_message}\nuser\n{user_message}\nassistant"
-
-        # Encode the prompt
+        prompt = f"<|im_start|>system\n{self.system_message}<|im_end|>\n<|im_start|>user\n{user_message}<|im_end|>\n<|im_start|>assistant" if self.conversation_history is None else self.conversation_history + f"<|im_end|>\n<|im_start|>user\n{user_message}<|im_end|>\n<|im_start|>assistant"
         inputs = self.tokenizer(prompt, return_tensors='pt', add_special_tokens=False)
         input_ids = inputs["input_ids"].to(self.model.device)
 
-        # Generate a response
         output_ids = self.model.generate(
             input_ids,
             max_length=input_ids.shape[1] + max_new_tokens,
@@ -45,10 +41,9 @@ class OrcaChatBot:
             top_p=top_p,
             repetition_penalty=repetition_penalty,
             pad_token_id=self.tokenizer.eos_token_id,
-            do_sample=True  # Enable sampling-based generation
+            do_sample=True 
         )
 
-        # Decode the generated response
         response = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
 
         return response
@@ -75,5 +70,4 @@ iface = gr.Interface(
     theme="ParityError/Anime"
 )
 
-# Launch the Gradio interface
 iface.launch()