Update app.py

app.py

@@ -19,10 +19,10 @@ def multimodal_prompt(user_input, system_prompt="You are an expert medical analy
         user_input: The user's input text to generate a response for.
         system_prompt: Optional system prompt.
     Returns:
-        A string containing the generated text.
+        A string containing the generated text in the Falcon-like format.
     """
     # Combine user input and system prompt
-    formatted_input = f"Falcon: {system_prompt} \n User: {user_input} \n Falcon:"
+    formatted_input = f"{{{{ {system_prompt} }}}}\nUser: {user_input}\nFalcon:"
 
     # Encode the input text
     encodeds = tokenizer(formatted_input, return_tensors="pt", add_special_tokens=False)
@@ -31,13 +31,13 @@ def multimodal_prompt(user_input, system_prompt="You are an expert medical analy
     # Generate a response using the model
     output = model.generate(
         **model_inputs,
-        max_length=max_length,
+        max_length=500,
         use_cache=True,
         early_stopping=True,
         bos_token_id=model.config.bos_token_id,
         eos_token_id=model.config.eos_token_id,
         pad_token_id=model.config.eos_token_id,
-        temperature=0.1,
+        temperature=0.4,
         do_sample=True
     )
 
@@ -64,7 +64,8 @@ tokenizer.padding_side = 'left'
 peft_config = PeftConfig.from_pretrained("Tonic/GaiaMiniMed")
 
 # Using Optimum
-peft_model.to_bettertransformer()
+model.to_bettertransformer()
+
 peft_model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-7b-instruct")
 peft_model = PeftModel.from_pretrained(peft_model, "Tonic/GaiaMiniMed")
 
@@ -80,28 +81,28 @@ peft_model = PeftModel.from_pretrained(peft_model, "Tonic/GaiaMiniMed")
 # peft_model = PeftModel.from_pretrained(peft_model, "Tonic/mistralmed", token="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")
 
 class ChatBot:
-    def __init__(self):
-        self.history = []
-
-class ChatBot:
-    def __init__(self):
-        # Initialize the ChatBot class with an empty history
+    def __init__(self, system_prompt="You are an expert medical analyst:"):
+        self.system_prompt = system_prompt
         self.history = []
 
-    def predict(self, user_input, system_prompt="You are an expert medical analyst:"):
-        # Combine the user's input with the system prompt
-        formatted_input = f"{{{ {system_prompt} }}} \n User: {user_input} \n Falcon:"
+    def predict(self, user_input):
+        # Combine the user's input with the system prompt in Falcon format
+        formatted_input = f"{{{{ {self.system_prompt} }}}}\nUser: {user_input}\nFalcon:"
 
         # Encode the formatted input using the tokenizer
-        user_input_ids = tokenizer.encode(formatted_input, return_tensors="pt")
+        input_ids = tokenizer.encode(formatted_input, return_tensors="pt", add_special_tokens=False)
 
-        # Generate a response using the PEFT model
-        response = peft_model.generate(input_ids=user_input_ids, max_length=512, pad_token_id=tokenizer.eos_token_id)
+        # Generate a response using the model
+        response = model.generate(input_ids, max_length=max_length, use_cache=True, early_stopping=True, bos_token_id=model.config.bos_token_id, eos_token_id=model.config.eos_token_id, pad_token_id=model.config.eos_token_id, temperature=0.1, do_sample=True)
 
         # Decode the generated response to text
         response_text = tokenizer.decode(response[0], skip_special_tokens=True)
-        
-        return response_text  # Return the generated response
+
+        # Append the Falcon-like conversation to the history
+        self.history.append(formatted_input)
+        self.history.append(response_text)
+
+        return response_text
 
 bot = ChatBot()