Update app.py

app.py

@@ -323,42 +323,24 @@ def wrap_text(text, width=90):
 tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen-1_8B-Chat", trust_remote_code=True) # TruEra
 model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-1_8B-Chat", device_map="auto", trust_remote_code=True).eval()
 
-def multimodal_prompt(user_input, system_prompt="You are an expert medical analyst:"): # TruEra
-    formatted_input = f"{user_input}{system_prompt}"
-
-    # Encode the input text
-    encoded_input = tokenizer(formatted_input, return_tensors="pt").to(device)
-
-    # Generate a response using the model
-    output = model.generate(
-        **encoded_input,
-        max_length=512,
-        use_cache=True,
-        early_stopping=True,
-        pad_token_id=tokenizer.eos_token_id,
-        temperature=0.1,
-        do_sample=True
-    )
-
-    response_text = tokenizer.decode(output[0], skip_special_tokens=True)
-
-    return response_text
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen-1_8B-Chat", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-1_8B-Chat", device_map="auto", trust_remote_code=True).eval()
 
+# TruEra
 class ChatBot:
     def __init__(self):
-        self.history = []
-        
-    @staticmethod
-    def doctor(user_input, system_prompt="You are an expert medical analyst:"): # TruEra
-        formatted_input = f"{system_prompt}{user_input}"
-        user_input_ids = tokenizer.encode(formatted_input, return_tensors="pt")
-        response = model.generate(input_ids=user_input_ids, max_length=512, pad_token_id=tokenizer.eos_token_id)
-        response_text = tokenizer.decode(response[0], skip_special_tokens=True)
-        return response_text
+        self.history = None
 
+    def predict(self, user_input, system_prompt=""):
+        response, self.history = model.chat(tokenizer, user_input, history=self.history, system=system_prompt)
+        return response
 
 bot = ChatBot()
 
+# TruEra
+def multimodal_prompt(user_input, system_prompt="You are an expert medical analyst:"):
+    return bot.predict(user_input, system_prompt)
+
 
 def process_summary_with_qwen(summary): # TruEra
     system_prompt = "You are a medical instructor . Assess and describe the proper options to your students in minute detail. Propose a course of action for them to base their recommendations on based on your description."