included new user and updated buttons

interface.py

@@ -16,7 +16,8 @@ REPORT_DIR.mkdir(exist_ok=True)
 SALT = b'$2b$12$MC7djiqmIR7154Syul5Wme'
 
 USERS = {
-    'test_user': b'$2b$12$MC7djiqmIR7154Syul5WmeQwebwsNOK5svMX08zMYhvpF9P9IVXe6'
+    'test_user': b'$2b$12$MC7djiqmIR7154Syul5WmeQwebwsNOK5svMX08zMYhvpF9P9IVXe6',
+    'pna': b'$2b$12$MC7djiqmIR7154Syul5WmeWTzYft1UnOV4uGVn54FGfmbH3dRNq1C'
 }
 
 class ChatInterface:
@@ -73,7 +74,7 @@ class ChatInterface:
         else:
             self.display_file_path = str(saved_path)
 
-        return self.display_file_path, gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=True)
+        return self.display_file_path, gr.update(interactive=True), gr.update(interactive=True)
 
     def add_message(
         self, message: str, display_image: str, history: List[dict]
@@ -266,7 +267,7 @@ def create_demo(agent, tools_dict):
                                 )
                             with gr.Row():
                                 analyze_btn = gr.Button("Analyze", interactive=False)
-                                ground_btn = gr.Button("Ground", interactive=False)
+                                # ground_btn = gr.Button("Ground", interactive=False)
                                 segment_btn = gr.Button("Segment", interactive=False)
                             with gr.Row():
                                 clear_btn = gr.Button("Clear Chat")
@@ -394,38 +395,38 @@ def create_demo(agent, tools_dict):
         bot_msg.then(lambda: gr.Textbox(interactive=True), None, [txt])
 
         analyze_btn.click(
-            lambda: gr.update(value="Analyze the above image and identify the probabilites of occurrence of various diseases along with proper reason"), None, txt
+            lambda: gr.update(value="Analyze this xray and give me a detailed response. Use the medgemma_xray_expert tool"), None, txt
         ).then(
-            interface.add_message, inputs=[txt, image_display, chatbot], outputs=[chatbot, txt]                         # add message & clear box
+            interface.add_message, inputs=[txt, image_display, chatbot], outputs=[chatbot, txt]
         ).then(
             interface.process_message,
             inputs=[txt, image_display, chatbot],
             outputs=[chatbot, image_display, txt],
         ).then(lambda: gr.Textbox(interactive=True), None, [txt])
         
-        ground_btn.click(
-            lambda: gr.update(value="Ground the main disease in this CXR"), None, txt
-        ).then(
-            interface.add_message, inputs=[txt, image_display, chatbot], outputs=[chatbot, txt]                         # add message & clear box
-        ).then(
-            interface.process_message,
-            inputs=[txt, image_display, chatbot],
-            outputs=[chatbot, image_display, txt],
-        ).then(lambda: gr.Textbox(interactive=True), None, [txt])
+        # ground_btn.click(
+        #     lambda: gr.update(value="Ground the main disease in this CXR"), None, txt
+        # ).then(
+        #     interface.add_message, inputs=[txt, image_display, chatbot], outputs=[chatbot, txt]  
+        # ).then(
+        #     interface.process_message,
+        #     inputs=[txt, image_display, chatbot],
+        #     outputs=[chatbot, image_display, txt],
+        # ).then(lambda: gr.Textbox(interactive=True), None, [txt])
 
         segment_btn.click(
             lambda: gr.update(value="Segment the major affected lung"), None, txt
         ).then(
-            interface.add_message, inputs=[txt, image_display, chatbot], outputs=[chatbot, txt]                         # add message & clear box
+            interface.add_message, inputs=[txt, image_display, chatbot], outputs=[chatbot, txt] 
         ).then(
             interface.process_message,
             inputs=[txt, image_display, chatbot],
             outputs=[chatbot, image_display, txt],
         ).then(lambda: gr.Textbox(interactive=True), None, [txt])
 
-        upload_button.upload(handle_file_upload, inputs=upload_button, outputs=[image_display, analyze_btn, ground_btn, segment_btn])
+        upload_button.upload(handle_file_upload, inputs=upload_button, outputs=[image_display, analyze_btn, segment_btn])
 
-        dicom_upload.upload(handle_file_upload, inputs=dicom_upload, outputs=[image_display, analyze_btn, ground_btn, segment_btn])
+        dicom_upload.upload(handle_file_upload, inputs=dicom_upload, outputs=[image_display, analyze_btn, segment_btn])
 
         clear_btn.click(clear_chat, outputs=[chatbot, image_display])
         new_thread_btn.click(new_thread, outputs=[chatbot, image_display])

medrax/tools/llava_med.py

@@ -56,6 +56,7 @@ class LlavaMedTool(BaseTool):
     def __init__(
         self,
         model_path: str = "microsoft/llava-med-v1.5-mistral-7b",
+        # model_path: str = "microsoft/llava-rad",
         cache_dir: str = "/model-weights",
         low_cpu_mem_usage: bool = True,
         torch_dtype: torch.dtype = torch.bfloat16,
@@ -68,6 +69,7 @@ class LlavaMedTool(BaseTool):
         self.tokenizer, self.model, self.image_processor, self.context_len = load_pretrained_model(
             model_path=model_path,
             model_base=None,
+            # model_base="lmsys/vicuna-7b-v1.5",
             model_name=model_path,
             load_in_4bit=load_in_4bit,
             load_in_8bit=load_in_8bit,