merge

benchmarking/llm_providers/medrax_provider.py

@@ -35,28 +35,19 @@ class MedRAXProvider(LLMProvider):
             print("Starting server...")
 
             selected_tools = [
-                # To be tested
-                # "DicomProcessorTool",  # For processing DICOM medical image files
-                # "TorchXRayVisionClassifierTool",  # For classifying chest X-ray images using TorchXRayVision
-                "ChestXRaySegmentationTool",  # For segmenting anatomical regions in chest X-rays
-                # "WebBrowserTool",  # For web browsing and search capabilities
-                
-                # These tools are working
-                # "MedicalRAGTool",  # For retrieval-augmented generation with medical knowledge
-                # "ChestXRayReportGeneratorTool",  # For generating medical reports from X-rays
-                # "XRayVQATool",  # For visual question answering on X-rays
-                        
-
-                # Couldn't test these tools
                 # "ImageVisualizerTool",  # For displaying images in the UI
-                # "PythonSandboxTool",  # Add the Python sandbox tool
+                # "DicomProcessorTool",  # For processing DICOM medical image files
+                "TorchXRayVisionClassifierTool",  # For classifying chest X-ray images using TorchXRayVision
+                "ArcPlusClassifierTool",  # For advanced chest X-ray classification using ArcPlus
+                # "ChestXRaySegmentationTool",  # For segmenting anatomical regions in chest X-rays
+                "ChestXRayReportGeneratorTool",  # For generating medical reports from X-rays
+                "XRayVQATool",  # For visual question answering on X-rays
                 # "LlavaMedTool",  # For multimodal medical image understanding
+                "XRayPhraseGroundingTool",  # For locating described features in X-rays
                 # "ChestXRayGeneratorTool",  # For generating synthetic chest X-rays
-                # "ArcPlusClassifierTool",  # For advanced chest X-ray classification using ArcPlus
-
-
-                # Something fishy is going on here
-                # "XRayPhraseGroundingTool",  # For locating described features in X-rays
+                "WebBrowserTool",  # For web browsing and search capabilities
+                "MedicalRAGTool",  # For retrieval-augmented generation with medical knowledge
+                # "PythonSandboxTool",  # Add the Python sandbox tool
             ]
 
             rag_config = RAGConfig(
@@ -67,7 +58,7 @@ class MedRAXProvider(LLMProvider):
                 pinecone_index_name="medrax2",  # Name for the Pinecone index
                 chunk_size=1500,
                 chunk_overlap=300,
-                retriever_k=7,
+                retriever_k=3,
                 local_docs_dir="rag_docs",  # Change this to the path of the documents for RAG
                 huggingface_datasets=["VictorLJZ/medrax2"],  # List of HuggingFace datasets to load
                 dataset_split="train",  # Which split of the datasets to use
@@ -79,11 +70,11 @@ class MedRAXProvider(LLMProvider):
             agent, tools_dict = initialize_agent(
                 prompt_file="medrax/docs/system_prompts.txt",
                 tools_to_use=selected_tools,
-                model_dir="model-weights",
+                model_dir="/model-weights",
                 temp_dir="temp",  # Change this to the path of the temporary directory
-                device="cpu",
+                device="cuda:0",
                 model=self.model_name,  # Change this to the model you want to use, e.g. gpt-4.1-2025-04-14, gemini-2.5-pro
-                temperature=0.7,
+                temperature=0.3,
                 top_p=0.95,
                 model_kwargs=model_kwargs,
                 rag_config=rag_config,

benchmarking/runner.py

@@ -262,9 +262,15 @@ class BenchmarkRunner:
         Returns:
             str: The extracted answer
         """
-        # First, look for the '<|A|>' format
-        final_answer_pattern = r'\s*<\|([A-F])\|>'
-        match = re.search(final_answer_pattern, response_text)
+        # Look for the '\boxed{A}' format
+        boxed_pattern = r'\\boxed\{([A-Fa-f])\}'
+        match = re.search(boxed_pattern, response_text)
+        if match:
+            return match.group(1).upper()
+        
+        # Fallback: look for the '<|A|>' format (legacy code, will remove later on)
+        legacy_pattern = r'\s*<\|([A-F])\|>'
+        match = re.search(legacy_pattern, response_text)
         if match:
             return match.group(1).upper()
         

main.py

@@ -143,15 +143,15 @@ if __name__ == "__main__":
     selected_tools = [
         "ImageVisualizerTool",  # For displaying images in the UI
         # "DicomProcessorTool",  # For processing DICOM medical image files
-        # "TorchXRayVisionClassifierTool",  # For classifying chest X-ray images using TorchXRayVision
-        # "ArcPlusClassifierTool",  # For advanced chest X-ray classification using ArcPlus
-        # "ChestXRaySegmentationTool",  # For segmenting anatomical regions in chest X-rays
-        # "ChestXRayReportGeneratorTool",  # For generating medical reports from X-rays
-        # "XRayVQATool",  # For visual question answering on X-rays
+        "TorchXRayVisionClassifierTool",  # For classifying chest X-ray images using TorchXRayVision
+        "ArcPlusClassifierTool",  # For advanced chest X-ray classification using ArcPlus
+        "ChestXRaySegmentationTool",  # For segmenting anatomical regions in chest X-rays
+        "ChestXRayReportGeneratorTool",  # For generating medical reports from X-rays
+        "XRayVQATool",  # For visual question answering on X-rays
         # "LlavaMedTool",  # For multimodal medical image understanding
-        # "XRayPhraseGroundingTool",  # For locating described features in X-rays
+        "XRayPhraseGroundingTool",  # For locating described features in X-rays
         # "ChestXRayGeneratorTool",  # For generating synthetic chest X-rays
-        "MedSAM2Tool",  # For advanced medical image segmentation using MedSAM2
+        # "MedSAM2Tool",  # For advanced medical image segmentation using MedSAM2
         "WebBrowserTool",  # For web browsing and search capabilities
         "MedicalRAGTool",  # For retrieval-augmented generation with medical knowledge
         # "PythonSandboxTool",  # Add the Python sandbox tool
@@ -167,7 +167,7 @@ if __name__ == "__main__":
         pinecone_index_name="medrax2",  # Name for the Pinecone index
         chunk_size=1500,
         chunk_overlap=300,
-        retriever_k=7,
+        retriever_k=3,
         local_docs_dir="rag_docs",  # Change this to the path of the documents for RAG
         huggingface_datasets=["VictorLJZ/medrax2"],  # List of HuggingFace datasets to load
         dataset_split="train",  # Which split of the datasets to use
@@ -179,10 +179,10 @@ if __name__ == "__main__":
     agent, tools_dict = initialize_agent(
         prompt_file="medrax/docs/system_prompts.txt",
         tools_to_use=selected_tools,
-        model_dir="model-weights",
+        model_dir="/model-weights",
         temp_dir="temp",  # Change this to the path of the temporary directory
-        device="cuda",
-        model="grok-4",  # Change this to the model you want to use, e.g. gpt-4.1-2025-04-14, gemini-2.5-pro
+        device="cuda:0",
+        model="gpt-4.1-2025-04-14",  # Change this to the model you want to use, e.g. gpt-4.1-2025-04-14, gemini-2.5-pro
         temperature=0.7,
         top_p=0.95,
         model_kwargs=model_kwargs,

medrax/docs/system_prompts.txt

@@ -22,5 +22,5 @@ Solve using your own vision and reasoning and use tools (if available) to comple
 You can make multiple tool calls in parallel or in sequence as needed for comprehensive answers.
 Think critically about and criticize the tool outputs.
 If you need to look up some information before asking a follow up question, you are allowed to do that.
-When encountering a multiple-choice question, your final response should end with "Final answer: <|A|>" from list of possible choices A, B, C, D, E, F.
+When encountering a multiple-choice question, your final response should end with "Final answer: \boxed{A}" from list of possible choices A, B, C, D, E, F.
 It is extremely important that you strictly answer in the format mentioned above.