Update app.py

app.py

@@ -53,16 +53,13 @@ tokenizer = AutoTokenizer.from_pretrained(model_id)
 
 model = AutoModelForCausalLM.from_pretrained(model_id , torch_dtype=torch.float16 , device_map= "auto" )
 
-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:"):
+    def predict(self, user_input, system_prompt="You are an expert medical analyst:" , example_instruction="produce a json", example_answer = "please dont make small talk "):
         # Combine the user's input with the system prompt
         formatted_input = f"<s> [INST] {example_instruction} [/INST] {example_answer}</s> [INST] {system_prompt} [/INST]"
 
@@ -70,7 +67,7 @@ class ChatBot:
         user_input_ids = tokenizer.encode(formatted_input, return_tensors="pt")
 
         # 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)
+        response = model.generate(input_ids=user_input_ids, max_length=512, pad_token_id=tokenizer.eos_token_id)
 
         # Decode the generated response to text
         response_text = tokenizer.decode(response[0], skip_special_tokens=True)