Update app.py

app.py

@@ -60,34 +60,23 @@ class FalconChatBot:
                     filtered_history.append({"user": user_message, "assistant": assistant_message})
         return filtered_history
 
-    def predict(self, system_prompt, user_message, assistant_message, history, temperature=temperature, max_new_tokens=max_new_tokens, top_p=top_p, repetition_penalty=repetition_penalty):
+    def predict(self, user_message, assistant_message, history, temperature=temperature, max_new_tokens=max_new_tokens, top_p=top_p, repetition_penalty=repetition_penalty):
 
         # Process the history to remove special commands
-            processed_history = self.process_history(history)
-
+        processed_history = self.process_history(history)
         # Combine the user and assistant messages into a conversation
-            conversation = f"{system_prompt}\nFalcon: {assistant_message if assistant_message else ''} User: {user_message}\nFalcon:\n"
-    
+        conversation = f"{self.system_prompt}\nFalcon: {assistant_message if assistant_message else ''} User: {user_message}\nFalcon:\n"
         # Encode the conversation using the tokenizer
-            input_ids = tokenizer.encode(conversation, return_tensors="pt", add_special_tokens=False)
-
+        input_ids = tokenizer.encode(conversation, return_tensors="pt", add_special_tokens=False)
         # Generate a response using the Falcon model
-            response_text = peft_model.generate(input_ids=input_ids, max_length=max_length, use_cache=True, early_stopping=True, bos_token_id=peft_model.config.bos_token_id, eos_token_id=peft_model.config.eos_token_id, pad_token_id=peft_model.config.eos_token_id, temperature=0.4, do_sample=True)
-
-        # Generate the formatted conversation in Falcon message format
-            conversation = f"{system_prompt}\n"
-            for message in processed_history:
-                user_message = message["user"]
-                assistant_message = message["assistant"]
-                conversation += f"Falcon:{' ' + assistant_message if assistant_message else ''} User: {user_message}\n Falcon:\n"
+        response = peft_model.generate(input_ids=input_ids, max_length=max_length, use_cache=False, early_stopping=False, bos_token_id=peft_model.config.bos_token_id, eos_token_id=peft_model.config.eos_token_id, pad_token_id=peft_model.config.eos_token_id, temperature=0.4, do_sample=True)
         # Decode the generated response to text
-            response_text = tokenizer.decode(response[0], skip_special_tokens=True)
-
+        response_text = tokenizer.decode(response[0], skip_special_tokens=True)
         # Append the Falcon-like conversation to the history
-            self.history.append(formatted_input)
-            self.history.append(response_text)
-
-            return response_text
+        self.history.append(conversation)
+        self.history.append(response_text)
+         
+        return response_text
 
 
 # Create the Falcon chatbot instance
@@ -99,17 +88,11 @@ description = "You can use this Space to test out the GaiaMiniMed model [(Tonic/
 
 examples = [  
     [  
-        "Assistant is a public health and medical expert ready to help the user.",   
+        "Assistant is a public health and medical expert named Gaia ready to help the user.",   
         "Hi there, I have a question!",   
         "My name is Gaia, I'm a health and sanitation expert ready to answer your medical questions.",   
-        "", 0.4, 500, 0.94, 1.9  
-    ],    
-    [  
-        "Assistant is a public health and medical expert ready to help the user.",   
-        "What is the proper treatment for buccal herpes?",   
-        None,   
-        "", 0.2, 500, 0.92, 1.7  
-    ]  
+        "Assistant is a medical and sanitation question expert trained to answer medical questions", 0.4, 700, 0.90, 1.9  
+    ]
 ]