Update app.py

app.py

@@ -70,23 +70,31 @@ peft_config = PeftConfig.from_pretrained("Tonic/mistralmed", token="hf_dQUWWpJJy
 peft_model = MistralForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1", trust_remote_code=True)
 peft_model = PeftModel.from_pretrained(peft_model, "Tonic/mistralmed", token="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")
 # Remove the memory function
+# ... (previous code)
+
+class ChatBot:
+    def __init__(self):
+        self.history = []
+
 class ChatBot:
     def __init__(self):
+        # Initialize the ChatBot class with an empty history
         self.history = []
 
     def predict(self, user_input, system_prompt="You are an expert medical analyst:"):
-        # Combine user input and system prompt
+        # Combine the user's input with the system prompt
         formatted_input = f"<s>[INST]{system_prompt} {user_input}[/INST]"
 
-        # Encode user input
+        # Encode the formatted input using the tokenizer
         user_input_ids = tokenizer.encode(formatted_input, return_tensors="pt")
 
         # Generate a response using the PEFT model
-        response = peft_model.generate(input_ids, max_length=512, pad_token_id=tokenizer.eos_token_id)
+        response = peft_model.generate(user_input_ids, max_length=512, pad_token_id=tokenizer.eos_token_id)
 
-        # Decode and return the response
+        # Decode the generated response to text
         response_text = tokenizer.decode(response[0], skip_special_tokens=True)
-        return response_text
+        
+        return response_text  # Return the generated response
 
 bot = ChatBot()