initial commit

.gitignore

@@ -0,0 +1,174 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# ruff
+ruff-cache/
+.ruff_cache/
+
+afallah/
+
+logs/
+
+temp/
+
+.gradio/

.vscode/launch.json

@@ -0,0 +1,15 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Python Debugger: main.py",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "main.py",
+            "console": "integratedTerminal"
+        }
+    ]
+}

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,12 +1,267 @@
----
-title: CRAX
-emoji: 🐢
-colorFrom: pink
-colorTo: green
-sdk: gradio
-sdk_version: 5.25.2
-app_file: app.py
-pinned: false
----
+<h1 align="center">
+🤖 MedRAX: Medical Reasoning Agent for Chest X-ray
+</h1>
+<p align="center"> <a href="https://arxiv.org/abs/2502.02673" target="_blank"><img src="https://img.shields.io/badge/arXiv-Paper-FF6B6B?style=for-the-badge&logo=arxiv&logoColor=white" alt="arXiv"></a> <a href="https://github.com/bowang-lab/MedRAX"><img src="https://img.shields.io/badge/GitHub-Code-4A90E2?style=for-the-badge&logo=github&logoColor=white" alt="GitHub"></a> <a href="https://huggingface.co/datasets/wanglab/chest-agent-bench"><img src="https://img.shields.io/badge/HuggingFace-Dataset-FFBF00?style=for-the-badge&logo=huggingface&logoColor=white" alt="HuggingFace Dataset"></a> </p>
+
+![](assets/demo_fast.gif?autoplay=1)
+
+<br>
+
+## Abstract
+Chest X-rays (CXRs) play an integral role in driving critical decisions in disease management and patient care. While recent innovations have led to specialized models for various CXR interpretation tasks, these solutions often operate in isolation, limiting their practical utility in clinical practice. We present MedRAX, the first versatile AI agent that seamlessly integrates state-of-the-art CXR analysis tools and multimodal large language models into a unified framework. MedRAX dynamically leverages these models to address complex medical queries without requiring additional training. To rigorously evaluate its capabilities, we introduce ChestAgentBench, a comprehensive benchmark containing 2,500 complex medical queries across 7 diverse categories. Our experiments demonstrate that MedRAX achieves state-of-the-art performance compared to both open-source and proprietary models, representing a significant step toward the practical deployment of automated CXR interpretation systems.
+<br><br>
+
+
+## MedRAX
+MedRAX is built on a robust technical foundation:
+- **Core Architecture**: Built on LangChain and LangGraph frameworks
+- **Language Model**: Uses GPT-4o with vision capabilities as the backbone LLM
+- **Deployment**: Supports both local and cloud-based deployments
+- **Interface**: Production-ready interface built with Gradio
+- **Modular Design**: Tool-agnostic architecture allowing easy integration of new capabilities
+
+### Integrated Tools
+- **Visual QA**: Utilizes CheXagent and LLaVA-Med for complex visual understanding and medical reasoning
+- **Segmentation**: Employs MedSAM and PSPNet model trained on ChestX-Det for precise anatomical structure identification
+- **Grounding**: Uses Maira-2 for localizing specific findings in medical images
+- **Report Generation**: Implements SwinV2 Transformer trained on CheXpert Plus for detailed medical reporting
+- **Disease Classification**: Leverages DenseNet-121 from TorchXRayVision for detecting 18 pathology classes
+- **X-ray Generation**: Utilizes RoentGen for synthetic CXR generation
+- **Utilities**: Includes DICOM processing, visualization tools, and custom plotting capabilities
+
+Note the current version of MedRAX is experimentally released and does not support vision for GPT-4o and MedSAM. We will be integrating these shortly.
+<br><br>
+
+
+## ChestAgentBench
+We introduce ChestAgentBench, a comprehensive evaluation framework with 2,500 complex medical queries across 7 categories, built from 675 expert-curated clinical cases. The benchmark evaluates complex multi-step reasoning in CXR interpretation through:
+
+- Detection
+- Classification
+- Localization
+- Comparison
+- Relationship
+- Diagnosis
+- Characterization
+
+Download the benchmark: [ChestAgentBench on Hugging Face](https://huggingface.co/datasets/wanglab/chest-agent-bench)
+```
+huggingface-cli download wanglab/chestagentbench --repo-type dataset --local-dir chestagentbench
+```
+
+Unzip the Eurorad figures to your local `MedMAX` directory.
+```
+unzip chestagentbench/figures.zip
+```
+
+To evaluate with GPT-4o, set your OpenAI API key and run the quickstart script.
+```
+export OPENAI_API_KEY="<your-openai-api-key>"
+python quickstart.py \
+    --model chatgpt-4o-latest \
+    --temperature 0.2 \
+    --max-cases 2 \
+    --log-prefix chatgpt-4o-latest \
+    --use-urls
+```
+
+
+<br>
+
+## Installation
+### Prerequisites
+- Python 3.8+
+- CUDA/GPU for best performance
+
+### Installation Steps
+```bash
+# Clone the repository
+git clone https://github.com/bowang-lab/MedRAX.git
+cd MedRAX
+
+# Install package
+pip install -e .
+```
+
+### Getting Started
+```bash
+# Start the Gradio interface
+python main.py
+```
+or if you run into permission issues
+```bash
+sudo -E env "PATH=$PATH" python main.py
+```
+You need to setup the `model_dir` inside `main.py` to the directory where you want to download or already have the weights of above tools from Hugging Face.
+Comment out the tools that you do not have access to.
+Make sure to setup your OpenAI API key in `.env` file!
+<br><br><br>
+
+
+## Tool Selection and Initialization
+
+MedRAX supports selective tool initialization, allowing you to use only the tools you need. Tools can be specified when initializing the agent (look at `main.py`):
+
+```python
+selected_tools = [
+    "ImageVisualizerTool",
+    "ChestXRayClassifierTool",
+    "ChestXRaySegmentationTool",
+    # Add or remove tools as needed
+]
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+agent, tools_dict = initialize_agent(
+    "medrax/docs/system_prompts.txt",
+    tools_to_use=selected_tools,
+    model_dir="/model-weights"
+)
+```
+
+<br><br>
+## Automatically Downloaded Models
+
+The following tools will automatically download their model weights when initialized:
+
+### Classification Tool
+```python
+ChestXRayClassifierTool(device=device)
+```
+
+### Segmentation Tool
+```python
+ChestXRaySegmentationTool(device=device)
+```
+
+### Grounding Tool
+```python
+XRayPhraseGroundingTool(
+    cache_dir=model_dir, 
+    temp_dir=temp_dir, 
+    load_in_8bit=True, 
+    device=device
+)
+```
+- Maira-2 weights download to specified `cache_dir`
+- 8-bit and 4-bit quantization available for reduced memory usage
+
+### LLaVA-Med Tool
+```python
+LlavaMedTool(
+    cache_dir=model_dir, 
+    device=device, 
+    load_in_8bit=True
+)
+```
+- Automatic weight download to `cache_dir`
+- 8-bit and 4-bit quantization available for reduced memory usage
+
+### Report Generation Tool
+```python
+ChestXRayReportGeneratorTool(
+    cache_dir=model_dir, 
+    device=device
+)
+```
+
+### Visual QA Tool
+```python
+XRayVQATool(
+    cache_dir=model_dir, 
+    device=device
+)
+```
+- CheXagent weights download automatically
+
+### MedSAM Tool
+```
+Support for MedSAM segmentation will be added in a future update.
+```
+
+### Utility Tools
+No additional model weights required:
+```python
+ImageVisualizerTool()
+DicomProcessorTool(temp_dir=temp_dir)
+```
+<br>
+
+## Manual Setup Required
+
+### Image Generation Tool
+```python
+ChestXRayGeneratorTool(
+    model_path=f"{model_dir}/roentgen", 
+    temp_dir=temp_dir, 
+    device=device
+)
+```
+- RoentGen weights require manual setup:
+  1. Contact authors: https://github.com/StanfordMIMI/RoentGen
+  2. Place weights in `{model_dir}/roentgen`
+  3. Optional tool, can be excluded if not needed
+<br>
+
+## Configuration Notes
+
+### Required Parameters
+- `model_dir` or `cache_dir`: Base directory for model weights that Hugging Face uses
+- `temp_dir`: Directory for temporary files
+- `device`: "cuda" for GPU, "cpu" for CPU-only
+
+### Memory Management
+- Consider selective tool initialization for resource constraints
+- Use 8-bit quantization where available
+- Some tools (LLaVA-Med, Grounding) are more resource-intensive
+<br>
+
+### Local LLMs
+If you are running a local LLM using frameworks like [Ollama](https://ollama.com/) or [LM Studio](https://lmstudio.ai/), you need to configure your environment variables accordingly. For example:
+```
+export OPENAI_BASE_URL="http://localhost:11434/v1"
+export OPENAI_API_KEY="ollama"
+```
+<br>
+
+## Star History
+<div align="center">
+  
+[![Star History Chart](https://api.star-history.com/svg?repos=bowang-lab/MedRAX&type=Date)](https://star-history.com/#bowang-lab/MedRAX&Date)
+
+</div>
+<br>
+
+
+## Authors
+- **Adibvafa Fallahpour**¹²³ * ([email protected])
+- **Jun Ma**²³ * 
+- **Alif Munim**³⁴ *
+- **Hongwei Lyu**³
+- **Bo Wang**¹²³⁵
+
+¹ Department of Computer Science, University of Toronto, Toronto, Canada  
+² Vector Institute, Toronto, Canada  
+³ University Health Network, Toronto, Canada  
+⁴ Cohere For AI, Toronto, Canada  
+⁵ Department of Laboratory Medicine and Pathobiology, University of Toronto, Toronto, Canada <br>
+\* Equal contribution
+<br><br>
+
+
+## Citation
+If you find this work useful, please cite our paper:
+```bibtex
+@misc{fallahpour2025medraxmedicalreasoningagent,
+      title={MedRAX: Medical Reasoning Agent for Chest X-ray}, 
+      author={Adibvafa Fallahpour and Jun Ma and Alif Munim and Hongwei Lyu and Bo Wang},
+      year={2025},
+      eprint={2502.02673},
+      archivePrefix={arXiv},
+      primaryClass={cs.LG},
+      url={https://arxiv.org/abs/2502.02673}, 
+}
+```
+
+---
+<p align="center">
+Made with ❤️ at University of Toronto, Vector Institute, and University Health Network
+</p>

app.py

@@ -1,7 +0,0 @@
-import gradio as gr
-
-def greet(name):
-    return "Hello " + name + "!!"
-
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()

benchmark/create_benchmark.py

@@ -0,0 +1,352 @@
+#!/usr/bin/env python3
+"""
+Medical X-ray Question Generation Benchmark aka ChestAgentBench
+
+This script generates clinical questions from X-ray case data of Eurorad dataset using GPT-4o.
+It structures questions across different analytical categories and saves them as JSON.
+"""
+
+import os
+import re
+import json
+from typing import *
+from pprint import pprint
+
+import openai
+import numpy as np
+from scipy import stats
+import plotly.graph_objects as go
+from tqdm import tqdm
+
+from benchmark.utils import load_eurorad_dataset
+from benchmark.llm import get_llm_response
+
+# Constants
+DATA_DIR = "set your data directory here, e.g. /home/MedRAX/data"
+DATASET_PATH = os.path.join(DATA_DIR, "eurorad_metadata.json")
+
+SYSTEM_PROMPT = """
+You are an expert medical benchmark creation assistant.
+Your goal is to generate questions that evaluate a multimodal medical AI agent's ability to interpret and reason about chest X-rays.
+""".strip()
+
+CATEGORIES_META = {
+    "detection": "Identify and locate specific findings in the chest X-ray.",
+    "classification": "Determine whether specific findings are present or absent in the chest X-ray.",
+    "enumeration": "Count the number of target findings in the chest X-ray.",
+    "localization": "Locate a given finding in the chest X-ray.",
+    "comparison": "Compare the size or position of a specific finding in the chest X-ray.",
+    "relationship": "Determine the relationship between two or more findings in the chest X-ray.",
+    "diagnosis": "Make a diagnosis or determine a treatment plan by interpreting the chest X-ray.",
+    "characterization": "Describe specific attributes (shape, density, margins, etc.) of findings.",
+    "reasoning": "Explain the medical rationale and thought process behind findings and conclusions.",
+}
+CATEGORIES = list(CATEGORIES_META.keys())
+
+CATEGORY_COMBINATIONS = [
+    ["detection", "localization", "characterization", "reasoning"],  # Detailed Finding Analysis
+    ["detection", "classification", "relationship", "reasoning"],  # Pattern Recognition & Relations
+    ["localization", "comparison", "relationship", "reasoning"],  # Spatial Understanding
+    ["classification", "comparison", "diagnosis", "reasoning"],  # Clinical Decision Making
+    ["classification", "characterization", "diagnosis", "reasoning"],  # Diagnostic Characterization
+]
+
+DEFAULT_SECTIONS = [
+    "history",
+    "image_finding",
+    "discussion",
+    "differential_diagnosis",
+    "diagnosis",
+    "figures",
+]
+
+
+class Question:
+    """A class to generate clinical questions from case data.
+
+    This class handles creating structured clinical questions by combining case data with
+    specified categories and difficulty levels.
+
+    Attributes:
+        type (str): The type of question (e.g. multiple choice)
+        difficulty (str): Difficulty level of the question
+        case_data (Dict[str, Any]): Dictionary containing the clinical case data
+        case_content (str): Formatted case data from selected sections
+        case_id (str): Unique identifier for the case
+        categories (List[str]): List of analytical categories this question tests
+        sections (List[str]): Case sections to include in question
+        raw_content (Optional[str]): Raw LLM response to the question prompt
+        content (Optional[Dict[str, str]]): Extracted content from the raw LLM response
+    """
+
+    def __init__(
+        self,
+        type: str,
+        difficulty: str,
+        case_data: Dict[str, Any],
+        categories: List[str],
+        sections: List[str] = [
+            "history",
+            "image_finding",
+            "discussion",
+            "differential_diagnosis",
+            "diagnosis",
+            "figures",
+        ],
+        system_prompt: str = "You are an expert medical benchmark creation assistant.",
+    ) -> None:
+        self.type = type
+        self.difficulty = difficulty
+        self.case_data = case_data
+        self.case_id = case_data["case_id"]
+        self.categories = categories
+        self.sections = sections
+        self.system_prompt = system_prompt
+        self.case_content = self.select_case_sections()
+        self.raw_content: Optional[str] = None
+        self.content: Optional[Dict[str, str]] = None
+
+    def create_question_prompt(self) -> str:
+        """Creates a formatted prompt for generating a clinical question.
+
+        Returns:
+            str: A structured prompt containing the question parameters and clinical data
+        """
+        category_descriptions = "\n".join(
+            f"{category}: {desc}"
+            for category, desc in CATEGORIES_META.items()
+            if category in self.categories
+        )
+
+        return f"""
+        You must follow these guidelines:
+        1. Questions must be answerable using only context and chest X-rays.
+        - Questions must explicitly mention the referenced figures
+        - Questions can only reference the chest X-ray figures
+
+        2. Questions must have unambiguous, verifiable answers, and should:
+        - Challenge the agent's analytical capabilities
+        - Require multi-step reasoning
+        - Test ability to make precise observations
+        - Evaluate capability to derive insights and findings from the chest X-ray
+
+        3. The agent has access to tools like classification, report generation, segmentation, grounding, visual question answering, etc. Your question should be complex to require the use of such tools.
+
+
+        Create a {self.difficulty} {self.type} clinical question that integrates the following:
+
+        {category_descriptions}
+
+        based on the following clinical case:
+
+        {self.case_content}
+
+        Do not use any infomration derived from the CT and MRI images. Do not provide any information and findings about the chest X-rays.
+        Your question should require the agent to derive insights and findings from the chest X-ray by itself.
+        Your answer should be verifiable directly in the context of the case.
+        You can only use the image findings that come from the chest X-ray figures.
+
+        Your response must follow this exact format:
+        THOUGHTS: [Think about different reasoning steps and tools the agent should use to answer the question]
+        QUESTION: [complete question with relevant context. Incorrect choices should be very close to the correct answer.]
+        FIGURES: [list of required figures, e.g. ["Figure 1", "Figure 2a"]]
+        EXPLANATION: [short explanation of why your answer is verifiable in the case]
+        ANSWER: [correct answer e.g. "A"]
+        """.strip().replace(
+            "        ", ""
+        )  # remove tabs
+
+    def select_case_sections(self) -> str:
+        """Extract and format selected sections from case data into paragraphs.
+
+        Returns:
+            str: Formatted string with case sections and content
+        """
+        section_mapping = {
+            "history": ("history", "No history provided."),
+            "image_finding": ("image_finding", "No findings provided."),
+            "discussion": ("discussion", "No discussion provided."),
+            "differential_diagnosis": (
+                "differential_diagnosis",
+                "No differential diagnosis provided.",
+            ),
+            "diagnosis": ("diagnosis", "No diagnosis provided."),
+            "figures": ("figures", "No figures provided."),
+        }
+
+        formatted = []
+        for section in self.sections:
+            if section in section_mapping:
+                key, default = section_mapping[section]
+                content = self.case_data.get(key, default)
+
+                if key == "figures":
+                    figures_text = []
+                    for figure in content:
+                        for subfig in figure["subfigures"]:
+                            figures_text.append(f"{subfig['number']}: {subfig['caption']}")
+                    content = "\n".join(figures_text)
+
+                formatted.append(f"{section}:\n{content}")
+
+        return "\n\n".join(formatted)
+
+    def create_question(
+        self,
+        client: openai.OpenAI,
+        temperature: float = 0.7,
+        top_p: float = 0.95,
+        max_tokens: int = 500,
+        model: str = "gpt-4o",
+    ) -> str:
+        """Create a clinical question using LLM.
+
+        Args:
+            client (openai.OpenAI): OpenAI client instance
+            temperature (float): Controls randomness in responses. Defaults to 0.7.
+            top_p (float): Controls diversity via nucleus sampling. Defaults to 0.95.
+            max_tokens (int): Max tokens in model response. Defaults to 500.
+            model (str): OpenAI model to use. Defaults to "gpt-4o".
+
+        Returns:
+            str: LLM response containing formatted question components
+        """
+        self.raw_content = get_llm_response(
+            client=client,
+            prompt=self.create_question_prompt(),
+            system_prompt=self.system_prompt,
+            temperature=temperature,
+            top_p=top_p,
+            max_tokens=max_tokens,
+            model=model,
+        )
+        self.content = self.extract_content()
+
+        return self.raw_content
+
+    def extract_content(self) -> Dict[str, str]:
+        """Extract sections from raw LLM response using regex patterns.
+
+        Returns:
+            Dict[str, str]: Extracted sections including thoughts, question, figures, explanation, and answer
+        """
+        keywords = ["THOUGHTS", "QUESTION", "FIGURES", "EXPLANATION", "ANSWER"]
+
+        content = {}
+        for kw in keywords:
+            pattern = rf"{kw}:\s*(.*?)(?=\n[A-Z]+:|$)"
+            match = re.search(pattern, self.raw_content, re.DOTALL)
+            content[kw.lower()] = match.group(1).strip() if match else None
+
+        return content
+
+    def save(self, output_path: str) -> Dict[str, Any]:
+        """Save question content and metadata as a JSON file.
+
+        Args:
+            output_path (str): Directory path where the JSON file will be saved
+
+        Returns:
+            Dict[str, Any]: Question data including content (thoughts, question, figures, options,
+                explanation, answer) and metadata (type, difficulty, categories, etc.)
+        """
+        question_metadata = self.content.copy()
+
+        # Add metadata
+        question_metadata["metadata"] = {
+            "case_id": self.case_id,
+            "type": self.type,
+            "difficulty": self.difficulty,
+            "categories": self.categories,
+            "sections": self.sections,
+        }
+
+        # Create a directory for the case
+        case_dir = os.path.join(output_path, str(self.case_id))
+        os.makedirs(case_dir, exist_ok=True)
+
+        # Save the question metadata to a JSON file
+        output_file = os.path.join(case_dir, f"{self.case_id}_{self.__hash__()}.json")
+        with open(output_file, "w") as f:
+            json.dump(question_metadata, f, indent=2)
+
+        return question_metadata
+
+
+def generate_questions(
+    dataset: Dict[str, Any],
+    client: openai.OpenAI,
+    output_dir: str,
+    skip_first: int = 100,
+    temperature: float = 0.7,
+    top_p: float = 0.95,
+    max_tokens: int = 1200,
+    model: str = "gpt-4o",
+) -> None:
+    """Generate questions for each case and category combination.
+
+    Args:
+        dataset: Dictionary of case data
+        client: OpenAI client instance
+        output_dir: Directory to save generated questions
+        skip_first: Number of initial cases to skip
+        temperature: LLM temperature parameter
+        top_p: LLM top_p parameter
+        max_tokens: Maximum tokens for LLM response
+        model: LLM model name
+    """
+    target_cases = sorted(list(dataset.keys()), key=int)[-len(dataset) : -skip_first]
+
+    for case_id in tqdm(target_cases, desc="Processing cases"):
+        case_data = dataset[case_id]
+
+        for category in tqdm(CATEGORY_COMBINATIONS, desc=f"Categories for case {case_id}"):
+            question = Question(
+                type="multiple choice (A/B/C/D/E/F)",
+                difficulty="complex",
+                case_data=case_data,
+                categories=category,
+                sections=DEFAULT_SECTIONS,
+                system_prompt=SYSTEM_PROMPT,
+            )
+
+            response = question.create_question(
+                client=client,
+                temperature=temperature,
+                top_p=top_p,
+                max_tokens=max_tokens,
+                model=model,
+            )
+            question.save(output_dir)
+
+
+def main():
+    """Main execution function."""
+    client = openai.OpenAI()
+
+    # Load and verify dataset
+    dataset = load_eurorad_dataset(
+        DATASET_PATH,
+        section="Chest Imaging",
+        as_dict=True,
+        filter_by_caption=[
+            "xray",
+            "x-ray",
+            "x ray",
+            "ray",
+            "xr",
+            "radiograph",
+        ],
+    )
+    print(f"\n---\nFound {len(dataset)} cases with X-ray mentions\n---\n")
+
+    # Optional: Print sample case for verification
+    case_data = dataset["16798"]
+    pprint(case_data, sort_dicts=False)
+
+    # Generate questions
+    generate_questions(dataset=dataset, client=client, output_dir="benchmark/questions")
+
+
+if __name__ == "__main__":
+    main()

benchmark/llm.py

@@ -0,0 +1,42 @@
+import openai
+from typing import List
+
+
+def get_llm_response(
+    client: openai.OpenAI,
+    prompt: str,
+    system_prompt: str = "You are a helpful assistant.",
+    model: str = "gpt-4o-mini",
+    temperature: float = 0.7,
+    top_p: float = 0.95,
+    max_tokens: int = 500,
+) -> str:
+    """
+    Get response from OpenAI language model.
+
+    Args:
+        client (openai.OpenAI): OpenAI client
+        prompt (str): The user prompt/question to send to the model
+        system_prompt (str, optional): System prompt to set model behavior.
+        model (str, optional): OpenAI model to use. Defaults to "gpt-4o-mini".
+        temperature (float, optional): Controls randomness in responses. Defaults to 0.7.
+        top_p (float, optional): Controls diversity via nucleus sampling. Defaults to 0.95.
+        max_tokens (int, optional): Max tokens in model response. Defaults to 200.
+
+    Returns:
+        str: The model's response text
+    """
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": prompt},
+    ]
+
+    response = client.chat.completions.create(
+        model=model,
+        messages=messages,
+        temperature=temperature,
+        top_p=top_p,
+        max_tokens=max_tokens,
+    )
+
+    return response.choices[0].message.content

benchmark/utils.py

@@ -0,0 +1,78 @@
+import os
+import json
+from typing import Dict, List
+
+
+def load_eurorad_dataset(
+    dataset_path: str,
+    section: str = "any",
+    as_dict: bool = False,
+    filter_by_caption: List[str] = [
+        "xray",
+        "x-ray",
+        "x ray",
+        "ray",
+        "xr",
+        "radiograph",
+        "radiogram",
+        "plain film",
+    ],
+) -> List[Dict] | Dict[str, Dict]:
+    """
+    Load a dataset from a JSON file.
+
+    Args:
+        dataset_path (str): Path to the JSON dataset file.
+        section (str, optional): Section of the dataset to load. Defaults to "any".
+        as_dict (bool, optional): Whether to return data as dict. Defaults to False.
+        filter_by_caption (List[str], optional): List of strings to filter cases by caption content. Defaults to [].
+
+    Returns:
+        List[Dict] | Dict[str, Dict]: The loaded dataset as a list of dictionaries or dict if as_dict=True.
+
+    Raises:
+        FileNotFoundError: If dataset_path does not exist
+        json.JSONDecodeError: If file is not valid JSON
+    """
+
+    with open(dataset_path, "r", encoding="utf-8") as file:
+        data = json.load(file)
+
+    if filter_by_caption:
+        filtered_data = {}
+        for case_id, case in data.items():
+            if any(
+                any(x in subfig["caption"].lower() for x in filter_by_caption)
+                for figure in case["figures"]
+                for subfig in figure["subfigures"]
+            ) or any(x in case["image_finding"].lower() for x in filter_by_caption):
+                filtered_data[case_id] = case
+        data = filtered_data
+
+    if section != "any":
+        section = section.strip().lower()
+        if not as_dict:
+            data = [
+                item for item in data.values() if item.get("section", "").strip().lower() == section
+            ]
+        else:
+            data = {
+                k: v for k, v in data.items() if v.get("section", "").strip().lower() == section
+            }
+
+    elif not as_dict:
+        data = list(data.values())
+
+    return data
+
+
+def save_dataset(dataset: Dict | List[Dict], dataset_path: str):
+    """
+    Save a dataset to a JSON file.
+
+    Args:
+        dataset (Dict | List[Dict]): The dataset to save as a dictionary or list of dictionaries.
+        dataset_path (str): Path where the JSON dataset file will be saved.
+    """
+    with open(dataset_path, "w", encoding="utf-8") as file:
+        json.dump(dataset, file)

data/figures.py

@@ -0,0 +1,74 @@
+import json
+import os
+from pathlib import Path
+import requests
+from tqdm import tqdm
+
+
+def download_eurorad_figures(metadata_path: str, output_dir: str) -> None:
+    """
+    Download figures from Eurorad dataset and save them organized by case_id.
+
+    Args:
+        metadata_path: Path to the eurorad_metadata.json file
+        output_dir: Base directory where figures will be saved
+
+    The figures will be saved as:
+        {output_dir}/{case_id}/{figure_number}.jpg
+    Example:
+        figures/189/Figure_1a.jpg
+    """
+    # Create output directory if it doesn't exist
+    output_path = Path(output_dir)
+    output_path.mkdir(exist_ok=True)
+
+    # Load metadata
+    with open(metadata_path) as f:
+        metadata = json.load(f)
+
+    # Iterate through all cases with progress bar
+    for case_id in tqdm(metadata, desc="Downloading cases", unit="case"):
+        case = metadata[case_id]
+        case_dir = output_path / str(case["case_id"])
+        case_dir.mkdir(exist_ok=True)
+
+        # Process all figures and their subfigures
+        for figure in case["figures"]:
+            for subfig in figure["subfigures"]:
+
+                # Remove leading and trailing whitespace and convert to lowercase
+                subfig_name = f"{subfig['number'].strip().replace(' ', '_').lower()}.jpg"
+                subfig_path = Path(case_dir) / subfig_name
+
+                save_figure(
+                    url=subfig["url"],
+                    output_path=subfig_path,
+                )
+
+
+def save_figure(url: str, output_path: Path) -> None:
+    """
+    Download and save a single figure.
+
+    Args:
+        url: URL of the figure to download
+        output_path: Path where the figure should be saved
+    """
+    if output_path.exists():
+        return
+
+    try:
+        response = requests.get(url, timeout=10)
+        response.raise_for_status()
+        with open(output_path, "wb") as f:
+            f.write(response.content)
+    except Exception as e:
+        print(f"Error downloading {url}: {e}")
+
+
+if __name__ == "__main__":
+    root = os.path.dirname(os.path.abspath(__file__))
+    download_eurorad_figures(
+        metadata_path=os.path.join(root, "eurorad_metadata.json"),
+        output_dir=os.path.join(root, "figures"),
+    )

data/get_cases.py

@@ -0,0 +1,51 @@
+import requests
+from bs4 import BeautifulSoup
+import time
+import json
+from tqdm import tqdm
+
+
+def get_response(url):
+    headers = {
+        "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/108.0.0.0 Safari/537.36 Edg/108.0.1462.54"
+    }
+    return requests.get(url, headers=headers)
+
+def get_case_numbers_from_page(page):
+    url = f"https://www.eurorad.org/advanced-search?sort_by=published_at&sort_order=ASC&page={page}&filter%5B0%5D=section%3A40"
+
+    # Remove proxy usage since it's likely triggering the protection
+    response = get_response(url)
+    print(response.text)
+
+    soup = BeautifulSoup(response.text, "html.parser")
+    spans = soup.find_all("span", class_="case__number small")
+
+    # Remove '#' from the span text and strip extra whitespace
+    numbers = [span.text.strip().replace("#", "").strip() for span in spans]
+    return numbers
+
+
+def main():
+    total_pages = 107  # Pages 0 through 106
+    all_numbers = []
+
+    for page in tqdm(range(total_pages)):
+        numbers = get_case_numbers_from_page(page)
+        all_numbers.extend(numbers)
+
+        if page != total_pages - 1 and len(numbers) != 9:
+            print(f"Warning: Page {page} returned {len(numbers)} cases instead of 9")
+
+        # Be kind to the server – avoid hitting it too fast
+        time.sleep(1)
+        break
+
+    with open('case_numbers.json', 'w') as f:
+        json.dump(all_numbers, f)
+
+    print(f"Saved {len(all_numbers)} case numbers to case_numbers.json")
+
+
+if __name__ == "__main__":
+    main()

experiments/README.md

@@ -0,0 +1,63 @@
+# Experiments
+Below are the instructions for running experiments using our novel ChestAgentBench and the previous SoTA CheXbench. ChestAgentBench is a comprehensive benchmark containing over 2,500 complex medical queries across 8 diverse categories.
+
+### ChestAgentBench
+
+To run gpt-4o on ChestAgentBench, enter the `experiments` directory and run the following script:
+```bash
+python benchmark_gpt4o.py
+```
+
+To run llama 3.2 vision 90B on ChestAgentBench, run the following:
+```bash
+python benchmark_llama.py
+```
+
+To run chexagent on ChestAgentBench, run the following:
+```bash
+python benchmark_chexagent.py
+```
+
+To run llava-med on ChestAgentBench, you'll need to clone their repo and copy the following script into it, after you follow their setup instructions.
+```bash
+mv benchmark_llavamed.py ~/LLaVA-Med/llava/serve
+python -m llava.serve.benchmark_llavamed --model-name llava-med-v1.5-mistral-7b --controller http://localhost:10000
+```
+
+If you want to inspect the logs, you can run the following. It will select the most recent log file by default.
+```bash
+python inspect_logs.py [optional: log-file] -n [num-logs]
+```
+
+Finally, to analyze results, run:
+```bash
+python analyze_axes.py results/[logfile].json ../benchmark/questions/ --model [gpt4|llama|chexagent|llava-med] --max-questions [optional:int]
+```
+
+### CheXbench
+
+To run the models on chexbench, you can use `chexbench_gpt4.py` as a reference. You'll need to download the dataset files locally, and upload them for each request. Rad-ReStruct and Open-I use the same set of images, so you can download the `NLMCXR.zip` file just once and copy the images to both directories.
+
+You can find the datasets here:
+1. [SLAKE: A Semantically-Labeled Knowledge-Enhanced Dataset for Medical Visual Question Answering](https://www.med-vqa.com/slake/). Save this to `MedMAX/data/slake`.
+2. [Rad-ReStruct: A Novel VQA Benchmark and Method for Structured Radiology Reporting](https://github.com/ChantalMP/Rad-ReStruct). Save the images to `MedMAX/data/rad-restruct/images`.
+3. [Open-I Service of the National Library of Medicine](https://openi.nlm.nih.gov/faq). Save the images to `MedMAX/data/openi/images`.
+
+Once you're finished, you'll want to fix the paths in the `chexbench.json` file to your local paths using the `MedMax/data/fix_chexbench.py` script.
+
+
+### Compare Runs
+Analyze a single file based on overall accuracy and along different axes
+```
+python compare_runs.py results/medmax.json
+```
+
+For a direct evaluation comparing **2** models, on the exact same questions 
+```
+python compare_runs.py results/medmax.json results/gpt4o.json
+```
+
+For a direct evaluation comparing **ALL** models, on the exact same questions (add as many model log files as you want).
+```
+python compare_runs.py results/medmax.json results/gpt4o.json results/llama.json results/chexagent.json results/llavamed.json
+```

experiments/analyze_axes.py

@@ -0,0 +1,385 @@
+from typing import Dict, List, Optional, Tuple, Union, Any
+import json
+import os
+import sys
+import argparse
+from collections import defaultdict
+from tqdm import tqdm
+
+QUESTION_TYPES = {
+    "Detailed Finding Analysis": ["detection", "localization", "characterization"],
+    "Pattern Recognition & Relations": ["detection", "classification", "relationship"],
+    "Spatial Understanding": ["localization", "comparison", "relationship"],
+    "Clinical Decision Making": ["classification", "comparison", "diagnosis"],
+    "Diagnostic Classification": ["classification", "characterization", "diagnosis"],
+}
+
+
+def extract_answer_letter(answer: Optional[Union[str, Any]]) -> Optional[str]:
+    """
+    Extract just the letter from various answer formats.
+
+    Args:
+        answer: The answer text to extract letter from
+
+    Returns:
+        Optional[str]: The extracted letter in uppercase, or None if no letter found
+    """
+    if not answer:
+        return None
+
+    # Convert to string and clean
+    answer = str(answer).strip()
+
+    # If it's just a single letter, return it
+    if len(answer) == 1 and answer.isalpha():
+        return answer.upper()
+
+    # Try to extract letter from format like "A)" or "A."
+    if len(answer) >= 2 and answer[0].isalpha() and answer[1] in ").:- ":
+        return answer[0].upper()
+
+    # Try to extract letter from format like "A) Some text"
+    if answer.startswith(("A)", "B)", "C)", "D)", "E)", "F)")):
+        return answer[0].upper()
+
+    return None
+
+
+def analyze_gpt4_results(
+    results_file: str, max_questions: Optional[int] = None
+) -> Tuple[float, Dict, Dict, List[str], List[str]]:
+    """
+    Analyze results in GPT-4 format.
+
+    Args:
+        results_file: Path to results file
+        max_questions: Maximum number of questions to analyze
+
+    Returns:
+        Tuple containing:
+            - overall_accuracy (float)
+            - category_accuracies (Dict)
+            - question_type_stats (Dict)
+            - correct_ids (List[str])
+            - incorrect_ids (List[str])
+    """
+    category_performance = defaultdict(lambda: {"total": 0, "correct": 0})
+    all_questions = 0
+    all_correct = 0
+    correct_ids = []
+    incorrect_ids = []
+
+    with open(results_file, "r") as f:
+        lines = f.readlines()
+
+    processed_questions = 0
+
+    for line in tqdm(lines, desc="Analyzing Benchmark Results"):
+        # Check if we've hit the maximum questions
+        if max_questions is not None and processed_questions >= max_questions:
+            break
+        if line.startswith("HTTP Request:"):
+            continue
+
+        try:
+            entry = json.loads(line)
+            metadata = entry.get("input", {}).get("question_data", {}).get("metadata", {})
+            question_id = entry.get("question_id")
+
+            model_letter = extract_answer_letter(entry.get("model_answer"))
+            correct_letter = extract_answer_letter(entry.get("correct_answer"))
+
+            if model_letter and correct_letter:
+                all_questions += 1
+                processed_questions += 1
+                is_correct = model_letter == correct_letter
+
+                if is_correct:
+                    all_correct += 1
+                    correct_ids.append(question_id)
+                else:
+                    incorrect_ids.append(question_id)
+
+                for category in metadata.get("categories", []):
+                    category_performance[category]["total"] += 1
+                    if is_correct:
+                        category_performance[category]["correct"] += 1
+
+        except json.JSONDecodeError:
+            continue
+
+    return process_results(
+        category_performance, all_questions, all_correct, correct_ids, incorrect_ids
+    )
+
+
+def analyze_llama_results(
+    results_file: str, max_questions: Optional[int] = None
+) -> Tuple[float, Dict, Dict, List[str], List[str]]:
+    """
+    Analyze results in Llama format.
+
+    Args:
+        results_file: Path to results file
+        max_questions: Maximum number of questions to analyze
+
+    Returns:
+        Tuple containing:
+            - overall_accuracy (float)
+            - category_accuracies (Dict)
+            - question_type_stats (Dict)
+            - correct_ids (List[str])
+            - incorrect_ids (List[str])
+    """
+    category_performance = defaultdict(lambda: {"total": 0, "correct": 0})
+    all_questions = 0
+    all_correct = 0
+    correct_ids = []
+    incorrect_ids = []
+
+    with open(results_file, "r") as f:
+        lines = f.readlines()
+
+    # If max_questions is set, limit the number of lines processed
+    if max_questions is not None:
+        lines = lines[:max_questions]
+
+    for line in tqdm(lines, desc="Analyzing Benchmark Results"):
+        if line.startswith("HTTP Request:"):
+            continue
+
+        try:
+            entry = json.loads(line)
+            metadata = entry.get("input", {}).get("question_data", {}).get("metadata", {})
+            question_id = entry.get("question_id")
+
+            model_letter = extract_answer_letter(entry.get("model_answer"))
+            correct_letter = extract_answer_letter(entry.get("correct_answer"))
+
+            if model_letter and correct_letter:
+                all_questions += 1
+                is_correct = model_letter == correct_letter
+
+                if is_correct:
+                    all_correct += 1
+                    correct_ids.append(question_id)
+                else:
+                    incorrect_ids.append(question_id)
+
+                for category in metadata.get("categories", []):
+                    category_performance[category]["total"] += 1
+                    if is_correct:
+                        category_performance[category]["correct"] += 1
+
+        except json.JSONDecodeError:
+            continue
+
+    return process_results(
+        category_performance, all_questions, all_correct, correct_ids, incorrect_ids
+    )
+
+
+def analyze_chexagent_results(
+    results_file: str, max_questions: Optional[int] = None
+) -> Tuple[float, Dict, Dict, List[str], List[str]]:
+    """
+    Analyze results in CheXagent format.
+
+    Args:
+        results_file: Path to results file
+        max_questions: Maximum number of questions to analyze
+
+    Returns:
+        Tuple containing:
+            - overall_accuracy (float)
+            - category_accuracies (Dict)
+            - question_type_stats (Dict)
+            - correct_ids (List[str])
+            - incorrect_ids (List[str])
+    """
+    category_performance = defaultdict(lambda: {"total": 0, "correct": 0})
+    all_questions = 0
+    all_correct = 0
+    correct_ids = []
+    incorrect_ids = []
+
+    with open(results_file, "r") as f:
+        lines = f.readlines()
+
+    # If max_questions is set, limit the number of lines processed
+    if max_questions is not None:
+        lines = lines[:max_questions]
+
+    for line in tqdm(lines, desc="Analyzing Benchmark Results"):
+        try:
+            entry = json.loads(line)
+            metadata = entry.get("input", {}).get("question_data", {}).get("metadata", {})
+            question_id = entry.get("question_id")
+
+            model_letter = extract_answer_letter(entry.get("model_answer"))
+            correct_letter = extract_answer_letter(entry.get("correct_answer"))
+
+            if model_letter and correct_letter:
+                all_questions += 1
+                is_correct = model_letter == correct_letter
+
+                if is_correct:
+                    all_correct += 1
+                    correct_ids.append(question_id)
+                else:
+                    incorrect_ids.append(question_id)
+
+                for category in metadata.get("categories", []):
+                    category_performance[category]["total"] += 1
+                    if is_correct:
+                        category_performance[category]["correct"] += 1
+
+        except json.JSONDecodeError:
+            continue
+
+    return process_results(
+        category_performance, all_questions, all_correct, correct_ids, incorrect_ids
+    )
+
+
+def process_results(
+    category_performance: Dict,
+    all_questions: int,
+    all_correct: int,
+    correct_ids: Optional[List[str]] = None,
+    incorrect_ids: Optional[List[str]] = None,
+) -> Tuple[float, Dict, Dict, List[str], List[str]]:
+    """
+    Process raw results into final statistics.
+
+    Args:
+        category_performance: Dict containing performance by category
+        all_questions: Total number of questions
+        all_correct: Total number of correct answers
+        correct_ids: List of IDs for correctly answered questions
+        incorrect_ids: List of IDs for incorrectly answered questions
+
+    Returns:
+        Tuple containing:
+            - overall_accuracy (float)
+            - category_accuracies (Dict)
+            - question_type_stats (Dict)
+            - correct_ids (List[str])
+            - incorrect_ids (List[str])
+    """
+    category_accuracies = {
+        category: {
+            "accuracy": stats["correct"] / stats["total"] * 100 if stats["total"] > 0 else 0,
+            "total": stats["total"],
+            "correct": stats["correct"],
+        }
+        for category, stats in category_performance.items()
+    }
+
+    question_type_stats = {}
+    for qtype, categories in QUESTION_TYPES.items():
+        total = sum(
+            category_performance[cat]["total"] for cat in categories if cat in category_performance
+        )
+        correct = sum(
+            category_performance[cat]["correct"]
+            for cat in categories
+            if cat in category_performance
+        )
+
+        question_type_stats[qtype] = {
+            "accuracy": (correct / total * 100) if total > 0 else 0,
+            "total": total,
+            "correct": correct,
+        }
+
+    overall_accuracy = (all_correct / all_questions * 100) if all_questions > 0 else 0
+
+    return (
+        overall_accuracy,
+        category_accuracies,
+        question_type_stats,
+        correct_ids or [],
+        incorrect_ids or [],
+    )
+
+
+def print_analysis(
+    overall_accuracy: float,
+    category_accuracies: Dict,
+    question_type_stats: Dict,
+    correct_ids: List[str],
+    incorrect_ids: List[str],
+    model_name: str,
+) -> None:
+    """
+    Print analysis results.
+
+    Args:
+        overall_accuracy: Overall accuracy percentage
+        category_accuracies: Dict containing accuracy metrics by category
+        question_type_stats: Dict containing stats by question type
+        correct_ids: List of IDs for correctly answered questions
+        incorrect_ids: List of IDs for incorrectly answered questions
+        model_name: Name of the model being analyzed
+    """
+    total_questions = len(correct_ids) + len(incorrect_ids)
+    print(
+        f"\nOverall Accuracy: {overall_accuracy:.2f}% ({len(correct_ids)} correct out of {total_questions} questions)"
+    )
+
+    print("\nCategory Performance:")
+    sorted_categories = sorted(
+        category_accuracies.items(), key=lambda x: x[1]["accuracy"], reverse=True
+    )
+    for category, metrics in sorted_categories:
+        print(f"{category}:")
+        print(f"  Accuracy: {metrics['accuracy']:.2f}%")
+        print(f"  Total Questions: {metrics['total']}")
+        print(f"  Correct Questions: {metrics['correct']}")
+
+    print("\nQuestion Type Performance:")
+    sorted_types = sorted(question_type_stats.items(), key=lambda x: x[1]["accuracy"], reverse=True)
+    for qtype, metrics in sorted_types:
+        print(f"\n{qtype}:")
+        print(f"  Accuracy: {metrics['accuracy']:.2f}%")
+        print(f"  Total Questions: {metrics['total']}")
+        print(f"  Correct Questions: {metrics['correct']}")
+        print(f"  Categories: {', '.join(QUESTION_TYPES[qtype])}")
+
+    # Save question IDs to JSON
+    question_ids = {"correct_ids": correct_ids, "incorrect_ids": incorrect_ids}
+
+    output_filename = f"{model_name}_question_ids.json"
+    with open(output_filename, "w") as f:
+        json.dump(question_ids, f, indent=2)
+
+    print(f"\nQuestion IDs have been saved to {output_filename}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Analyze benchmark results")
+    parser.add_argument("results_file", help="Path to results file")
+    parser.add_argument("benchmark_dir", nargs="?", help="Path to benchmark questions directory")
+    parser.add_argument(
+        "--model",
+        choices=["llava-med", "chexagent", "llama", "gpt4", "medrax"],
+        default="gpt4",
+        help="Specify model format (default: gpt4)",
+    )
+    parser.add_argument("--max-questions", type=int, help="Maximum number of questions to analyze")
+    args = parser.parse_args()
+
+    if args.model == "gpt4":
+        results = analyze_gpt4_results(args.results_file, args.max_questions)
+    elif args.model == "llama":
+        results = analyze_llama_results(args.results_file, args.max_questions)
+    elif args.model == "chexagent":
+        results = analyze_chexagent_results(args.results_file, args.max_questions)
+    elif args.model == "medrax":
+        results = analyze_gpt4_results(args.results_file, args.max_questions)
+    else:
+        parser.error(f"Unsupported model: {args.model}")
+
+    print_analysis(*results, args.model)

experiments/benchmark_chexagent.py

@@ -0,0 +1,316 @@
+import re
+import json
+import os
+import glob
+import time
+import logging
+from datetime import datetime
+import torch
+from PIL import Image
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from tqdm import tqdm
+
+# Configure model settings
+MODEL_NAME = "StanfordAIMI/CheXagent-2-3b"
+DTYPE = torch.bfloat16
+DEVICE = "cuda"
+
+# Configure logging
+log_filename = f"model_inference_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
+logging.basicConfig(filename=log_filename, level=logging.INFO, format="%(message)s")
+
+
+def initialize_model() -> tuple[AutoModelForCausalLM, AutoTokenizer]:
+    """Initialize the CheXagent model and tokenizer.
+
+    Returns:
+        tuple containing:
+            - AutoModelForCausalLM: The initialized CheXagent model
+            - AutoTokenizer: The initialized tokenizer
+    """
+    print("Loading model and tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME, device_map="auto", trust_remote_code=True
+    )
+    model = model.to(DTYPE)
+    model.eval()
+    return model, tokenizer
+
+
+def create_inference_request(
+    question_data: dict,
+    case_details: dict,
+    case_id: str,
+    question_id: str,
+    model: AutoModelForCausalLM,
+    tokenizer: AutoTokenizer,
+) -> str | None:
+    """Create and execute an inference request for the CheXagent model.
+
+    Args:
+        question_data: Dictionary containing question details and metadata
+        case_details: Dictionary containing case information and image paths
+        case_id: Unique identifier for the medical case
+        question_id: Unique identifier for the question
+        model: The initialized CheXagent model
+        tokenizer: The initialized tokenizer
+
+    Returns:
+        str | None: Single letter answer (A-F) if successful, None if failed
+    """
+    system_prompt = """You are a medical imaging expert. Your task is to provide ONLY a single letter answer.
+Rules:
+1. Respond with exactly one uppercase letter (A/B/C/D/E/F)
+2. Do not add periods, explanations, or any other text
+3. Do not use markdown or formatting
+4. Do not restate the question
+5. Do not explain your reasoning
+
+Examples of valid responses:
+A
+B
+C
+
+Examples of invalid responses:
+"A."
+"Answer: B"
+"C) This shows..."
+"The answer is D"
+"""
+
+    prompt = f"""Given the following medical case:
+Please answer this multiple choice question:
+{question_data['question']}
+Base your answer only on the provided images and case information."""
+
+    # Parse required figures
+    try:
+        if isinstance(question_data["figures"], str):
+            try:
+                required_figures = json.loads(question_data["figures"])
+            except json.JSONDecodeError:
+                required_figures = [question_data["figures"]]
+        elif isinstance(question_data["figures"], list):
+            required_figures = question_data["figures"]
+        else:
+            required_figures = [str(question_data["figures"])]
+    except Exception as e:
+        print(f"Error parsing figures: {e}")
+        required_figures = []
+
+    required_figures = [
+        fig if fig.startswith("Figure ") else f"Figure {fig}" for fig in required_figures
+    ]
+
+    # Get image paths
+    image_paths = []
+    for figure in required_figures:
+        base_figure_num = "".join(filter(str.isdigit, figure))
+        figure_letter = "".join(filter(str.isalpha, figure.split()[-1])) or None
+
+        matching_figures = [
+            case_figure
+            for case_figure in case_details.get("figures", [])
+            if case_figure["number"] == f"Figure {base_figure_num}"
+        ]
+
+        for case_figure in matching_figures:
+            subfigures = []
+            if figure_letter:
+                subfigures = [
+                    subfig
+                    for subfig in case_figure.get("subfigures", [])
+                    if subfig.get("number", "").lower().endswith(figure_letter.lower())
+                    or subfig.get("label", "").lower() == figure_letter.lower()
+                ]
+            else:
+                subfigures = case_figure.get("subfigures", [])
+
+            for subfig in subfigures:
+                if "local_path" in subfig:
+                    image_paths.append("medrax/data/" + subfig["local_path"])
+
+    if not image_paths:
+        print(f"No local images found for case {case_id}, question {question_id}")
+        return None
+
+    try:
+        start_time = time.time()
+
+        # Prepare input for the model
+        query = tokenizer.from_list_format(
+            [*[{"image": path} for path in image_paths], {"text": prompt}]
+        )
+        conv = [{"from": "system", "value": system_prompt}, {"from": "human", "value": query}]
+        input_ids = tokenizer.apply_chat_template(
+            conv, add_generation_prompt=True, return_tensors="pt"
+        )
+
+        # Generate response
+        with torch.no_grad():
+            output = model.generate(
+                input_ids.to(DEVICE),
+                do_sample=False,
+                num_beams=1,
+                temperature=1.0,
+                top_p=1.0,
+                use_cache=True,
+                max_new_tokens=512,
+            )[0]
+
+        response = tokenizer.decode(output[input_ids.size(1) : -1])
+        duration = time.time() - start_time
+
+        # Clean response
+        clean_answer = validate_answer(response)
+
+        # Log response
+        log_entry = {
+            "case_id": case_id,
+            "question_id": question_id,
+            "timestamp": datetime.now().isoformat(),
+            "model": MODEL_NAME,
+            "duration": round(duration, 2),
+            "model_answer": clean_answer,
+            "correct_answer": question_data["answer"],
+            "input": {
+                "question_data": {
+                    "question": question_data["question"],
+                    "explanation": question_data["explanation"],
+                    "metadata": question_data.get("metadata", {}),
+                    "figures": question_data["figures"],
+                },
+                "image_paths": image_paths,
+            },
+        }
+        logging.info(json.dumps(log_entry))
+        return clean_answer
+
+    except Exception as e:
+        print(f"Error processing case {case_id}, question {question_id}: {str(e)}")
+        log_entry = {
+            "case_id": case_id,
+            "question_id": question_id,
+            "timestamp": datetime.now().isoformat(),
+            "model": MODEL_NAME,
+            "status": "error",
+            "error": str(e),
+            "input": {
+                "question_data": {
+                    "question": question_data["question"],
+                    "explanation": question_data["explanation"],
+                    "metadata": question_data.get("metadata", {}),
+                    "figures": question_data["figures"],
+                },
+                "image_paths": image_paths,
+            },
+        }
+        logging.info(json.dumps(log_entry))
+        return None
+
+
+def validate_answer(response_text: str) -> str | None:
+    """Enforce strict single-letter response format.
+
+    Args:
+        response_text: Raw response text from the model
+
+    Returns:
+        str | None: Single uppercase letter (A-F) if valid, None if invalid
+    """
+    if not response_text:
+        return None
+
+    # Remove all whitespace and convert to uppercase
+    cleaned = response_text.strip().upper()
+
+    # Check if it's exactly one valid letter
+    if len(cleaned) == 1 and cleaned in "ABCDEF":
+        return cleaned
+
+    # If not, try to extract just the letter
+    match = re.search(r"([A-F])", cleaned)
+    return match.group(1) if match else None
+
+
+def load_benchmark_questions(case_id: str) -> list[str]:
+    """Find all question files for a given case ID.
+
+    Args:
+        case_id: Unique identifier for the medical case
+
+    Returns:
+        list[str]: List of paths to question JSON files
+    """
+    benchmark_dir = "../benchmark/questions"
+    return glob.glob(f"{benchmark_dir}/{case_id}/{case_id}_*.json")
+
+
+def count_total_questions() -> tuple[int, int]:
+    """Count total number of cases and questions in benchmark.
+
+    Returns:
+        tuple containing:
+            - int: Total number of cases
+            - int: Total number of questions
+    """
+    total_cases = len(glob.glob("../benchmark/questions/*"))
+    total_questions = sum(
+        len(glob.glob(f"../benchmark/questions/{case_id}/*.json"))
+        for case_id in os.listdir("../benchmark/questions")
+    )
+    return total_cases, total_questions
+
+
+def main():
+    # Load the cases with local paths
+    with open("medrax/data/updated_cases.json", "r") as file:
+        data = json.load(file)
+
+    # Initialize model and tokenizer
+    model, tokenizer = initialize_model()
+
+    total_cases, total_questions = count_total_questions()
+    cases_processed = 0
+    questions_processed = 0
+    skipped_questions = 0
+
+    print(f"\nBeginning inference with {MODEL_NAME}")
+    print(f"Found {total_cases} cases with {total_questions} total questions")
+
+    # Process each case with progress bar
+    for case_id, case_details in tqdm(data.items(), desc="Processing cases"):
+        question_files = load_benchmark_questions(case_id)
+        if not question_files:
+            continue
+
+        cases_processed += 1
+        for question_file in tqdm(
+            question_files, desc=f"Processing questions for case {case_id}", leave=False
+        ):
+            with open(question_file, "r") as file:
+                question_data = json.load(file)
+                question_id = os.path.basename(question_file).split(".")[0]
+
+            questions_processed += 1
+            answer = create_inference_request(
+                question_data, case_details, case_id, question_id, model, tokenizer
+            )
+
+            if answer is None:
+                skipped_questions += 1
+                continue
+
+            print(f"\nCase {case_id}, Question {question_id}")
+            print(f"Model Answer: {answer}")
+            print(f"Correct Answer: {question_data['answer']}")
+
+    print(f"\nInference Summary:")
+    print(f"Total Cases Processed: {cases_processed}")
+    print(f"Total Questions Processed: {questions_processed}")
+    print(f"Total Questions Skipped: {skipped_questions}")
+
+
+if __name__ == "__main__":
+    main()

experiments/benchmark_gpt4o.py

@@ -0,0 +1,331 @@
+import json
+import openai
+import os
+import glob
+import time
+import logging
+from datetime import datetime
+from tenacity import retry, wait_exponential, stop_after_attempt
+
+model_name = "chatgpt-4o-latest"
+temperature = 0.2
+log_filename = f"api_usage_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
+logging.basicConfig(filename=log_filename, level=logging.INFO, format="%(message)s")
+
+
+def calculate_cost(
+    prompt_tokens: int, completion_tokens: int, model: str = "chatgpt-4o-latest"
+) -> float:
+    """Calculate the cost of API usage based on token counts.
+
+    Args:
+        prompt_tokens: Number of tokens in the prompt
+        completion_tokens: Number of tokens in the completion
+        model: Model name to use for pricing, defaults to chatgpt-4o-latest
+
+    Returns:
+        float: Cost in USD
+    """
+    pricing = {"chatgpt-4o-latest": {"prompt": 5.0, "completion": 15.0}}
+    rates = pricing.get(model, {"prompt": 5.0, "completion": 15.0})
+    return (prompt_tokens * rates["prompt"] + completion_tokens * rates["completion"]) / 1000000
+
+
+@retry(wait=wait_exponential(multiplier=1, min=4, max=10), stop=stop_after_attempt(3))
+def create_multimodal_request(
+    question_data: dict, case_details: dict, case_id: str, question_id: str, client: openai.OpenAI
+) -> openai.types.chat.ChatCompletion:
+    """Create and send a multimodal request to the OpenAI API.
+
+    Args:
+        question_data: Dictionary containing question details and figures
+        case_details: Dictionary containing case information and figures
+        case_id: Identifier for the medical case
+        question_id: Identifier for the specific question
+        client: OpenAI client instance
+
+    Returns:
+        openai.types.chat.ChatCompletion: API response object, or None if request fails
+    """
+    prompt = f"""Given the following medical case:
+Please answer this multiple choice question:
+{question_data['question']}
+Base your answer only on the provided images and case information."""
+
+    content = [{"type": "text", "text": prompt}]
+
+    # Parse required figures
+    try:
+        # Try multiple ways of parsing figures
+        if isinstance(question_data["figures"], str):
+            try:
+                required_figures = json.loads(question_data["figures"])
+            except json.JSONDecodeError:
+                required_figures = [question_data["figures"]]
+        elif isinstance(question_data["figures"], list):
+            required_figures = question_data["figures"]
+        else:
+            required_figures = [str(question_data["figures"])]
+    except Exception as e:
+        print(f"Error parsing figures: {e}")
+        required_figures = []
+
+    # Ensure each figure starts with "Figure "
+    required_figures = [
+        fig if fig.startswith("Figure ") else f"Figure {fig}" for fig in required_figures
+    ]
+
+    subfigures = []
+    for figure in required_figures:
+        # Handle both regular figures and those with letter suffixes
+        base_figure_num = "".join(filter(str.isdigit, figure))
+        figure_letter = "".join(filter(str.isalpha, figure.split()[-1])) or None
+
+        # Find matching figures in case details
+        matching_figures = [
+            case_figure
+            for case_figure in case_details.get("figures", [])
+            if case_figure["number"] == f"Figure {base_figure_num}"
+        ]
+
+        if not matching_figures:
+            print(f"No matching figure found for {figure} in case {case_id}")
+            continue
+
+        for case_figure in matching_figures:
+            # If a specific letter is specified, filter subfigures
+            if figure_letter:
+                matching_subfigures = [
+                    subfig
+                    for subfig in case_figure.get("subfigures", [])
+                    if subfig.get("number", "").lower().endswith(figure_letter.lower())
+                    or subfig.get("label", "").lower() == figure_letter.lower()
+                ]
+                subfigures.extend(matching_subfigures)
+            else:
+                # If no letter specified, add all subfigures
+                subfigures.extend(case_figure.get("subfigures", []))
+
+    # Add images to content
+    for subfig in subfigures:
+        if "url" in subfig:
+            content.append({"type": "image_url", "image_url": {"url": subfig["url"]}})
+        else:
+            print(f"Subfigure missing URL: {subfig}")
+
+    # If no images found, log and return None
+    if len(content) == 1:  # Only the text prompt exists
+        print(f"No images found for case {case_id}, question {question_id}")
+        return None
+
+    messages = [
+        {
+            "role": "system",
+            "content": "You are a medical imaging expert. Provide only the letter corresponding to your answer choice (A/B/C/D/E/F).",
+        },
+        {"role": "user", "content": content},
+    ]
+
+    if len(content) == 1:  # Only the text prompt exists
+        print(f"No images found for case {case_id}, question {question_id}")
+        log_entry = {
+            "case_id": case_id,
+            "question_id": question_id,
+            "timestamp": datetime.now().isoformat(),
+            "model": model_name,
+            "temperature": temperature,
+            "status": "skipped",
+            "reason": "no_images",
+            "cost": 0,
+            "input": {
+                "messages": messages,
+                "question_data": {
+                    "question": question_data["question"],
+                    "explanation": question_data["explanation"],
+                    "metadata": question_data.get("metadata", {}),
+                    "figures": question_data["figures"],
+                },
+                "image_urls": [subfig["url"] for subfig in subfigures if "url" in subfig],
+                "image_captions": [subfig.get("caption", "") for subfig in subfigures],
+            },
+        }
+        logging.info(json.dumps(log_entry))
+        return None
+
+    try:
+        start_time = time.time()
+
+        response = client.chat.completions.create(
+            model=model_name, messages=messages, max_tokens=50, temperature=temperature
+        )
+        duration = time.time() - start_time
+
+        log_entry = {
+            "case_id": case_id,
+            "question_id": question_id,
+            "timestamp": datetime.now().isoformat(),
+            "model": model_name,
+            "temperature": temperature,
+            "duration": round(duration, 2),
+            "usage": {
+                "prompt_tokens": response.usage.prompt_tokens,
+                "completion_tokens": response.usage.completion_tokens,
+                "total_tokens": response.usage.total_tokens,
+            },
+            "cost": calculate_cost(response.usage.prompt_tokens, response.usage.completion_tokens),
+            "model_answer": response.choices[0].message.content,
+            "correct_answer": question_data["answer"],
+            "input": {
+                "messages": messages,
+                "question_data": {
+                    "question": question_data["question"],
+                    "explanation": question_data["explanation"],
+                    "metadata": question_data.get("metadata", {}),
+                    "figures": question_data["figures"],
+                },
+                "image_urls": [subfig["url"] for subfig in subfigures if "url" in subfig],
+                "image_captions": [subfig.get("caption", "") for subfig in subfigures],
+            },
+        }
+        logging.info(json.dumps(log_entry))
+        return response
+
+    except openai.RateLimitError:
+        log_entry = {
+            "case_id": case_id,
+            "question_id": question_id,
+            "timestamp": datetime.now().isoformat(),
+            "model": model_name,
+            "temperature": temperature,
+            "status": "error",
+            "reason": "rate_limit",
+            "cost": 0,
+            "input": {
+                "messages": messages,
+                "question_data": {
+                    "question": question_data["question"],
+                    "explanation": question_data["explanation"],
+                    "metadata": question_data.get("metadata", {}),
+                    "figures": question_data["figures"],
+                },
+                "image_urls": [subfig["url"] for subfig in subfigures if "url" in subfig],
+                "image_captions": [subfig.get("caption", "") for subfig in subfigures],
+            },
+        }
+        logging.info(json.dumps(log_entry))
+        print(
+            f"\nRate limit hit for case {case_id}, question {question_id}. Waiting 20s...",
+            flush=True,
+        )
+        time.sleep(20)
+        raise
+    except Exception as e:
+        log_entry = {
+            "case_id": case_id,
+            "question_id": question_id,
+            "timestamp": datetime.now().isoformat(),
+            "model": model_name,
+            "temperature": temperature,
+            "status": "error",
+            "error": str(e),
+            "cost": 0,
+            "input": {
+                "messages": messages,
+                "question_data": {
+                    "question": question_data["question"],
+                    "explanation": question_data["explanation"],
+                    "metadata": question_data.get("metadata", {}),
+                    "figures": question_data["figures"],
+                },
+                "image_urls": [subfig["url"] for subfig in subfigures if "url" in subfig],
+                "image_captions": [subfig.get("caption", "") for subfig in subfigures],
+            },
+        }
+        logging.info(json.dumps(log_entry))
+        print(f"Error processing case {case_id}, question {question_id}: {str(e)}")
+        raise
+
+
+def load_benchmark_questions(case_id: str) -> list:
+    """Load benchmark questions for a given case.
+
+    Args:
+        case_id: Identifier for the medical case
+
+    Returns:
+        list: List of paths to question files
+    """
+    benchmark_dir = "../benchmark/questions"
+    return glob.glob(f"{benchmark_dir}/{case_id}/{case_id}_*.json")
+
+
+def count_total_questions() -> tuple[int, int]:
+    """Count total number of cases and questions in benchmark.
+
+    Returns:
+        tuple: (total_cases, total_questions)
+    """
+    total_cases = len(glob.glob("../benchmark/questions/*"))
+    total_questions = sum(
+        len(glob.glob(f"../benchmark/questions/{case_id}/*.json"))
+        for case_id in os.listdir("../benchmark/questions")
+    )
+    return total_cases, total_questions
+
+
+def main() -> None:
+    """Main function to run the benchmark evaluation."""
+    with open("../data/eurorad_metadata.json", "r") as file:
+        data = json.load(file)
+
+    api_key = os.getenv("OPENAI_API_KEY")
+    if not api_key:
+        raise ValueError("OPENAI_API_KEY environment variable is not set.")
+    global client
+    client = openai.OpenAI(api_key=api_key)
+
+    total_cases, total_questions = count_total_questions()
+    cases_processed = 0
+    questions_processed = 0
+    skipped_questions = 0
+
+    print(f"Beginning benchmark evaluation for model {model_name} with temperature {temperature}")
+
+    for case_id, case_details in data.items():
+        question_files = load_benchmark_questions(case_id)
+        if not question_files:
+            continue
+
+        cases_processed += 1
+        for question_file in question_files:
+            with open(question_file, "r") as file:
+                question_data = json.load(file)
+                question_id = os.path.basename(question_file).split(".")[0]
+
+            questions_processed += 1
+            response = create_multimodal_request(
+                question_data, case_details, case_id, question_id, client
+            )
+
+            # Handle cases where response is None
+            if response is None:
+                skipped_questions += 1
+                print(f"Skipped question: Case ID {case_id}, Question ID {question_id}")
+                continue
+
+            print(
+                f"Progress: Case {cases_processed}/{total_cases}, Question {questions_processed}/{total_questions}"
+            )
+            print(f"Case ID: {case_id}")
+            print(f"Question ID: {question_id}")
+            print(f"Model Answer: {response.choices[0].message.content}")
+            print(f"Correct Answer: {question_data['answer']}\n")
+
+    print(f"\nBenchmark Summary:")
+    print(f"Total Cases Processed: {cases_processed}")
+    print(f"Total Questions Processed: {questions_processed}")
+    print(f"Total Questions Skipped: {skipped_questions}")
+
+
+if __name__ == "__main__":
+    main()

experiments/benchmark_llama.py

@@ -0,0 +1,443 @@
+from typing import Dict, List, Optional, Any, Union
+import re
+import json
+import os
+import glob
+import time
+import logging
+import socket
+import requests
+import httpx
+import backoff
+from datetime import datetime
+from tenacity import retry, wait_exponential, stop_after_attempt
+from openai import OpenAI
+
+# Configure model settings
+MODEL_NAME = "meta-llama/llama-3.2-90b-vision-instruct"
+temperature = 0.2
+
+# Configure logging
+log_filename = f"api_usage_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
+logging.basicConfig(filename=log_filename, level=logging.INFO, format="%(message)s")
+
+
+def verify_dns() -> bool:
+    """Verify DNS resolution and connectivity.
+
+    Returns:
+        bool: True if DNS resolution succeeds, False otherwise
+    """
+    try:
+        # Try to resolve openrouter.ai
+        socket.gethostbyname("openrouter.ai")
+        return True
+    except socket.gaierror:
+        print("DNS resolution failed. Trying to use Google DNS (8.8.8.8)...")
+        # Modify resolv.conf to use Google DNS
+        try:
+            with open("/etc/resolv.conf", "w") as f:
+                f.write("nameserver 8.8.8.8\n")
+            return True
+        except Exception as e:
+            print(f"Failed to update DNS settings: {e}")
+            return False
+
+
+def verify_connection() -> bool:
+    """Verify connection to OpenRouter API.
+
+    Returns:
+        bool: True if connection succeeds, False otherwise
+    """
+    try:
+        response = requests.get("https://openrouter.ai/api/v1/status", timeout=10)
+        return response.status_code == 200
+    except Exception as e:
+        print(f"Connection test failed: {e}")
+        return False
+
+
+def initialize_client() -> OpenAI:
+    """Initialize the OpenRouter client with proper timeout settings and connection verification.
+
+    Returns:
+        OpenAI: Configured OpenAI client for OpenRouter
+
+    Raises:
+        ValueError: If OPENROUTER_API_KEY environment variable is not set
+        ConnectionError: If DNS verification or connection test fails
+    """
+    api_key = os.getenv("OPENROUTER_API_KEY")
+    if not api_key:
+        raise ValueError("OPENROUTER_API_KEY environment variable is not set.")
+
+    # Configure timeout settings for the client
+    timeout_settings = 120  # Increased timeout for large images/responses
+
+    # Verify DNS and connection
+    if not verify_dns():
+        raise ConnectionError("DNS verification failed. Please check your network settings.")
+
+    if not verify_connection():
+        raise ConnectionError(
+            "Cannot connect to OpenRouter. Please check your internet connection."
+        )
+
+    # Set up client with retry and timeout settings
+    return OpenAI(
+        base_url="https://openrouter.ai/api/v1",
+        api_key=api_key,
+        timeout=timeout_settings,
+        http_client=httpx.Client(
+            timeout=timeout_settings, transport=httpx.HTTPTransport(retries=3)
+        ),
+    )
+
+
+@backoff.on_exception(
+    backoff.expo,
+    (ConnectionError, TimeoutError, socket.gaierror, httpx.ConnectError),
+    max_tries=5,
+    max_time=300,  # Maximum total time to try in seconds
+)
+def create_multimodal_request(
+    question_data: Dict[str, Any],
+    case_details: Dict[str, Any],
+    case_id: str,
+    question_id: str,
+    client: OpenAI,
+) -> Optional[Any]:
+    """Create and send a multimodal request to the model.
+
+    Args:
+        question_data: Dictionary containing question details
+        case_details: Dictionary containing case information
+        case_id: ID of the medical case
+        question_id: ID of the specific question
+        client: OpenAI client instance
+
+    Returns:
+        Optional[Any]: Model response if successful, None if skipped
+
+    Raises:
+        ConnectionError: If connection fails
+        TimeoutError: If request times out
+        Exception: For other errors
+    """
+
+    system_prompt = """You are a medical imaging expert. Your task is to provide ONLY a single letter answer.
+Rules:
+1. Respond with exactly one uppercase letter (A/B/C/D/E/F)
+2. Do not add periods, explanations, or any other text
+3. Do not use markdown or formatting
+4. Do not restate the question
+5. Do not explain your reasoning
+
+Examples of valid responses:
+A
+B
+C
+
+Examples of invalid responses:
+"A."
+"Answer: B"
+"C) This shows..."
+"The answer is D"
+"""
+
+    prompt = f"""Given the following medical case:
+Please answer this multiple choice question:
+{question_data['question']}
+Base your answer only on the provided images and case information."""
+
+    # Parse required figures
+    try:
+        if isinstance(question_data["figures"], str):
+            try:
+                required_figures = json.loads(question_data["figures"])
+            except json.JSONDecodeError:
+                required_figures = [question_data["figures"]]
+        elif isinstance(question_data["figures"], list):
+            required_figures = question_data["figures"]
+        else:
+            required_figures = [str(question_data["figures"])]
+    except Exception as e:
+        print(f"Error parsing figures: {e}")
+        required_figures = []
+
+    required_figures = [
+        fig if fig.startswith("Figure ") else f"Figure {fig}" for fig in required_figures
+    ]
+
+    # Process subfigures and prepare content
+    content = [{"type": "text", "text": prompt}]
+    image_urls = []
+    image_captions = []
+
+    for figure in required_figures:
+        base_figure_num = "".join(filter(str.isdigit, figure))
+        figure_letter = "".join(filter(str.isalpha, figure.split()[-1])) or None
+
+        matching_figures = [
+            case_figure
+            for case_figure in case_details.get("figures", [])
+            if case_figure["number"] == f"Figure {base_figure_num}"
+        ]
+
+        for case_figure in matching_figures:
+            subfigures = []
+            if figure_letter:
+                subfigures = [
+                    subfig
+                    for subfig in case_figure.get("subfigures", [])
+                    if subfig.get("number", "").lower().endswith(figure_letter.lower())
+                    or subfig.get("label", "").lower() == figure_letter.lower()
+                ]
+            else:
+                subfigures = case_figure.get("subfigures", [])
+
+            for subfig in subfigures:
+                if "url" in subfig:
+                    content.append({"type": "image_url", "image_url": {"url": subfig["url"]}})
+                    image_urls.append(subfig["url"])
+                    image_captions.append(subfig.get("caption", ""))
+
+    if len(content) == 1:  # Only the text prompt exists
+        print(f"No images found for case {case_id}, question {question_id}")
+        # Log the skipped question
+        log_entry = {
+            "case_id": case_id,
+            "question_id": question_id,
+            "timestamp": datetime.now().isoformat(),
+            "model": MODEL_NAME,
+            "status": "skipped",
+            "reason": "no_images",
+            "input": {
+                "question_data": {
+                    "question": question_data["question"],
+                    "explanation": question_data["explanation"],
+                    "metadata": question_data.get("metadata", {}),
+                    "figures": question_data["figures"],
+                },
+                "image_urls": image_urls,
+            },
+        }
+        logging.info(json.dumps(log_entry))
+        return None
+
+    try:
+        start_time = time.time()
+
+        response = client.chat.completions.create(
+            model=MODEL_NAME,
+            temperature=temperature,
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": content},
+            ],
+        )
+        duration = time.time() - start_time
+
+        # Get raw response
+        raw_answer = response.choices[0].message.content
+
+        # Validate and clean
+        clean_answer = validate_answer(raw_answer)
+
+        if not clean_answer:
+            print(f"Warning: Invalid response format for case {case_id}, question {question_id}")
+            print(f"Raw response: {raw_answer}")
+
+        # Update response object with cleaned answer
+        response.choices[0].message.content = clean_answer
+
+        # Log response
+        log_entry = {
+            "case_id": case_id,
+            "question_id": question_id,
+            "timestamp": datetime.now().isoformat(),
+            "model": MODEL_NAME,
+            "temperature": temperature,
+            "duration": round(duration, 2),
+            "usage": {
+                "prompt_tokens": response.usage.prompt_tokens,
+                "completion_tokens": response.usage.completion_tokens,
+                "total_tokens": response.usage.total_tokens,
+            },
+            "model_answer": response.choices[0].message.content,
+            "correct_answer": question_data["answer"],
+            "input": {
+                "question_data": {
+                    "question": question_data["question"],
+                    "explanation": question_data["explanation"],
+                    "metadata": question_data.get("metadata", {}),
+                    "figures": question_data["figures"],
+                },
+                "image_urls": image_urls,
+            },
+        }
+        logging.info(json.dumps(log_entry))
+        return response
+
+    except ConnectionError as e:
+        print(f"Connection error for case {case_id}, question {question_id}: {str(e)}")
+        print("Retrying after a longer delay...")
+        time.sleep(30)  # Add a longer delay before retry
+        raise
+    except TimeoutError as e:
+        print(f"Timeout error for case {case_id}, question {question_id}: {str(e)}")
+        print("Retrying with increased timeout...")
+        raise
+    except Exception as e:
+        # Log failed requests too
+        log_entry = {
+            "case_id": case_id,
+            "question_id": question_id,
+            "timestamp": datetime.now().isoformat(),
+            "model": MODEL_NAME,
+            "temperature": temperature,
+            "status": "error",
+            "error": str(e),
+            "input": {
+                "question_data": {
+                    "question": question_data["question"],
+                    "explanation": question_data["explanation"],
+                    "metadata": question_data.get("metadata", {}),
+                    "figures": question_data["figures"],
+                },
+                "image_urls": image_urls,
+            },
+        }
+        logging.info(json.dumps(log_entry))
+        raise
+
+
+def extract_answer(response_text: str) -> Optional[str]:
+    """Extract single letter answer from model response.
+
+    Args:
+        response_text: Raw text response from model
+
+    Returns:
+        Optional[str]: Single letter answer if found, None otherwise
+    """
+    # Convert to uppercase and remove periods
+    text = response_text.upper().replace(".", "")
+
+    # Look for common patterns
+    patterns = [
+        r"ANSWER:\s*([A-F])",  # Matches "ANSWER: X"
+        r"OPTION\s*([A-F])",  # Matches "OPTION X"
+        r"([A-F])\)",  # Matches "X)"
+        r"\b([A-F])\b",  # Matches single letter
+    ]
+
+    for pattern in patterns:
+        matches = re.findall(pattern, text)
+        if matches:
+            return matches[0]
+
+    return None
+
+
+def validate_answer(response_text: str) -> Optional[str]:
+    """Enforce strict single-letter response format.
+
+    Args:
+        response_text: Raw text response from model
+
+    Returns:
+        Optional[str]: Valid single letter answer if found, None otherwise
+    """
+    if not response_text:
+        return None
+
+    # Remove all whitespace and convert to uppercase
+    cleaned = response_text.strip().upper()
+
+    # Check if it's exactly one valid letter
+    if len(cleaned) == 1 and cleaned in "ABCDEF":
+        return cleaned
+
+    # If not, try to extract just the letter
+    match = re.search(r"([A-F])", cleaned)
+    return match.group(1) if match else None
+
+
+def load_benchmark_questions(case_id: str) -> List[str]:
+    """Find all question files for a given case ID.
+
+    Args:
+        case_id: ID of the medical case
+
+    Returns:
+        List[str]: List of paths to question files
+    """
+    benchmark_dir = "../benchmark/questions"
+    return glob.glob(f"{benchmark_dir}/{case_id}/{case_id}_*.json")
+
+
+def count_total_questions() -> Tuple[int, int]:
+    """Count total number of cases and questions.
+
+    Returns:
+        Tuple[int, int]: (total_cases, total_questions)
+    """
+    total_cases = len(glob.glob("../benchmark/questions/*"))
+    total_questions = sum(
+        len(glob.glob(f"../benchmark/questions/{case_id}/*.json"))
+        for case_id in os.listdir("../benchmark/questions")
+    )
+    return total_cases, total_questions
+
+
+def main():
+    with open("../data/eurorad_metadata.json", "r") as file:
+        data = json.load(file)
+
+    client = initialize_client()
+    total_cases, total_questions = count_total_questions()
+    cases_processed = 0
+    questions_processed = 0
+    skipped_questions = 0
+
+    print(f"Beginning benchmark evaluation for {MODEL_NAME} with temperature {temperature}")
+
+    for case_id, case_details in data.items():
+        question_files = load_benchmark_questions(case_id)
+        if not question_files:
+            continue
+
+        cases_processed += 1
+        for question_file in question_files:
+            with open(question_file, "r") as file:
+                question_data = json.load(file)
+                question_id = os.path.basename(question_file).split(".")[0]
+
+            questions_processed += 1
+            response = create_multimodal_request(
+                question_data, case_details, case_id, question_id, client
+            )
+
+            if response is None:
+                skipped_questions += 1
+                print(f"Skipped question: Case ID {case_id}, Question ID {question_id}")
+                continue
+
+            print(
+                f"Progress: Case {cases_processed}/{total_cases}, Question {questions_processed}/{total_questions}"
+            )
+            print(f"Case ID: {case_id}")
+            print(f"Question ID: {question_id}")
+            print(f"Model Answer: {response.choices[0].message.content}")
+            print(f"Correct Answer: {question_data['answer']}\n")
+
+    print(f"\nBenchmark Summary:")
+    print(f"Total Cases Processed: {cases_processed}")
+    print(f"Total Questions Processed: {questions_processed}")
+    print(f"Total Questions Skipped: {skipped_questions}")
+
+
+if __name__ == "__main__":
+    main()

experiments/benchmark_llavamed.py

@@ -0,0 +1,541 @@
+import argparse
+import json
+import requests
+import base64
+from PIL import Image
+from io import BytesIO
+from llava.conversation import conv_templates
+import time
+import os
+import glob
+import logging
+from datetime import datetime
+from tqdm import tqdm
+import re
+from typing import Dict, List, Optional, Union, Any, Tuple
+
+
+def process_image(image_path: str, target_size: int = 640) -> Image.Image:
+    """Process and resize an image to match model requirements.
+
+    Args:
+        image_path: Path to the input image file
+        target_size: Target size for both width and height in pixels
+
+    Returns:
+        PIL.Image: Processed and padded image with dimensions (target_size, target_size)
+    """
+    image = Image.open(image_path)
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+
+    # Calculate scaling to maintain aspect ratio
+    ratio = min(target_size / image.width, target_size / image.height)
+    new_size = (int(image.width * ratio), int(image.height * ratio))
+
+    # Resize image
+    image = image.resize(new_size, Image.LANCZOS)
+
+    # Create new image with padding
+    new_image = Image.new("RGB", (target_size, target_size), (0, 0, 0))
+    # Paste resized image in center
+    offset = ((target_size - new_size[0]) // 2, (target_size - new_size[1]) // 2)
+    new_image.paste(image, offset)
+
+    return new_image
+
+
+def validate_answer(response_text: str) -> Optional[str]:
+    """Extract and validate a single-letter response from the model's output.
+    Handles multiple response formats and edge cases.
+
+    Args:
+        response_text: The full text output from the model
+
+    Returns:
+        A single letter answer (A-F) or None if no valid answer found
+    """
+    if not response_text:
+        return None
+
+    # Clean the response text
+    cleaned = response_text.strip()
+
+    # Comprehensive set of patterns to extract the answer
+    extraction_patterns = [
+        # Strict format with explicit letter answer
+        r"(?:THE\s*)?(?:SINGLE\s*)?LETTER\s*(?:ANSWER\s*)?(?:IS:?)\s*([A-F])\b",
+        # Patterns for extracting from longer descriptions
+        r"(?:correct\s+)?(?:answer|option)\s*(?:is\s*)?([A-F])\b",
+        r"\b(?:answer|option)\s*([A-F])[):]\s*",
+        # Patterns for extracting from descriptive sentences
+        r"(?:most\s+likely\s+)?(?:answer|option)\s*(?:is\s*)?([A-F])\b",
+        r"suggest[s]?\s+(?:that\s+)?(?:the\s+)?(?:answer\s+)?(?:is\s*)?([A-F])\b",
+        # Patterns with contextual words
+        r"characteriz[e]?d?\s+by\s+([A-F])\b",
+        r"indicat[e]?s?\s+([A-F])\b",
+        # Fallback to Option X or Letterr X formats
+        r"Option\s*([A-F])\b",
+        r"\b([A-F])\)\s*",
+        # Fallback to standalone letter
+        r"^\s*([A-F])\s*$",
+    ]
+
+    # Try each pattern
+    for pattern in extraction_patterns:
+        matches = re.findall(pattern, cleaned, re.IGNORECASE)
+        for match in matches:
+            # Ensure match is a single valid letter
+            if isinstance(match, tuple):
+                match = match[0] if match[0] in "ABCDEF" else None
+            if match and match.upper() in "ABCDEF":
+                return match.upper()
+
+    # Final fallback: look for standalone letters in context
+    context_matches = re.findall(r"\b([A-F])\b", cleaned.upper())
+    context_letters = [m for m in context_matches if m in "ABCDEF"]
+    if context_letters:
+        return context_letters[0]
+
+    # No valid answer found
+    return None
+
+
+def load_benchmark_questions(case_id: str) -> List[str]:
+    """Find all question files for a given case ID.
+
+    Args:
+        case_id: The ID of the medical case
+
+    Returns:
+        List of paths to question JSON files
+    """
+    benchmark_dir = "MedMAX/benchmark/questions"
+    return glob.glob(f"{benchmark_dir}/{case_id}/{case_id}_*.json")
+
+
+def count_total_questions() -> Tuple[int, int]:
+    """Count total number of cases and questions in benchmark.
+
+    Returns:
+        Tuple containing (total_cases, total_questions)
+    """
+    total_cases = len(glob.glob("MedMAX/benchmark/questions/*"))
+    total_questions = sum(
+        len(glob.glob(f"MedMAX/benchmark/questions/{case_id}/*.json"))
+        for case_id in os.listdir("MedMAX/benchmark/questions")
+    )
+    return total_cases, total_questions
+
+
+def create_inference_request(
+    question_data: Dict[str, Any],
+    case_details: Dict[str, Any],
+    case_id: str,
+    question_id: str,
+    worker_addr: str,
+    model_name: str,
+    raw_output: bool = False,
+) -> Union[Tuple[Optional[str], Optional[float]], Dict[str, Any]]:
+    """Create and send inference request to worker.
+
+    Args:
+        question_data: Dictionary containing question details and figures
+        case_details: Dictionary containing case information and figures
+        case_id: Identifier for the medical case
+        question_id: Identifier for the specific question
+        worker_addr: Address of the worker endpoint
+        model_name: Name of the model to use
+        raw_output: Whether to return raw model output
+
+    Returns:
+        If raw_output is False: Tuple of (validated_answer, duration)
+        If raw_output is True: Dictionary with full inference details
+    """
+    system_prompt = """You are a medical imaging expert. Your answer MUST be a SINGLE LETTER (A/B/C/D/E/F), provided in this format: 'The SINGLE LETTER answer is: X'.
+"""
+
+    prompt = f"""Given the following medical case:
+Please answer this multiple choice question:
+{question_data['question']}
+Base your answer only on the provided images and case information. Respond with your SINGLE LETTER answer: """
+
+    try:
+        # Parse required figures
+        if isinstance(question_data["figures"], str):
+            try:
+                required_figures = json.loads(question_data["figures"])
+            except json.JSONDecodeError:
+                required_figures = [question_data["figures"]]
+        elif isinstance(question_data["figures"], list):
+            required_figures = question_data["figures"]
+        else:
+            required_figures = [str(question_data["figures"])]
+    except Exception as e:
+        print(f"Error parsing figures: {e}")
+        required_figures = []
+
+    required_figures = [
+        fig if fig.startswith("Figure ") else f"Figure {fig}" for fig in required_figures
+    ]
+
+    # Get image paths
+    image_paths = []
+    for figure in required_figures:
+        base_figure_num = "".join(filter(str.isdigit, figure))
+        figure_letter = "".join(filter(str.isalpha, figure.split()[-1])) or None
+
+        matching_figures = [
+            case_figure
+            for case_figure in case_details.get("figures", [])
+            if case_figure["number"] == f"Figure {base_figure_num}"
+        ]
+
+        for case_figure in matching_figures:
+            subfigures = []
+            if figure_letter:
+                subfigures = [
+                    subfig
+                    for subfig in case_figure.get("subfigures", [])
+                    if subfig.get("number", "").lower().endswith(figure_letter.lower())
+                    or subfig.get("label", "").lower() == figure_letter.lower()
+                ]
+            else:
+                subfigures = case_figure.get("subfigures", [])
+
+            for subfig in subfigures:
+                if "local_path" in subfig:
+                    image_paths.append("MedMAX/data/" + subfig["local_path"])
+
+    if not image_paths:
+        print(f"No local images found for case {case_id}, question {question_id}")
+        return "skipped", 0.0  # Return a special 'skipped' marker
+
+    try:
+        start_time = time.time()
+
+        # Process each image
+        processed_images = [process_image(path) for path in image_paths]
+
+        # Create conversation
+        conv = conv_templates["mistral_instruct"].copy()
+
+        # Add image and message
+        if "<image>" not in prompt:
+            text = prompt + "\n<image>"
+        else:
+            text = prompt
+
+        message = (text, processed_images[0], "Default")  # Currently handling first image
+        conv.append_message(conv.roles[0], message)
+        conv.append_message(conv.roles[1], None)
+
+        prompt = conv.get_prompt()
+        headers = {"User-Agent": "LLaVA-Med Client"}
+        pload = {
+            "model": model_name,
+            "prompt": prompt,
+            "max_new_tokens": 150,  # Reduce this since we only need one letter
+            "temperature": 0.5,  # Lower temperature for more focused responses
+            "stop": conv.sep2,
+            "images": conv.get_images(),
+            "top_p": 1,  # Lower top_p for more focused sampling
+            "frequency_penalty": 0.0,
+            "presence_penalty": 0.0,
+        }
+
+        max_retries = 3
+        retry_delay = 5
+        response_text = None
+
+        for attempt in range(max_retries):
+            try:
+                response = requests.post(
+                    worker_addr + "/worker_generate_stream",
+                    headers=headers,
+                    json=pload,
+                    stream=True,
+                    timeout=30,
+                )
+
+                complete_output = ""
+                for chunk in response.iter_lines(
+                    chunk_size=8192, decode_unicode=False, delimiter=b"\0"
+                ):
+                    if chunk:
+                        data = json.loads(chunk.decode("utf-8"))
+                        if data["error_code"] == 0:
+                            output = data["text"].split("[/INST]")[-1]
+                            complete_output = output
+                        else:
+                            print(f"\nError: {data['text']} (error_code: {data['error_code']})")
+                            if attempt < max_retries - 1:
+                                time.sleep(retry_delay)
+                                break
+                            return None, None
+
+                if complete_output:
+                    response_text = complete_output
+                    break
+
+            except (requests.exceptions.RequestException, json.JSONDecodeError) as e:
+                if attempt < max_retries - 1:
+                    print(f"\nNetwork error: {str(e)}. Retrying in {retry_delay} seconds...")
+                    time.sleep(retry_delay)
+                else:
+                    print(f"\nFailed after {max_retries} attempts: {str(e)}")
+                    return None, None
+
+        duration = time.time() - start_time
+
+        if raw_output:
+            inference_details = {
+                "raw_output": response_text,
+                "validated_answer": validate_answer(response_text),
+                "duration": duration,
+                "prompt": prompt,
+                "system_prompt": system_prompt,
+                "image_paths": image_paths,
+                "payload": pload,
+            }
+            return inference_details
+
+        return validate_answer(response_text), duration
+
+    except Exception as e:
+        print(f"Error in inference request: {str(e)}")
+        return None, None
+
+
+def clean_payload(payload: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
+    """Remove image-related and large data from the payload to keep the log lean.
+
+    Args:
+        payload: Original request payload dictionary
+
+    Returns:
+        Cleaned payload dictionary with large data removed
+    """
+    if not payload:
+        return None
+
+    # Create a copy of the payload to avoid modifying the original
+    cleaned_payload = payload.copy()
+
+    # Remove large or sensitive data
+    if "images" in cleaned_payload:
+        del cleaned_payload["images"]
+
+    return cleaned_payload
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--controller-address", type=str, default="http://localhost:21001")
+    parser.add_argument("--worker-address", type=str)
+    parser.add_argument("--model-name", type=str, default="llava-med-v1.5-mistral-7b")
+    parser.add_argument("--output-dir", type=str, default="benchmark_results")
+    parser.add_argument(
+        "--raw-output", action="store_true", help="Return raw model output without validation"
+    )
+    parser.add_argument(
+        "--num-cases",
+        type=int,
+        help="Number of cases to process if looking at raw outputs",
+        default=2,
+    )
+    args = parser.parse_args()
+
+    # Setup output directory
+    os.makedirs(args.output_dir, exist_ok=True)
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+
+    # Setup live logging files
+    live_log_filename = os.path.join(args.output_dir, f"live_benchmark_log_{timestamp}.json")
+    final_results_filename = os.path.join(args.output_dir, f"final_results_{timestamp}.json")
+
+    # Initialize live log file
+    with open(live_log_filename, "w") as live_log_file:
+        live_log_file.write("[\n")  # Start of JSON array
+
+    # Setup logging
+    logging.basicConfig(
+        filename=os.path.join(args.output_dir, f"benchmark_{timestamp}.log"),
+        level=logging.INFO,
+        format="%(message)s",
+    )
+
+    # Get worker address
+    if args.worker_address:
+        worker_addr = args.worker_address
+    else:
+        try:
+            requests.post(args.controller_address + "/refresh_all_workers")
+            ret = requests.post(args.controller_address + "/list_models")
+            models = ret.json()["models"]
+            ret = requests.post(
+                args.controller_address + "/get_worker_address", json={"model": args.model_name}
+            )
+            worker_addr = ret.json()["address"]
+            print(f"Worker address: {worker_addr}")
+        except requests.exceptions.RequestException as e:
+            print(f"Failed to connect to controller: {e}")
+            return
+
+    if worker_addr == "":
+        print("No available worker")
+        return
+
+    # Load cases with local paths
+    with open("MedMAX/data/updated_cases.json", "r") as file:
+        data = json.load(file)
+
+    total_cases, total_questions = count_total_questions()
+    print(f"\nStarting benchmark with {args.model_name}")
+    print(f"Found {total_cases} cases with {total_questions} total questions")
+
+    results = {
+        "model": args.model_name,
+        "timestamp": datetime.now().isoformat(),
+        "total_cases": total_cases,
+        "total_questions": total_questions,
+        "results": [],
+    }
+
+    cases_processed = 0
+    questions_processed = 0
+    correct_answers = 0
+    skipped_questions = 0
+    total_processed_entries = 0
+
+    # Process each case
+    for case_id, case_details in tqdm(data.items(), desc="Processing cases"):
+        question_files = load_benchmark_questions(case_id)
+        if not question_files:
+            continue
+
+        cases_processed += 1
+        for question_file in tqdm(
+            question_files, desc=f"Processing questions for case {case_id}", leave=False
+        ):
+            with open(question_file, "r") as file:
+                question_data = json.load(file)
+                question_id = os.path.basename(question_file).split(".")[0]
+
+            questions_processed += 1
+
+            # Get model's answer
+            inference_result = create_inference_request(
+                question_data,
+                case_details,
+                case_id,
+                question_id,
+                worker_addr,
+                args.model_name,
+                raw_output=True,  # Always use raw output for detailed logging
+            )
+
+            # Handle skipped questions
+            if inference_result == ("skipped", 0.0):
+                skipped_questions += 1
+                print(f"\nCase {case_id}, Question {question_id}: Skipped (No images)")
+
+                # Log skipped question
+                skipped_entry = {
+                    "case_id": case_id,
+                    "question_id": question_id,
+                    "status": "skipped",
+                    "reason": "No images found",
+                }
+                with open(live_log_filename, "a") as live_log_file:
+                    json.dump(skipped_entry, live_log_file, indent=2)
+                    live_log_file.write(",\n")  # Add comma for next entry
+
+                continue
+
+            # Extract information
+            answer = inference_result["validated_answer"]
+            duration = inference_result["duration"]
+
+            # Prepare detailed logging entry
+            log_entry = {
+                "case_id": case_id,
+                "question_id": question_id,
+                "question": question_data["question"],
+                "correct_answer": question_data["answer"],
+                "raw_output": inference_result["raw_output"],
+                "validated_answer": answer,
+                "model_answer": answer,
+                "is_correct": answer == question_data["answer"] if answer else False,
+                "duration": duration,
+                "system_prompt": inference_result["system_prompt"],
+                "input_prompt": inference_result["prompt"],
+                "image_paths": inference_result["image_paths"],
+                "payload": clean_payload(inference_result["payload"]),
+            }
+
+            # Write to live log file
+            with open(live_log_filename, "a") as live_log_file:
+                json.dump(log_entry, live_log_file, indent=2)
+                live_log_file.write(",\n")  # Add comma for next entry
+
+            # Print to console
+            print(f"\nCase {case_id}, Question {question_id}")
+            print(f"Model Answer: {answer}")
+            print(f"Correct Answer: {question_data['answer']}")
+            print(f"Time taken: {duration:.2f}s")
+
+            # Track correct answers
+            if answer == question_data["answer"]:
+                correct_answers += 1
+
+            # Append to results
+            results["results"].append(log_entry)
+            total_processed_entries += 1
+
+            # Optional: break if reached specified number of cases
+            if args.raw_output and cases_processed == args.num_cases:
+                break
+
+        # Optional: break if reached specified number of cases
+        if args.raw_output and cases_processed == args.num_cases:
+            break
+
+    # Close live log file
+    with open(live_log_filename, "a") as live_log_file:
+        # Remove trailing comma and close JSON array
+        live_log_file.seek(live_log_file.tell() - 2, 0)  # Go back 2 chars to remove ',\n'
+        live_log_file.write("\n]")
+
+    # Calculate final statistics
+    results["summary"] = {
+        "cases_processed": cases_processed,
+        "questions_processed": questions_processed,
+        "total_processed_entries": total_processed_entries,
+        "correct_answers": correct_answers,
+        "skipped_questions": skipped_questions,
+        "accuracy": (
+            correct_answers / (questions_processed - skipped_questions)
+            if (questions_processed - skipped_questions) > 0
+            else 0
+        ),
+    }
+
+    # Save final results
+    with open(final_results_filename, "w") as f:
+        json.dump(results, f, indent=2)
+
+    print(f"\nBenchmark Summary:")
+    print(f"Total Cases Processed: {cases_processed}")
+    print(f"Total Questions Processed: {questions_processed}")
+    print(f"Total Processed Entries: {total_processed_entries}")
+    print(f"Correct Answers: {correct_answers}")
+    print(f"Skipped Questions: {skipped_questions}")
+    print(f"Accuracy: {(correct_answers / (questions_processed - skipped_questions) * 100):.2f}%")
+    print(f"\nResults saved to {args.output_dir}")
+    print(f"Live log: {live_log_filename}")
+    print(f"Final results: {final_results_filename}")
+
+
+if __name__ == "__main__":
+    main()

experiments/benchmark_medrax.ipynb

@@ -0,0 +1,374 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import operator\n",
+    "import warnings\n",
+    "from typing import *\n",
+    "import traceback\n",
+    "\n",
+    "import os\n",
+    "import torch\n",
+    "from dotenv import load_dotenv\n",
+    "from IPython.display import Image\n",
+    "from langgraph.checkpoint.memory import MemorySaver\n",
+    "from langgraph.graph import END, StateGraph\n",
+    "from langchain_core.messages import AnyMessage, HumanMessage, SystemMessage, ToolMessage\n",
+    "from langchain_openai import ChatOpenAI\n",
+    "from transformers import logging\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "import re\n",
+    "\n",
+    "from medrax.agent import *\n",
+    "from medrax.tools import *\n",
+    "from medrax.utils import *\n",
+    "\n",
+    "import json\n",
+    "import openai\n",
+    "import os\n",
+    "import glob\n",
+    "import time\n",
+    "import logging\n",
+    "from datetime import datetime\n",
+    "from tenacity import retry, wait_exponential, stop_after_attempt\n",
+    "\n",
+    "warnings.filterwarnings(\"ignore\")\n",
+    "_ = load_dotenv()\n",
+    "\n",
+    "\n",
+    "# Setup directory paths\n",
+    "ROOT = \"set this directory to where MedRAX is, .e.g /home/MedRAX\"\n",
+    "PROMPT_FILE = f\"{ROOT}/medrax/docs/system_prompts.txt\"\n",
+    "BENCHMARK_FILE = f\"{ROOT}/benchmark/questions\"\n",
+    "MODEL_DIR = f\"set this to where the tool models are, e.g /home/models\"\n",
+    "FIGURES_DIR = f\"{ROOT}/benchmark/figures\"\n",
+    "\n",
+    "model_name = \"medrax\"\n",
+    "temperature = 0.2\n",
+    "medrax_logs = f\"{ROOT}/experiments/medrax_logs\"\n",
+    "log_filename = f\"{medrax_logs}/{model_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json\"\n",
+    "logging.basicConfig(filename=log_filename, level=logging.INFO, format=\"%(message)s\", force=True)\n",
+    "device = \"cuda\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def get_tools():\n",
+    "    report_tool = ChestXRayReportGeneratorTool(cache_dir=MODEL_DIR, device=device)\n",
+    "    xray_classification_tool = ChestXRayClassifierTool(device=device)\n",
+    "    segmentation_tool = ChestXRaySegmentationTool(device=device)\n",
+    "    grounding_tool = XRayPhraseGroundingTool(\n",
+    "        cache_dir=MODEL_DIR, temp_dir=\"temp\", device=device, load_in_8bit=True\n",
+    "    )\n",
+    "    xray_vqa_tool = XRayVQATool(cache_dir=MODEL_DIR, device=device)\n",
+    "    llava_med_tool = LlavaMedTool(cache_dir=MODEL_DIR, device=device, load_in_8bit=True)\n",
+    "\n",
+    "    return [\n",
+    "        report_tool,\n",
+    "        xray_classification_tool,\n",
+    "        segmentation_tool,\n",
+    "        grounding_tool,\n",
+    "        xray_vqa_tool,\n",
+    "        llava_med_tool,\n",
+    "    ]\n",
+    "\n",
+    "\n",
+    "def get_agent(tools):\n",
+    "    prompts = load_prompts_from_file(PROMPT_FILE)\n",
+    "    prompt = prompts[\"MEDICAL_ASSISTANT\"]\n",
+    "\n",
+    "    checkpointer = MemorySaver()\n",
+    "    model = ChatOpenAI(model=\"gpt-4o\", temperature=temperature, top_p=0.95)\n",
+    "    agent = Agent(\n",
+    "        model,\n",
+    "        tools=tools,\n",
+    "        log_tools=True,\n",
+    "        log_dir=\"logs\",\n",
+    "        system_prompt=prompt,\n",
+    "        checkpointer=checkpointer,\n",
+    "    )\n",
+    "    thread = {\"configurable\": {\"thread_id\": \"1\"}}\n",
+    "    return agent, thread\n",
+    "\n",
+    "\n",
+    "def run_medrax(agent, thread, prompt, image_urls=[]):\n",
+    "    messages = [\n",
+    "        HumanMessage(\n",
+    "            content=[\n",
+    "                {\"type\": \"text\", \"text\": prompt},\n",
+    "            ]\n",
+    "            + [{\"type\": \"image_url\", \"image_url\": {\"url\": image_url}} for image_url in image_urls]\n",
+    "        )\n",
+    "    ]\n",
+    "\n",
+    "    final_response = None\n",
+    "    for event in agent.workflow.stream({\"messages\": messages}, thread):\n",
+    "        for v in event.values():\n",
+    "            final_response = v\n",
+    "\n",
+    "    final_response = final_response[\"messages\"][-1].content.strip()\n",
+    "    agent_state = agent.workflow.get_state(thread)\n",
+    "\n",
+    "    return final_response, str(agent_state)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def create_multimodal_request(question_data, case_details, case_id, question_id, agent, thread):\n",
+    "    # Parse required figures\n",
+    "    try:\n",
+    "        # Try multiple ways of parsing figures\n",
+    "        if isinstance(question_data[\"figures\"], str):\n",
+    "            try:\n",
+    "                required_figures = json.loads(question_data[\"figures\"])\n",
+    "            except json.JSONDecodeError:\n",
+    "                required_figures = [question_data[\"figures\"]]\n",
+    "        elif isinstance(question_data[\"figures\"], list):\n",
+    "            required_figures = question_data[\"figures\"]\n",
+    "        else:\n",
+    "            required_figures = [str(question_data[\"figures\"])]\n",
+    "    except Exception as e:\n",
+    "        print(f\"Error parsing figures: {e}\")\n",
+    "        required_figures = []\n",
+    "\n",
+    "    # Ensure each figure starts with \"Figure \"\n",
+    "    required_figures = [\n",
+    "        fig if fig.startswith(\"Figure \") else f\"Figure {fig}\" for fig in required_figures\n",
+    "    ]\n",
+    "\n",
+    "    subfigures = []\n",
+    "    for figure in required_figures:\n",
+    "        # Handle both regular figures and those with letter suffixes\n",
+    "        base_figure_num = \"\".join(filter(str.isdigit, figure))\n",
+    "        figure_letter = \"\".join(filter(str.isalpha, figure.split()[-1])) or None\n",
+    "\n",
+    "        # Find matching figures in case details\n",
+    "        matching_figures = [\n",
+    "            case_figure\n",
+    "            for case_figure in case_details.get(\"figures\", [])\n",
+    "            if case_figure[\"number\"] == f\"Figure {base_figure_num}\"\n",
+    "        ]\n",
+    "\n",
+    "        if not matching_figures:\n",
+    "            print(f\"No matching figure found for {figure} in case {case_id}\")\n",
+    "            continue\n",
+    "\n",
+    "        for case_figure in matching_figures:\n",
+    "            # If a specific letter is specified, filter subfigures\n",
+    "            if figure_letter:\n",
+    "                matching_subfigures = [\n",
+    "                    subfig\n",
+    "                    for subfig in case_figure.get(\"subfigures\", [])\n",
+    "                    if subfig.get(\"number\", \"\").lower().endswith(figure_letter.lower())\n",
+    "                    or subfig.get(\"label\", \"\").lower() == figure_letter.lower()\n",
+    "                ]\n",
+    "                subfigures.extend(matching_subfigures)\n",
+    "            else:\n",
+    "                # If no letter specified, add all subfigures\n",
+    "                subfigures.extend(case_figure.get(\"subfigures\", []))\n",
+    "\n",
+    "    # Add images to content\n",
+    "    figure_prompt = \"\"\n",
+    "    image_urls = []\n",
+    "\n",
+    "    for subfig in subfigures:\n",
+    "        if \"number\" in subfig:\n",
+    "            subfig_number = subfig[\"number\"].lower().strip().replace(\" \", \"_\") + \".jpg\"\n",
+    "            subfig_path = os.path.join(FIGURES_DIR, case_id, subfig_number)\n",
+    "            figure_prompt += f\"{subfig_number} located at {subfig_path}\\n\"\n",
+    "        if \"url\" in subfig:\n",
+    "            image_urls.append(subfig[\"url\"])\n",
+    "        else:\n",
+    "            print(f\"Subfigure missing URL: {subfig}\")\n",
+    "\n",
+    "    prompt = (\n",
+    "        f\"Answer this question correctly using chain of thought reasoning and \"\n",
+    "        \"carefully evaluating choices. Solve using our own vision and reasoning and then\"\n",
+    "        \"use tools to complement your reasoning. Trust your own judgement over any tools.\\n\"\n",
+    "        f\"{question_data['question']}\\n{figure_prompt}\"\n",
+    "    )\n",
+    "\n",
+    "    try:\n",
+    "        start_time = time.time()\n",
+    "\n",
+    "        final_response, agent_state = run_medrax(\n",
+    "            agent=agent, thread=thread, prompt=prompt, image_urls=image_urls\n",
+    "        )\n",
+    "        model_answer, agent_state = run_medrax(\n",
+    "            agent=agent,\n",
+    "            thread=thread,\n",
+    "            prompt=\"If you had to choose the best option, only respond with the letter of choice (only one of A, B, C, D, E, F)\",\n",
+    "        )\n",
+    "        duration = time.time() - start_time\n",
+    "\n",
+    "        log_entry = {\n",
+    "            \"case_id\": case_id,\n",
+    "            \"question_id\": question_id,\n",
+    "            \"timestamp\": datetime.now().isoformat(),\n",
+    "            \"model\": model_name,\n",
+    "            \"temperature\": temperature,\n",
+    "            \"duration\": round(duration, 2),\n",
+    "            \"usage\": \"\",\n",
+    "            \"cost\": 0,\n",
+    "            \"raw_response\": final_response,\n",
+    "            \"model_answer\": model_answer.strip(),\n",
+    "            \"correct_answer\": question_data[\"answer\"][0],\n",
+    "            \"input\": {\n",
+    "                \"messages\": prompt,\n",
+    "                \"question_data\": {\n",
+    "                    \"question\": question_data[\"question\"],\n",
+    "                    \"explanation\": question_data[\"explanation\"],\n",
+    "                    \"metadata\": question_data.get(\"metadata\", {}),\n",
+    "                    \"figures\": question_data[\"figures\"],\n",
+    "                },\n",
+    "                \"image_urls\": [subfig[\"url\"] for subfig in subfigures if \"url\" in subfig],\n",
+    "                \"image_captions\": [subfig.get(\"caption\", \"\") for subfig in subfigures],\n",
+    "            },\n",
+    "            \"agent_state\": agent_state,\n",
+    "        }\n",
+    "        logging.info(json.dumps(log_entry))\n",
+    "        return final_response, model_answer.strip()\n",
+    "\n",
+    "    except Exception as e:\n",
+    "        log_entry = {\n",
+    "            \"case_id\": case_id,\n",
+    "            \"question_id\": question_id,\n",
+    "            \"timestamp\": datetime.now().isoformat(),\n",
+    "            \"model\": model_name,\n",
+    "            \"temperature\": temperature,\n",
+    "            \"status\": \"error\",\n",
+    "            \"error\": str(e),\n",
+    "            \"cost\": 0,\n",
+    "            \"input\": {\n",
+    "                \"messages\": prompt,\n",
+    "                \"question_data\": {\n",
+    "                    \"question\": question_data[\"question\"],\n",
+    "                    \"explanation\": question_data[\"explanation\"],\n",
+    "                    \"metadata\": question_data.get(\"metadata\", {}),\n",
+    "                    \"figures\": question_data[\"figures\"],\n",
+    "                },\n",
+    "                \"image_urls\": [subfig[\"url\"] for subfig in subfigures if \"url\" in subfig],\n",
+    "                \"image_captions\": [subfig.get(\"caption\", \"\") for subfig in subfigures],\n",
+    "            },\n",
+    "        }\n",
+    "        logging.info(json.dumps(log_entry))\n",
+    "        print(f\"Error processing case {case_id}, question {question_id}: {str(e)}\")\n",
+    "        return \"\", \"\"\n",
+    "\n",
+    "\n",
+    "def load_benchmark_questions(case_id):\n",
+    "    benchmark_dir = \"../benchmark/questions\"\n",
+    "    return glob.glob(f\"{benchmark_dir}/{case_id}/{case_id}_*.json\")\n",
+    "\n",
+    "\n",
+    "def count_total_questions():\n",
+    "    total_cases = len(glob.glob(\"../benchmark/questions/*\"))\n",
+    "    total_questions = sum(\n",
+    "        len(glob.glob(f\"../benchmark/questions/{case_id}/*.json\"))\n",
+    "        for case_id in os.listdir(\"../benchmark/questions\")\n",
+    "    )\n",
+    "    return total_cases, total_questions\n",
+    "\n",
+    "\n",
+    "def main(tools):\n",
+    "    with open(\"../data/eurorad_metadata.json\", \"r\") as file:\n",
+    "        data = json.load(file)\n",
+    "\n",
+    "    total_cases, total_questions = count_total_questions()\n",
+    "    cases_processed = 0\n",
+    "    questions_processed = 0\n",
+    "    skipped_questions = 0\n",
+    "\n",
+    "    print(f\"Beginning benchmark evaluation for model {model_name} with temperature {temperature}\\n\")\n",
+    "\n",
+    "    for case_id, case_details in data.items():\n",
+    "        if int(case_details[\"case_id\"]) <= 17158:\n",
+    "            continue\n",
+    "\n",
+    "        print(f\"----------------------------------------------------------------\")\n",
+    "        agent, thread = get_agent(tools)\n",
+    "\n",
+    "        question_files = load_benchmark_questions(case_id)\n",
+    "        if not question_files:\n",
+    "            continue\n",
+    "\n",
+    "        cases_processed += 1\n",
+    "        for question_file in question_files:\n",
+    "            with open(question_file, \"r\") as file:\n",
+    "                question_data = json.load(file)\n",
+    "                question_id = os.path.basename(question_file).split(\".\")[0]\n",
+    "\n",
+    "            # agent, thread = get_agent(tools)\n",
+    "            questions_processed += 1\n",
+    "            final_response, model_answer = create_multimodal_request(\n",
+    "                question_data, case_details, case_id, question_id, agent, thread\n",
+    "            )\n",
+    "\n",
+    "            # Handle cases where response is None\n",
+    "            if final_response is None:\n",
+    "                skipped_questions += 1\n",
+    "                print(f\"Skipped question: Case ID {case_id}, Question ID {question_id}\")\n",
+    "                continue\n",
+    "\n",
+    "            print(\n",
+    "                f\"Progress: Case {cases_processed}/{total_cases}, Question {questions_processed}/{total_questions}\"\n",
+    "            )\n",
+    "            print(f\"Case ID: {case_id}\")\n",
+    "            print(f\"Question ID: {question_id}\")\n",
+    "            print(f\"Final Response: {final_response}\")\n",
+    "            print(f\"Model Answer: {model_answer}\")\n",
+    "            print(f\"Correct Answer: {question_data['answer']}\")\n",
+    "            print(f\"----------------------------------------------------------------\\n\")\n",
+    "\n",
+    "    print(f\"\\nBenchmark Summary:\")\n",
+    "    print(f\"Total Cases Processed: {cases_processed}\")\n",
+    "    print(f\"Total Questions Processed: {questions_processed}\")\n",
+    "    print(f\"Total Questions Skipped: {skipped_questions}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tools = get_tools()\n",
+    "main(tools)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "medmax",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.16"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

experiments/chexbench_gpt4.py

@@ -0,0 +1,405 @@
+import json
+import openai
+import os
+from datetime import datetime
+import base64
+import logging
+from pathlib import Path
+import time
+from tqdm import tqdm
+from typing import Dict, List, Optional, Union, Any
+
+# Configuration constants
+DEBUG_MODE = False
+OUTPUT_DIR = "results"
+MODEL_NAME = "gpt-4o-2024-05-13"
+TEMPERATURE = 0.2
+SUBSET = "Visual Question Answering"
+
+# Set up logging configuration
+logging_level = logging.DEBUG if DEBUG_MODE else logging.INFO
+logging.basicConfig(level=logging_level, format="%(asctime)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+
+
+def get_mime_type(file_path: str) -> str:
+    """
+    Determine MIME type based on file extension.
+
+    Args:
+        file_path (str): Path to the file
+
+    Returns:
+        str: MIME type string for the file
+    """
+    extension = os.path.splitext(file_path)[1].lower()
+    mime_types = {
+        ".png": "image/png",
+        ".jpg": "image/jpeg",
+        ".jpeg": "image/jpeg",
+        ".gif": "image/gif",
+    }
+    return mime_types.get(extension, "application/octet-stream")
+
+
+def encode_image(image_path: str) -> str:
+    """
+    Encode image to base64 with extensive error checking.
+
+    Args:
+        image_path (str): Path to the image file
+
+    Returns:
+        str: Base64 encoded image string
+
+    Raises:
+        FileNotFoundError: If image file does not exist
+        ValueError: If image file is empty or too large
+        Exception: For other image processing errors
+    """
+    logger.debug(f"Attempting to read image from: {image_path}")
+    if not os.path.exists(image_path):
+        raise FileNotFoundError(f"Image file not found: {image_path}")
+
+    # Add check for file size
+    file_size = os.path.getsize(image_path)
+    if file_size > 20 * 1024 * 1024:  # 20MB limit
+        raise ValueError("Image file size exceeds 20MB limit")
+    if file_size == 0:
+        raise ValueError("Image file is empty")
+    logger.debug(f"Image file size: {file_size / 1024:.2f} KB")
+
+    try:
+        from PIL import Image
+
+        # Try to open and verify the image
+        with Image.open(image_path) as img:
+            # Get image details
+            width, height = img.size
+            format = img.format
+            mode = img.mode
+            logger.debug(
+                f"Image verification - Format: {format}, Size: {width}x{height}, Mode: {mode}"
+            )
+
+            if format not in ["PNG", "JPEG", "GIF"]:
+                raise ValueError(f"Unsupported image format: {format}")
+
+        with open(image_path, "rb") as image_file:
+            # Read the first few bytes to verify it's a valid PNG
+            header = image_file.read(8)
+            # if header != b'\x89PNG\r\n\x1a\n':
+            #     logger.warning("File does not have a valid PNG signature")
+
+            # Reset file pointer and read entire file
+            image_file.seek(0)
+            encoded = base64.b64encode(image_file.read()).decode("utf-8")
+            encoded_length = len(encoded)
+            logger.debug(f"Base64 encoded length: {encoded_length} characters")
+
+            # Verify the encoded string is not empty and starts correctly
+            if encoded_length == 0:
+                raise ValueError("Base64 encoding produced empty string")
+            if not encoded.startswith("/9j/") and not encoded.startswith("iVBOR"):
+                logger.warning("Base64 string doesn't start with expected JPEG or PNG header")
+
+            return encoded
+    except Exception as e:
+        logger.error(f"Error reading/encoding image: {str(e)}")
+        raise
+
+
+def create_single_request(
+    image_path: str, question: str, options: Dict[str, str]
+) -> List[Dict[str, Any]]:
+    """
+    Create a single API request with image and question.
+
+    Args:
+        image_path (str): Path to the image file
+        question (str): Question text
+        options (Dict[str, str]): Dictionary containing options with keys 'option_0' and 'option_1'
+
+    Returns:
+        List[Dict[str, Any]]: List of message dictionaries for the API request
+
+    Raises:
+        Exception: For errors in request creation
+    """
+    if DEBUG_MODE:
+        logger.debug("Creating API request...")
+
+    prompt = f"""Given the following medical examination question:
+Please answer this multiple choice question:
+
+Question: {question}
+
+Options:
+A) {options['option_0']}
+B) {options['option_1']}
+
+Base your answer only on the provided image and select either A or B."""
+
+    try:
+        encoded_image = encode_image(image_path)
+        mime_type = get_mime_type(image_path)
+
+        if DEBUG_MODE:
+            logger.debug(f"Image encoded with MIME type: {mime_type}")
+
+        messages = [
+            {
+                "role": "system",
+                "content": "You are taking a medical exam. Answer ONLY with the letter (A/B) corresponding to your answer.",
+            },
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": prompt},
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:{mime_type};base64,{encoded_image}"},
+                    },
+                ],
+            },
+        ]
+
+        if DEBUG_MODE:
+            log_messages = json.loads(json.dumps(messages))
+            log_messages[1]["content"][1]["image_url"][
+                "url"
+            ] = f"data:{mime_type};base64,[BASE64_IMAGE_TRUNCATED]"
+            logger.debug(f"Complete API request payload:\n{json.dumps(log_messages, indent=2)}")
+
+        return messages
+
+    except Exception as e:
+        logger.error(f"Error creating request: {str(e)}")
+        raise
+
+
+def check_answer(model_answer: str, correct_answer: int) -> bool:
+    """
+    Check if the model's answer matches the correct answer.
+
+    Args:
+        model_answer (str): The model's answer (A or B)
+        correct_answer (int): The correct answer index (0 for A, 1 for B)
+
+    Returns:
+        bool: True if answer is correct, False otherwise
+    """
+    if not isinstance(model_answer, str):
+        return False
+
+    # Clean the model answer to get just the letter
+    model_letter = model_answer.strip().upper()
+    if model_letter.startswith("A"):
+        model_index = 0
+    elif model_letter.startswith("B"):
+        model_index = 1
+    else:
+        return False
+
+    return model_index == correct_answer
+
+
+def save_results_to_json(results: List[Dict[str, Any]], output_dir: str) -> str:
+    """
+    Save results to a JSON file with timestamp.
+
+    Args:
+        results (List[Dict[str, Any]]): List of result dictionaries
+        output_dir (str): Directory to save results
+
+    Returns:
+        str: Path to the saved file
+    """
+    Path(output_dir).mkdir(parents=True, exist_ok=True)
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    output_file = os.path.join(output_dir, f"batch_results_{timestamp}.json")
+
+    with open(output_file, "w") as f:
+        json.dump(results, f, indent=2)
+
+    logger.info(f"Batch results saved to {output_file}")
+    return output_file
+
+
+def calculate_accuracy(results: List[Dict[str, Any]]) -> tuple[float, int, int]:
+    """
+    Calculate accuracy from results, handling error cases.
+
+    Args:
+        results (List[Dict[str, Any]]): List of result dictionaries
+
+    Returns:
+        tuple[float, int, int]: Tuple containing (accuracy percentage, number correct, total)
+    """
+    if not results:
+        return 0.0, 0, 0
+
+    total = len(results)
+    valid_results = [r for r in results if "output" in r]
+    correct = sum(
+        1 for result in valid_results if result.get("output", {}).get("is_correct", False)
+    )
+
+    accuracy = (correct / total * 100) if total > 0 else 0
+    return accuracy, correct, total
+
+
+def calculate_batch_accuracy(results: List[Dict[str, Any]]) -> float:
+    """
+    Calculate accuracy for the current batch.
+
+    Args:
+        results (List[Dict[str, Any]]): List of result dictionaries
+
+    Returns:
+        float: Accuracy percentage for the batch
+    """
+    valid_results = [r for r in results if "output" in r]
+    if not valid_results:
+        return 0.0
+    return sum(1 for r in valid_results if r["output"]["is_correct"]) / len(valid_results) * 100
+
+
+def process_batch(
+    data: List[Dict[str, Any]], client: openai.OpenAI, start_idx: int = 0, batch_size: int = 50
+) -> List[Dict[str, Any]]:
+    """
+    Process a batch of examples and return results.
+
+    Args:
+        data (List[Dict[str, Any]]): List of data items to process
+        client (openai.OpenAI): OpenAI client instance
+        start_idx (int, optional): Starting index for batch. Defaults to 0
+        batch_size (int, optional): Size of batch to process. Defaults to 50
+
+    Returns:
+        List[Dict[str, Any]]: List of processed results
+    """
+    batch_results = []
+    end_idx = min(start_idx + batch_size, len(data))
+
+    pbar = tqdm(
+        range(start_idx, end_idx),
+        desc=f"Processing batch {start_idx//batch_size + 1}",
+        unit="example",
+    )
+
+    for index in pbar:
+        vqa_item = data[index]
+        options = {"option_0": vqa_item["option_0"], "option_1": vqa_item["option_1"]}
+
+        try:
+            messages = create_single_request(
+                image_path=vqa_item["image_path"], question=vqa_item["question"], options=options
+            )
+
+            response = client.chat.completions.create(
+                model=MODEL_NAME, messages=messages, max_tokens=50, temperature=TEMPERATURE
+            )
+
+            model_answer = response.choices[0].message.content.strip()
+            is_correct = check_answer(model_answer, vqa_item["answer"])
+
+            result = {
+                "timestamp": datetime.now().isoformat(),
+                "example_index": index,
+                "input": {
+                    "question": vqa_item["question"],
+                    "options": {"A": vqa_item["option_0"], "B": vqa_item["option_1"]},
+                    "image_path": vqa_item["image_path"],
+                },
+                "output": {
+                    "model_answer": model_answer,
+                    "correct_answer": "A" if vqa_item["answer"] == 0 else "B",
+                    "is_correct": is_correct,
+                    "usage": {
+                        "prompt_tokens": response.usage.prompt_tokens,
+                        "completion_tokens": response.usage.completion_tokens,
+                        "total_tokens": response.usage.total_tokens,
+                    },
+                },
+            }
+            batch_results.append(result)
+
+            # Update progress bar with current accuracy
+            current_accuracy = calculate_batch_accuracy(batch_results)
+            pbar.set_description(
+                f"Batch {start_idx//batch_size + 1} - Accuracy: {current_accuracy:.2f}% "
+                f"({len(batch_results)}/{index-start_idx+1} examples)"
+            )
+
+        except Exception as e:
+            error_result = {
+                "timestamp": datetime.now().isoformat(),
+                "example_index": index,
+                "error": str(e),
+                "input": {
+                    "question": vqa_item["question"],
+                    "options": {"A": vqa_item["option_0"], "B": vqa_item["option_1"]},
+                    "image_path": vqa_item["image_path"],
+                },
+            }
+            batch_results.append(error_result)
+            if DEBUG_MODE:
+                pbar.write(f"Error processing example {index}: {str(e)}")
+
+        time.sleep(1)  # Rate limiting
+
+    return batch_results
+
+
+def main() -> None:
+    """
+    Main function to process the entire dataset.
+
+    Raises:
+        ValueError: If OPENAI_API_KEY is not set
+        Exception: For other processing errors
+    """
+    logger.info("Starting full dataset processing...")
+    json_path = "../data/chexbench_updated.json"
+
+    try:
+        api_key = os.getenv("OPENAI_API_KEY")
+        if not api_key:
+            raise ValueError("OPENAI_API_KEY environment variable is not set.")
+        client = openai.OpenAI(api_key=api_key)
+
+        with open(json_path, "r") as f:
+            data = json.load(f)
+
+        subset_data = data[SUBSET]
+        total_examples = len(subset_data)
+        logger.info(f"Found {total_examples} examples in {SUBSET} subset")
+
+        all_results = []
+        batch_size = 50  # Process in batches of 50 examples
+
+        # Process all examples in batches
+        for start_idx in range(0, total_examples, batch_size):
+            batch_results = process_batch(subset_data, client, start_idx, batch_size)
+            all_results.extend(batch_results)
+
+            # Save intermediate results after each batch
+            output_file = save_results_to_json(all_results, OUTPUT_DIR)
+
+            # Calculate and log overall progress
+            overall_accuracy, correct, total = calculate_accuracy(all_results)
+            logger.info(f"Overall Progress: {len(all_results)}/{total_examples} examples processed")
+            logger.info(f"Current Accuracy: {overall_accuracy:.2f}% ({correct}/{total} correct)")
+
+        logger.info("Processing completed!")
+        logger.info(f"Final results saved to: {output_file}")
+
+    except Exception as e:
+        logger.error(f"Fatal error: {str(e)}")
+        raise
+
+
+if __name__ == "__main__":
+    main()

experiments/compare_runs.py

@@ -0,0 +1,290 @@
+import json
+import argparse
+import random
+from typing import List, Dict, Any, Tuple
+import re
+from collections import defaultdict
+
+# Define category order
+CATEGORY_ORDER = [
+    "detection",
+    "classification",
+    "localization",
+    "comparison",
+    "relationship",
+    "diagnosis",
+    "characterization",
+]
+
+
+def extract_letter_answer(answer: str) -> str:
+    """Extract just the letter answer from various answer formats.
+
+    Args:
+        answer: The answer string to extract a letter from
+
+    Returns:
+        str: The extracted letter in uppercase, or empty string if no letter found
+    """
+    if not answer:
+        return ""
+
+    # Convert to string and clean
+    answer = str(answer).strip()
+
+    # If it's just a single letter A-F, return it
+    if len(answer) == 1 and answer.upper() in "ABCDEF":
+        return answer.upper()
+
+    # Try to match patterns like "A)", "A.", "A ", etc.
+    match = re.match(r"^([A-F])[).\s]", answer, re.IGNORECASE)
+    if match:
+        return match.group(1).upper()
+
+    # Try to find any standalone A-F letters preceded by space or start of string
+    # and followed by space, period, parenthesis or end of string
+    matches = re.findall(r"(?:^|\s)([A-F])(?:[).\s]|$)", answer, re.IGNORECASE)
+    if matches:
+        return matches[0].upper()
+
+    # Last resort: just find any A-F letter
+    letters = re.findall(r"[A-F]", answer, re.IGNORECASE)
+    if letters:
+        return letters[0].upper()
+
+    # If no letter found, return original (cleaned)
+    return answer.strip().upper()
+
+
+def parse_json_lines(file_path: str) -> Tuple[str, List[Dict[str, Any]]]:
+    """Parse JSON Lines file and extract valid predictions.
+
+    Args:
+        file_path: Path to the JSON Lines file to parse
+
+    Returns:
+        Tuple containing:
+            - str: Model name or file path if model name not found
+            - List[Dict[str, Any]]: List of valid prediction entries
+    """
+    valid_predictions = []
+    model_name = None
+
+    # First try to parse as LLaVA format
+    try:
+        with open(file_path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+            if data.get("model") == "llava-med-v1.5-mistral-7b":
+                model_name = data["model"]
+                for result in data.get("results", []):
+                    if all(k in result for k in ["case_id", "question_id", "correct_answer"]):
+                        # Extract answer with priority: model_answer > validated_answer > raw_output
+                        model_answer = (
+                            result.get("model_answer")
+                            or result.get("validated_answer")
+                            or result.get("raw_output", "")
+                        )
+
+                        # Add default categories for LLaVA results
+                        prediction = {
+                            "case_id": result["case_id"],
+                            "question_id": result["question_id"],
+                            "model_answer": model_answer,
+                            "correct_answer": result["correct_answer"],
+                            "input": {
+                                "question_data": {
+                                    "metadata": {
+                                        "categories": [
+                                            "detection",
+                                            "classification",
+                                            "localization",
+                                            "comparison",
+                                            "relationship",
+                                            "diagnosis",
+                                            "characterization",
+                                        ]
+                                    }
+                                }
+                            },
+                        }
+                        valid_predictions.append(prediction)
+                return model_name, valid_predictions
+    except (json.JSONDecodeError, KeyError):
+        pass
+
+    # If not LLaVA format, process as original format
+    with open(file_path, "r", encoding="utf-8") as f:
+        for line in f:
+            if line.startswith("HTTP Request:"):
+                continue
+            try:
+                data = json.loads(line.strip())
+                if "model" in data:
+                    model_name = data["model"]
+                if all(
+                    k in data for k in ["model_answer", "correct_answer", "case_id", "question_id"]
+                ):
+                    valid_predictions.append(data)
+            except json.JSONDecodeError:
+                continue
+
+    return model_name if model_name else file_path, valid_predictions
+
+
+def filter_common_questions(
+    predictions_list: List[List[Dict[str, Any]]]
+) -> List[List[Dict[str, Any]]]:
+    """Ensure only questions that exist across all models are evaluated.
+
+    Args:
+        predictions_list: List of prediction lists from different models
+
+    Returns:
+        List[List[Dict[str, Any]]]: Filtered predictions containing only common questions
+    """
+    question_sets = [
+        set((p["case_id"], p["question_id"]) for p in preds) for preds in predictions_list
+    ]
+    common_questions = set.intersection(*question_sets)
+
+    return [
+        [p for p in preds if (p["case_id"], p["question_id"]) in common_questions]
+        for preds in predictions_list
+    ]
+
+
+def calculate_accuracy(
+    predictions: List[Dict[str, Any]]
+) -> Tuple[float, int, int, Dict[str, Dict[str, float]]]:
+    """Compute overall and category-level accuracy.
+
+    Args:
+        predictions: List of prediction entries to analyze
+
+    Returns:
+        Tuple containing:
+            - float: Overall accuracy percentage
+            - int: Number of correct predictions
+            - int: Total number of predictions
+            - Dict[str, Dict[str, float]]: Category-level accuracy statistics
+    """
+    if not predictions:
+        return 0.0, 0, 0, {}
+
+    category_performance = defaultdict(lambda: {"total": 0, "correct": 0})
+    correct = 0
+    total = 0
+    sample_size = min(5, len(predictions))
+    sampled_indices = random.sample(range(len(predictions)), sample_size)
+
+    print("\nSample extracted answers:")
+    for i in sampled_indices:
+        pred = predictions[i]
+        model_ans = extract_letter_answer(pred["model_answer"])
+        correct_ans = extract_letter_answer(pred["correct_answer"])
+        print(f"QID: {pred['question_id']}")
+        print(f"  Raw Model Answer: {pred['model_answer']}")
+        print(f"  Extracted Model Answer: {model_ans}")
+        print(f"  Raw Correct Answer: {pred['correct_answer']}")
+        print(f"  Extracted Correct Answer: {correct_ans}")
+        print("-" * 80)
+
+    for pred in predictions:
+        try:
+            model_ans = extract_letter_answer(pred["model_answer"])
+            correct_ans = extract_letter_answer(pred["correct_answer"])
+            categories = (
+                pred.get("input", {})
+                .get("question_data", {})
+                .get("metadata", {})
+                .get("categories", [])
+            )
+
+            if model_ans and correct_ans:
+                total += 1
+                is_correct = model_ans == correct_ans
+                if is_correct:
+                    correct += 1
+
+                for category in categories:
+                    category_performance[category]["total"] += 1
+                    if is_correct:
+                        category_performance[category]["correct"] += 1
+
+        except KeyError:
+            continue
+
+    category_accuracies = {
+        category: {
+            "accuracy": (stats["correct"] / stats["total"]) * 100 if stats["total"] > 0 else 0,
+            "total": stats["total"],
+            "correct": stats["correct"],
+        }
+        for category, stats in category_performance.items()
+    }
+
+    return (correct / total * 100 if total > 0 else 0.0, correct, total, category_accuracies)
+
+
+def compare_models(file_paths: List[str]) -> None:
+    """Compare accuracy between multiple model prediction files.
+
+    Args:
+        file_paths: List of paths to model prediction files to compare
+    """
+    # Parse all files
+    parsed_results = [parse_json_lines(file_path) for file_path in file_paths]
+    model_names, predictions_list = zip(*parsed_results)
+
+    # Get initial stats
+    print(f"\n📊 **Initial Accuracy**:")
+    results = []
+    category_results = []
+
+    for preds, name in zip(predictions_list, model_names):
+        acc, correct, total, category_acc = calculate_accuracy(preds)
+        results.append((acc, correct, total, name))
+        category_results.append(category_acc)
+        print(f"{name}: Accuracy = {acc:.2f}% ({correct}/{total} correct)")
+
+    # Get common questions across all models
+    filtered_predictions = filter_common_questions(predictions_list)
+    print(
+        f"\nQuestions per model after ensuring common questions: {[len(p) for p in filtered_predictions]}"
+    )
+
+    # Compute accuracy on common questions
+    print(f"\n📊 **Accuracy on Common Questions**:")
+    filtered_results = []
+    filtered_category_results = []
+
+    for preds, name in zip(filtered_predictions, model_names):
+        acc, correct, total, category_acc = calculate_accuracy(preds)
+        filtered_results.append((acc, correct, total, name))
+        filtered_category_results.append(category_acc)
+        print(f"{name}: Accuracy = {acc:.2f}% ({correct}/{total} correct)")
+
+    # Print category-wise accuracy
+    print("\nCategory Performance (Common Questions):")
+    for category in CATEGORY_ORDER:
+        print(f"\n{category.capitalize()}:")
+        for model_name, category_acc in zip(model_names, filtered_category_results):
+            stats = category_acc.get(category, {"accuracy": 0, "total": 0, "correct": 0})
+            print(f"  {model_name}: {stats['accuracy']:.2f}% ({stats['correct']}/{stats['total']})")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Compare accuracy across multiple model prediction files"
+    )
+    parser.add_argument("files", nargs="+", help="Paths to model prediction files")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed for sampling")
+
+    args = parser.parse_args()
+    random.seed(args.seed)
+
+    compare_models(args.files)
+
+
+if __name__ == "__main__":
+    main()

experiments/inspect_logs.py

@@ -0,0 +1,210 @@
+from typing import Optional, List
+import argparse
+import json
+import glob
+from pathlib import Path
+from datetime import datetime
+
+
+def get_latest_log() -> str:
+    """Find the most recently modified log file in the current directory.
+
+    Returns:
+        str: Path to the most recently modified log file
+
+    Raises:
+        FileNotFoundError: If no log files are found in the current directory
+    """
+    logs = list(Path(".").glob("api_usage_*.json"))
+    if not logs:
+        raise FileNotFoundError("No log files found in the current directory.")
+    return str(max(logs, key=lambda p: p.stat().st_mtime))
+
+
+def format_cost(entry: dict) -> str:
+    """Format cost if available, otherwise return 'N/A'
+
+    Args:
+        entry: Log entry dictionary containing cost information
+
+    Returns:
+        str: Formatted cost string with $ and 4 decimal places, or 'N/A' if cost not found
+    """
+    return f"${entry.get('cost', 'N/A'):.4f}" if "cost" in entry else "N/A"
+
+
+def print_gpt4_entry(entry: dict) -> None:
+    """Print entry for GPT-4 format
+
+    Args:
+        entry: Log entry dictionary in GPT-4 format containing model info, inputs and outputs
+    """
+    print("\n=== Log Entry ===")
+    print(f"Model: {entry['model']}")
+    print(f"Case ID: {entry['case_id']}")
+    print(f"Question ID: {entry['question_id']}")
+
+    print("\n=== Model Input ===")
+    messages = entry["input"]["messages"]
+    print("System message:", messages[0]["content"])
+    user_content = messages[1]["content"]
+    print("\nUser prompt:", user_content[0]["text"])
+    print("\nImages provided:")
+    for content in user_content[1:]:
+        print(f"  - {content['image_url']['url']}")
+
+    print("\n=== Model Output ===")
+    print(f"Answer: {entry['model_answer']}")
+    print(f"Correct: {entry['correct_answer']}")
+
+    print("\n=== Usage Stats ===")
+    print(f"Duration: {entry['duration']}s")
+    print(f"Cost: {format_cost(entry)}")
+    print(
+        f"Tokens: {entry['usage']['total_tokens']}",
+        f"(prompt: {entry['usage']['prompt_tokens']},",
+        f"completion: {entry['usage']['completion_tokens']})",
+    )
+
+
+def print_llama_entry(entry: dict) -> None:
+    """Print entry for Llama-3.2 format
+
+    Args:
+        entry: Log entry dictionary in Llama format containing model info, inputs and outputs
+    """
+    print("\n=== Log Entry ===")
+    print(f"Model: {entry['model']}")
+    print(f"Case ID: {entry['case_id']}")
+    print(f"Question ID: {entry['question_id']}")
+
+    print("\n=== Model Input ===")
+    print(f"Question: {entry['input']['question_data']['question']}")
+    print("\nImages provided:")
+    for url in entry["input"]["image_urls"]:
+        print(f"  - {url}")
+    if entry["input"]["image_captions"]:
+        print("\nImage captions:")
+        for caption in entry["input"]["image_captions"]:
+            if caption:
+                print(f"  - {caption}")
+
+    print("\n=== Model Output ===")
+    print(f"Answer: {entry['model_answer']}")
+    print(f"Correct: {entry['correct_answer']}")
+
+    print("\n=== Usage Stats ===")
+    print(f"Duration: {entry['duration']}s")
+    if "usage" in entry:
+        print(
+            f"Tokens: {entry['usage']['total_tokens']}",
+            f"(prompt: {entry['usage']['prompt_tokens']},",
+            f"completion: {entry['usage']['completion_tokens']})",
+        )
+
+
+def determine_model_type(entry: dict) -> str:
+    """Determine the model type from the entry
+
+    Args:
+        entry: Log entry dictionary containing model information
+
+    Returns:
+        str: Model type - 'gpt4', 'llama', or 'unknown'
+    """
+    model = entry.get("model", "").lower()
+    if "gpt-4" in model:
+        return "gpt4"
+    elif "llama" in model:
+        return "llama"
+    elif "chexagent" in model:
+        return "chexagent"
+    elif "medrax" in model:
+        return "medrax"
+    else:
+        return "unknown"
+
+
+def print_log_entry(
+    log_file: Optional[str] = None,
+    num_entries: Optional[int] = None,
+    model_filter: Optional[str] = None,
+) -> None:
+    """Print log entries from the specified log file or the latest log file.
+
+    Args:
+        log_file: Path to the log file. If None, uses the latest log file.
+        num_entries: Number of entries to print. If None, prints all entries.
+        model_filter: Filter entries by model type ('gpt4' or 'llama'). If None, prints all.
+    """
+    if log_file is None:
+        log_file = get_latest_log()
+        print(f"Using latest log file: {log_file}")
+
+    entries_printed = 0
+    total_entries = 0
+    filtered_entries = 0
+
+    with open(log_file, "r") as f:
+        for line in f:
+            if line.startswith("HTTP"):
+                continue
+            try:
+                total_entries += 1
+                entry = json.loads(line)
+
+                # Apply model filter if specified
+                model_type = determine_model_type(entry)
+                if model_filter and model_type != model_filter:
+                    filtered_entries += 1
+                    continue
+
+                if model_type == "gpt4":
+                    print_gpt4_entry(entry)
+                elif model_type == "llama":
+                    print_llama_entry(entry)
+                else:
+                    print(f"Unknown model type in entry: {entry['model']}")
+                    continue
+
+                print("=" * 50)
+                entries_printed += 1
+                if num_entries and entries_printed >= num_entries:
+                    break
+
+            except (json.JSONDecodeError, KeyError) as e:
+                print(f"Error processing entry: {e}")
+                continue
+
+    print(f"\nSummary:")
+    print(f"Total entries: {total_entries}")
+    print(f"Entries printed: {entries_printed}")
+    if model_filter:
+        print(f"Entries filtered: {filtered_entries}")
+
+
+def main() -> None:
+    """Main entry point for the script"""
+    parser = argparse.ArgumentParser(
+        description="Parse and display log entries from API usage logs."
+    )
+    parser.add_argument("-l", "--log_file", nargs="?", help="Path to the log file (optional)")
+    parser.add_argument("-n", "--num_entries", type=int, help="Number of entries to display")
+    parser.add_argument(
+        "-m",
+        "--model",
+        choices=["gpt4", "llama"],
+        default="gpt4",
+        help="Model type to display (default: gpt4)",
+    )
+    args = parser.parse_args()
+
+    try:
+        print_log_entry(args.log_file, args.num_entries, args.model)
+    except FileNotFoundError as e:
+        print(f"Error: {e}")
+        exit(1)
+
+
+if __name__ == "__main__":
+    main()

experiments/validate_logs.py

@@ -0,0 +1,162 @@
+from typing import Dict, List, Tuple, Optional
+import json
+import sys
+import glob
+from pathlib import Path
+from collections import defaultdict
+
+
+def get_latest_log() -> str:
+    """Find the most recently modified log file in the current directory.
+
+    Returns:
+        str: Path to the most recently modified log file
+
+    Raises:
+        SystemExit: If no log files are found in current directory
+    """
+    log_pattern = "api_usage_*.json"
+    logs = list(Path(".").glob(log_pattern))
+    if not logs:
+        print(f"No files matching pattern '{log_pattern}' found in current directory")
+        sys.exit(1)
+    return str(max(logs, key=lambda p: p.stat().st_mtime))
+
+
+def analyze_log_file(filename: str) -> Tuple[List[Dict], List[Dict], Dict[str, List[str]]]:
+    """Analyze a log file for entries missing images and errors.
+
+    Args:
+        filename: Path to the log file to analyze
+
+    Returns:
+        Tuple containing:
+            - List of entries with no images
+            - List of skipped/error entries
+            - Dict of processing errors by type
+
+    Raises:
+        SystemExit: If file cannot be found or read
+    """
+    no_images = []
+    errors = defaultdict(list)
+    skipped = []
+
+    try:
+        with open(filename, "r") as f:
+            for line_num, line in enumerate(f, 1):
+                # Skip HTTP request logs
+                if line.startswith("HTTP Request:") or line.strip() == "":
+                    continue
+                try:
+                    # Try to parse the JSON line
+                    if not line.strip().startswith("{"):
+                        continue
+                    entry = json.loads(line.strip())
+                    case_id = entry.get("case_id")
+                    question_id = entry.get("question_id")
+
+                    # Skip if we can't identify the question
+                    if not case_id or not question_id:
+                        continue
+
+                    # Check for explicit skip/error status
+                    if entry.get("status") in ["skipped", "error"]:
+                        skipped.append(
+                            {
+                                "case_id": case_id,
+                                "question_id": question_id,
+                                "reason": entry.get("reason"),
+                                "status": entry.get("status"),
+                            }
+                        )
+                        continue
+
+                    # Check user content for images
+                    messages = entry.get("input", {}).get("messages", [])
+                    has_image = False
+                    for msg in messages:
+                        content = msg.get("content", [])
+                        if isinstance(content, list):
+                            for item in content:
+                                if isinstance(item, dict) and item.get("type") == "image_url":
+                                    has_image = True
+                                    break
+                    if not has_image:
+                        no_images.append(
+                            {
+                                "case_id": case_id,
+                                "question_id": question_id,
+                                "question": entry.get("input", {})
+                                .get("question_data", {})
+                                .get("question", "")[:100]
+                                + "...",  # First 100 chars of question
+                            }
+                        )
+                except json.JSONDecodeError:
+                    errors["json_decode"].append(f"Line {line_num}: Invalid JSON")
+                    continue
+                except Exception as e:
+                    errors["other"].append(f"Line {line_num}: Error processing entry: {str(e)}")
+    except FileNotFoundError:
+        print(f"Error: Could not find log file: {filename}")
+        sys.exit(1)
+    except Exception as e:
+        print(f"Error reading file {filename}: {str(e)}")
+        sys.exit(1)
+
+    return no_images, skipped, errors
+
+
+def print_results(
+    filename: str, no_images: List[Dict], skipped: List[Dict], errors: Dict[str, List[str]]
+) -> None:
+    """Print analysis results.
+
+    Args:
+        filename: Name of the analyzed log file
+        no_images: List of entries with no images
+        skipped: List of skipped/error entries
+        errors: Dict of processing errors by type
+    """
+    print(f"\nAnalyzing log file: {filename}")
+    print("\n=== Questions with No Images ===")
+    if no_images:
+        for entry in no_images:
+            print(f"\nCase ID: {entry['case_id']}")
+            print(f"Question ID: {entry['question_id']}")
+            print(f"Question Preview: {entry['question']}")
+    print(f"\nTotal questions without images: {len(no_images)}")
+
+    print("\n=== Skipped/Error Questions ===")
+    if skipped:
+        for entry in skipped:
+            print(f"\nCase ID: {entry['case_id']}")
+            print(f"Question ID: {entry['question_id']}")
+            print(f"Status: {entry['status']}")
+            print(f"Reason: {entry.get('reason', 'unknown')}")
+    print(f"\nTotal skipped/error questions: {len(skipped)}")
+
+    if errors:
+        print("\n=== Processing Errors ===")
+        for error_type, messages in errors.items():
+            if messages:
+                print(f"\n{error_type}:")
+                for msg in messages:
+                    print(f"  {msg}")
+
+
+def main() -> None:
+    """Main entry point for log validation script."""
+    # If a file is specified as an argument, use it; otherwise find the latest log
+    if len(sys.argv) > 1:
+        log_file = sys.argv[1]
+    else:
+        log_file = get_latest_log()
+
+    no_images, skipped, errors = analyze_log_file(log_file)
+    print_results(log_file, no_images, skipped, errors)
+
+
+if __name__ == "__main__":
+    main()

interface.py

@@ -0,0 +1,279 @@
+import re
+import base64
+import gradio as gr
+from pathlib import Path
+import time
+import shutil
+from typing import AsyncGenerator, List, Optional, Tuple
+from gradio import ChatMessage
+
+
+class ChatInterface:
+    """
+    A chat interface for interacting with a medical AI agent through Gradio.
+
+    Handles file uploads, message processing, and chat history management.
+    Supports both regular image files and DICOM medical imaging files.
+    """
+
+    def __init__(self, agent, tools_dict):
+        """
+        Initialize the chat interface.
+
+        Args:
+            agent: The medical AI agent to handle requests
+            tools_dict (dict): Dictionary of available tools for image processing
+        """
+        self.agent = agent
+        self.tools_dict = tools_dict
+        self.upload_dir = Path("temp")
+        self.upload_dir.mkdir(exist_ok=True)
+        self.current_thread_id = None
+        # Separate storage for original and display paths
+        self.original_file_path = None  # For LLM (.dcm or other)
+        self.display_file_path = None  # For UI (always viewable format)
+
+    def handle_upload(self, file_path: str) -> str:
+        """
+        Handle new file upload and set appropriate paths.
+
+        Args:
+            file_path (str): Path to the uploaded file
+
+        Returns:
+            str: Display path for UI, or None if no file uploaded
+        """
+        if not file_path:
+            return None
+
+        source = Path(file_path)
+        timestamp = int(time.time())
+
+        # Save original file with proper suffix
+        suffix = source.suffix.lower()
+        saved_path = self.upload_dir / f"upload_{timestamp}{suffix}"
+        shutil.copy2(file_path, saved_path)  # Use file_path directly instead of source
+        self.original_file_path = str(saved_path)
+
+        # Handle DICOM conversion for display only
+        if suffix == ".dcm":
+            output, _ = self.tools_dict["DicomProcessorTool"]._run(str(saved_path))
+            self.display_file_path = output["image_path"]
+        else:
+            self.display_file_path = str(saved_path)
+
+        return self.display_file_path
+
+    def add_message(
+        self, message: str, display_image: str, history: List[dict]
+    ) -> Tuple[List[dict], gr.Textbox]:
+        """
+        Add a new message to the chat history.
+
+        Args:
+            message (str): Text message to add
+            display_image (str): Path to image being displayed
+            history (List[dict]): Current chat history
+
+        Returns:
+            Tuple[List[dict], gr.Textbox]: Updated history and textbox component
+        """
+        image_path = self.original_file_path or display_image
+        if image_path is not None:
+            history.append({"role": "user", "content": {"path": image_path}})
+        if message is not None:
+            history.append({"role": "user", "content": message})
+        return history, gr.Textbox(value=message, interactive=False)
+
+    async def process_message(
+        self, message: str, display_image: Optional[str], chat_history: List[ChatMessage]
+    ) -> AsyncGenerator[Tuple[List[ChatMessage], Optional[str], str], None]:
+        """
+        Process a message and generate responses.
+
+        Args:
+            message (str): User message to process
+            display_image (Optional[str]): Path to currently displayed image
+            chat_history (List[ChatMessage]): Current chat history
+
+        Yields:
+            Tuple[List[ChatMessage], Optional[str], str]: Updated chat history, display path, and empty string
+        """
+        chat_history = chat_history or []
+
+        # Initialize thread if needed
+        if not self.current_thread_id:
+            self.current_thread_id = str(time.time())
+
+        messages = []
+        image_path = self.original_file_path or display_image
+
+        if image_path is not None:
+            # Send path for tools
+            messages.append({"role": "user", "content": f"image_path: {image_path}"})
+
+            # Load and encode image for multimodal
+            with open(image_path, "rb") as img_file:
+                img_base64 = base64.b64encode(img_file.read()).decode("utf-8")
+
+            messages.append(
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "image_url",
+                            "image_url": {"url": f"data:image/jpeg;base64,{img_base64}"},
+                        }
+                    ],
+                }
+            )
+
+        if message is not None:
+            messages.append({"role": "user", "content": [{"type": "text", "text": message}]})
+
+        try:
+            for event in self.agent.workflow.stream(
+                {"messages": messages}, {"configurable": {"thread_id": self.current_thread_id}}
+            ):
+                if isinstance(event, dict):
+                    if "process" in event:
+                        content = event["process"]["messages"][-1].content
+                        if content:
+                            content = re.sub(r"temp/[^\s]*", "", content)
+                            chat_history.append(ChatMessage(role="assistant", content=content))
+                            yield chat_history, self.display_file_path, ""
+
+                    elif "execute" in event:
+                        for message in event["execute"]["messages"]:
+                            tool_name = message.name
+                            tool_result = eval(message.content)[0]
+
+                            if tool_result:
+                                metadata = {"title": f"🖼️ Image from tool: {tool_name}"}
+                                formatted_result = " ".join(
+                                    line.strip() for line in str(tool_result).splitlines()
+                                ).strip()
+                                metadata["description"] = formatted_result
+                                chat_history.append(
+                                    ChatMessage(
+                                        role="assistant",
+                                        content=formatted_result,
+                                        metadata=metadata,
+                                    )
+                                )
+
+                            # For image_visualizer, use display path
+                            if tool_name == "image_visualizer":
+                                self.display_file_path = tool_result["image_path"]
+                                chat_history.append(
+                                    ChatMessage(
+                                        role="assistant",
+                                        # content=gr.Image(value=self.display_file_path),
+                                        content={"path": self.display_file_path},
+                                    )
+                                )
+
+                            yield chat_history, self.display_file_path, ""
+
+        except Exception as e:
+            chat_history.append(
+                ChatMessage(
+                    role="assistant", content=f"❌ Error: {str(e)}", metadata={"title": "Error"}
+                )
+            )
+            yield chat_history, self.display_file_path
+
+
+def create_demo(agent, tools_dict):
+    """
+    Create a Gradio demo interface for the medical AI agent.
+
+    Args:
+        agent: The medical AI agent to handle requests
+        tools_dict (dict): Dictionary of available tools for image processing
+
+    Returns:
+        gr.Blocks: Gradio Blocks interface
+    """
+    interface = ChatInterface(agent, tools_dict)
+
+    with gr.Blocks(theme=gr.themes.Soft()) as demo:
+        with gr.Column():
+            gr.Markdown(
+                """
+            # 🏥 MedRAX
+            Medical Reasoning Agent for Chest X-ray
+            """
+            )
+
+            with gr.Row():
+                with gr.Column(scale=3):
+                    chatbot = gr.Chatbot(
+                        [],
+                        height=800,
+                        container=True,
+                        show_label=True,
+                        elem_classes="chat-box",
+                        type="messages",
+                        label="Agent",
+                        avatar_images=(
+                            None,
+                            "assets/medrax_logo.jpg",
+                        ),
+                    )
+                    with gr.Row():
+                        with gr.Column(scale=3):
+                            txt = gr.Textbox(
+                                show_label=False,
+                                placeholder="Ask about the X-ray...",
+                                container=False,
+                            )
+
+                with gr.Column(scale=3):
+                    image_display = gr.Image(
+                        label="Image", type="filepath", height=700, container=True
+                    )
+                    with gr.Row():
+                        upload_button = gr.UploadButton(
+                            "📎 Upload X-Ray",
+                            file_types=["image"],
+                        )
+                        dicom_upload = gr.UploadButton(
+                            "📄 Upload DICOM",
+                            file_types=["file"],
+                        )
+                    with gr.Row():
+                        clear_btn = gr.Button("Clear Chat")
+                        new_thread_btn = gr.Button("New Thread")
+
+        # Event handlers
+        def clear_chat():
+            interface.original_file_path = None
+            interface.display_file_path = None
+            return [], None
+
+        def new_thread():
+            interface.current_thread_id = str(time.time())
+            return [], interface.display_file_path
+
+        def handle_file_upload(file):
+            return interface.handle_upload(file.name)
+
+        chat_msg = txt.submit(
+            interface.add_message, inputs=[txt, image_display, chatbot], outputs=[chatbot, txt]
+        )
+        bot_msg = chat_msg.then(
+            interface.process_message,
+            inputs=[txt, image_display, chatbot],
+            outputs=[chatbot, image_display, txt],
+        )
+        bot_msg.then(lambda: gr.Textbox(interactive=True), None, [txt])
+
+        upload_button.upload(handle_file_upload, inputs=upload_button, outputs=image_display)
+
+        dicom_upload.upload(handle_file_upload, inputs=dicom_upload, outputs=image_display)
+
+        clear_btn.click(clear_chat, outputs=[chatbot, image_display])
+        new_thread_btn.click(new_thread, outputs=[chatbot, image_display])
+
+    return demo

main.py

@@ -0,0 +1,141 @@
+import os
+import warnings
+from typing import *
+from dotenv import load_dotenv
+from transformers import logging
+
+from langgraph.checkpoint.memory import MemorySaver
+from langchain_openai import ChatOpenAI
+from langgraph.checkpoint.memory import MemorySaver
+from langchain_openai import ChatOpenAI
+
+from interface import create_demo
+from medrax.agent import *
+from medrax.tools import *
+from medrax.utils import *
+
+warnings.filterwarnings("ignore")
+logging.set_verbosity_error()
+_ = load_dotenv()
+
+
+def initialize_agent(
+    prompt_file,
+    tools_to_use=None,
+    model_dir="/model-weights",
+    temp_dir="temp",
+    device="cuda",
+    model="google/gemini-2.5-pro-exp-03-25:free",
+    temperature=0.7,
+    top_p=0.95, 
+    openai_kwargs=openai_kwargs
+
+):
+    """Initialize the MedRAX agent with specified tools and configuration.
+
+    Args:
+        prompt_file (str): Path to file containing system prompts
+        tools_to_use (List[str], optional): List of tool names to initialize. If None, all tools are initialized.
+        model_dir (str, optional): Directory containing model weights. Defaults to "/model-weights".
+        temp_dir (str, optional): Directory for temporary files. Defaults to "temp".
+        device (str, optional): Device to run models on. Defaults to "cuda".
+        model (str, optional): Model to use. Defaults to "chatgpt-4o-latest".
+        temperature (float, optional): Temperature for the model. Defaults to 0.7.
+        top_p (float, optional): Top P for the model. Defaults to 0.95.
+        openai_kwargs (dict, optional): Additional keyword arguments for OpenAI API, such as API key and base URL.
+
+    Returns:
+        Tuple[Agent, Dict[str, BaseTool]]: Initialized agent and dictionary of tool instances
+    """
+    prompts = load_prompts_from_file(prompt_file)
+    prompt = prompts["MEDICAL_ASSISTANT"]
+
+    all_tools = {
+        "ChestXRayClassifierTool": lambda: ChestXRayClassifierTool(device=device),
+        "ChestXRaySegmentationTool": lambda: ChestXRaySegmentationTool(device=device),
+        "LlavaMedTool": lambda: LlavaMedTool(cache_dir=model_dir, device=device, load_in_8bit=True),
+        "XRayVQATool": lambda: XRayVQATool(cache_dir=model_dir, device=device),
+        "ChestXRayReportGeneratorTool": lambda: ChestXRayReportGeneratorTool(
+            cache_dir=model_dir, device=device
+        ),
+        "XRayPhraseGroundingTool": lambda: XRayPhraseGroundingTool(
+            cache_dir=model_dir, temp_dir=temp_dir, load_in_8bit=True, device=device
+        ),
+        "ChestXRayGeneratorTool": lambda: ChestXRayGeneratorTool(
+            model_path=f"{model_dir}/roentgen", temp_dir=temp_dir, device=device
+        ),
+        "ImageVisualizerTool": lambda: ImageVisualizerTool(),
+        "DicomProcessorTool": lambda: DicomProcessorTool(temp_dir=temp_dir),
+    }
+
+    # Initialize only selected tools or all if none specified
+    tools_dict = {}
+    tools_to_use = tools_to_use or all_tools.keys()
+    for tool_name in tools_to_use:
+        if tool_name in all_tools:
+            tools_dict[tool_name] = all_tools[tool_name]()
+
+    checkpointer = MemorySaver()
+    model = ChatOpenAI(model=model, temperature=temperature, top_p=top_p, **openai_kwargs)
+    agent = Agent(
+        model,
+        tools=list(tools_dict.values()),
+        log_tools=True,
+        log_dir="logs",
+        system_prompt=prompt,
+        checkpointer=checkpointer,
+    )
+
+    print("Agent initialized")
+    return agent, tools_dict
+
+
+if __name__ == "__main__":
+    """
+    This is the main entry point for the MedRAX application.
+    It initializes the agent with the selected tools and creates the demo.
+    """
+    print("Starting server...")
+
+    # Example: initialize with only specific tools
+    # Here three tools are commented out, you can uncomment them to use them
+    selected_tools = [
+        "ImageVisualizerTool",
+        "DicomProcessorTool",
+        "ChestXRayClassifierTool",
+        "ChestXRaySegmentationTool",
+        "ChestXRayReportGeneratorTool",
+        "XRayVQATool",
+        # "LlavaMedTool",
+        # "XRayPhraseGroundingTool",
+        # "ChestXRayGeneratorTool",
+    ]
+
+    # Collect the ENV variables
+    openai_kwargs = {}
+    if api_key := os.getenv("OPENAI_API_KEY"):
+        openai_kwargs["api_key"] = api_key
+
+    if base_url := os.getenv("OPENAI_BASE_URL"):
+        openai_kwargs["base_url"] = base_url
+    
+    # openai_kwargs = {
+    # "openai_api_key": os.environ.get("OPENAI_API_KEY"),
+    # "openai_api_base": os.environ.get("OPENAI_BASE_URL"),
+    # }
+
+
+    agent, tools_dict = initialize_agent(
+        "medrax/docs/system_prompts.txt",
+        tools_to_use=selected_tools,
+        model_dir="/model-weights",  # Change this to the path of the model weights
+        temp_dir="temp",  # Change this to the path of the temporary directory
+        device="cuda",  # Change this to the device you want to use
+        model="gpt-4o",  # Change this to the model you want to use, e.g. gpt-4o-mini
+        temperature=0.7,
+        top_p=0.95,
+        openai_kwargs=openai_kwargs
+    )
+    demo = create_demo(agent, tools_dict)
+
+    demo.launch(server_name="0.0.0.0", server_port=8585, share=True)

medrax/agent/__init__.py

@@ -0,0 +1 @@
+from .agent import AgentState, Agent

medrax/agent/agent.py

@@ -0,0 +1,193 @@
+import json
+import operator
+from pathlib import Path
+from dotenv import load_dotenv
+from datetime import datetime
+from typing import List, Dict, Any, TypedDict, Annotated, Optional
+
+from langgraph.graph import StateGraph, END
+from langchain_core.messages import AnyMessage, SystemMessage, ToolMessage
+from langchain_core.language_models import BaseLanguageModel
+from langchain_core.tools import BaseTool
+
+_ = load_dotenv()
+
+
+class ToolCallLog(TypedDict):
+    """
+    A TypedDict representing a log entry for a tool call.
+
+    Attributes:
+        timestamp (str): The timestamp of when the tool call was made.
+        tool_call_id (str): The unique identifier for the tool call.
+        name (str): The name of the tool that was called.
+        args (Any): The arguments passed to the tool.
+        content (str): The content or result of the tool call.
+    """
+
+    timestamp: str
+    tool_call_id: str
+    name: str
+    args: Any
+    content: str
+
+
+class AgentState(TypedDict):
+    """
+    A TypedDict representing the state of an agent.
+
+    Attributes:
+        messages (Annotated[List[AnyMessage], operator.add]): A list of messages
+            representing the conversation history. The operator.add annotation
+            indicates that new messages should be appended to this list.
+    """
+
+    messages: Annotated[List[AnyMessage], operator.add]
+
+
+class Agent:
+    """
+    A class representing an agent that processes requests and executes tools based on
+    language model responses.
+
+    Attributes:
+        model (BaseLanguageModel): The language model used for processing.
+        tools (Dict[str, BaseTool]): A dictionary of available tools.
+        checkpointer (Any): Manages and persists the agent's state.
+        system_prompt (str): The system instructions for the agent.
+        workflow (StateGraph): The compiled workflow for the agent's processing.
+        log_tools (bool): Whether to log tool calls.
+        log_path (Path): Path to save tool call logs.
+    """
+
+    def __init__(
+        self,
+        model: BaseLanguageModel,
+        tools: List[BaseTool],
+        checkpointer: Any = None,
+        system_prompt: str = "",
+        log_tools: bool = True,
+        log_dir: Optional[str] = "logs",
+    ):
+        """
+        Initialize the Agent.
+
+        Args:
+            model (BaseLanguageModel): The language model to use.
+            tools (List[BaseTool]): A list of available tools.
+            checkpointer (Any, optional): State persistence manager. Defaults to None.
+            system_prompt (str, optional): System instructions. Defaults to "".
+            log_tools (bool, optional): Whether to log tool calls. Defaults to True.
+            log_dir (str, optional): Directory to save logs. Defaults to 'logs'.
+        """
+        self.system_prompt = system_prompt
+        self.log_tools = log_tools
+
+        if self.log_tools:
+            self.log_path = Path(log_dir or "logs")
+            self.log_path.mkdir(exist_ok=True)
+
+        # Define the agent workflow
+        workflow = StateGraph(AgentState)
+        workflow.add_node("process", self.process_request)
+        workflow.add_node("execute", self.execute_tools)
+        workflow.add_conditional_edges(
+            "process", self.has_tool_calls, {True: "execute", False: END}
+        )
+        workflow.add_edge("execute", "process")
+        workflow.set_entry_point("process")
+
+        self.workflow = workflow.compile(checkpointer=checkpointer)
+        self.tools = {t.name: t for t in tools}
+        self.model = model.bind_tools(tools)
+
+    def process_request(self, state: AgentState) -> Dict[str, List[AnyMessage]]:
+        """
+        Process the request using the language model.
+
+        Args:
+            state (AgentState): The current state of the agent.
+
+        Returns:
+            Dict[str, List[AnyMessage]]: A dictionary containing the model's response.
+        """
+        messages = state["messages"]
+        if self.system_prompt:
+            messages = [SystemMessage(content=self.system_prompt)] + messages
+        response = self.model.invoke(messages)
+        return {"messages": [response]}
+
+    def has_tool_calls(self, state: AgentState) -> bool:
+        """
+        Check if the response contains any tool calls.
+
+        Args:
+            state (AgentState): The current state of the agent.
+
+        Returns:
+            bool: True if tool calls exist, False otherwise.
+        """
+        response = state["messages"][-1]
+        return len(response.tool_calls) > 0
+
+    def execute_tools(self, state: AgentState) -> Dict[str, List[ToolMessage]]:
+        """
+        Execute tool calls from the model's response.
+
+        Args:
+            state (AgentState): The current state of the agent.
+
+        Returns:
+            Dict[str, List[ToolMessage]]: A dictionary containing tool execution results.
+        """
+        tool_calls = state["messages"][-1].tool_calls
+        results = []
+
+        for call in tool_calls:
+            print(f"Executing tool: {call}")
+            if call["name"] not in self.tools:
+                print("\n....invalid tool....")
+                result = "invalid tool, please retry"
+            else:
+                result = self.tools[call["name"]].invoke(call["args"])
+
+            results.append(
+                ToolMessage(
+                    tool_call_id=call["id"],
+                    name=call["name"],
+                    args=call["args"],
+                    content=str(result),
+                )
+            )
+
+        self._save_tool_calls(results)
+        print("Returning to model processing!")
+
+        return {"messages": results}
+
+    def _save_tool_calls(self, tool_calls: List[ToolMessage]) -> None:
+        """
+        Save tool calls to a JSON file with timestamp-based naming.
+
+        Args:
+            tool_calls (List[ToolMessage]): List of tool calls to save.
+        """
+        if not self.log_tools:
+            return
+
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        filename = self.log_path / f"tool_calls_{timestamp}.json"
+
+        logs: List[ToolCallLog] = []
+        for call in tool_calls:
+            log_entry = {
+                "tool_call_id": call.tool_call_id,
+                "name": call.name,
+                "args": call.args,
+                "content": call.content,
+                "timestamp": datetime.now().isoformat(),
+            }
+            logs.append(log_entry)
+
+        with open(filename, "w") as f:
+            json.dump(logs, f, indent=4)

medrax/docs/system_prompts.txt

@@ -0,0 +1,9 @@
+[MEDICAL_ASSISTANT]
+You are an expert medical AI assistant who can answer any medical questions and analyze medical images similar to a doctor.
+Solve using your own vision and reasoning and use tools to complement your reasoning.
+Make multiple tool calls in parallel or sequence as needed for comprehensive answers.
+Critically think 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.
+
+[GENERAL_ASSISTANT]
+You are a helpful AI assistant. Your role is to assist users with a wide range of tasks and questions, providing accurate and useful information on various topics.

medrax/llava/constants.py

@@ -0,0 +1,13 @@
+CONTROLLER_HEART_BEAT_EXPIRATION = 30
+WORKER_HEART_BEAT_INTERVAL = 15
+
+LOGDIR = "."
+
+# Model Constants
+IGNORE_INDEX = -100
+IMAGE_TOKEN_INDEX = -200
+DEFAULT_IMAGE_TOKEN = "<image>"
+DEFAULT_IMAGE_PATCH_TOKEN = "<im_patch>"
+DEFAULT_IM_START_TOKEN = "<im_start>"
+DEFAULT_IM_END_TOKEN = "<im_end>"
+IMAGE_PLACEHOLDER = "<image-placeholder>"

medrax/llava/conversation.py

@@ -0,0 +1,448 @@
+import dataclasses
+from enum import auto, Enum
+from typing import List
+
+
+class SeparatorStyle(Enum):
+    """Different separator style."""
+
+    SINGLE = auto()
+    TWO = auto()
+    MPT = auto()
+    PLAIN = auto()
+    LLAMA_2 = auto()
+    MISTRAL = auto()
+
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that keeps all conversation history."""
+
+    system: str
+    roles: List[str]
+    messages: List[List[str]]
+    offset: int
+    sep_style: SeparatorStyle = SeparatorStyle.SINGLE
+    sep: str = "###"
+    sep2: str = None
+    version: str = "Unknown"
+
+    skip_next: bool = False
+
+    def get_prompt(self):
+        messages = self.messages
+        if len(messages) > 0 and type(messages[0][1]) is tuple:
+            messages = self.messages.copy()
+            init_role, init_msg = messages[0].copy()
+            init_msg = init_msg[0].replace("<image>", "").strip()
+            if "mmtag" in self.version:
+                messages[0] = (init_role, init_msg)
+                messages.insert(0, (self.roles[0], "<Image><image></Image>"))
+                messages.insert(1, (self.roles[1], "Received."))
+            else:
+                messages[0] = (init_role, "<image>\n" + init_msg)
+
+        if self.sep_style == SeparatorStyle.SINGLE:
+            ret = self.system + self.sep
+            for role, message in messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + ": " + message + self.sep
+                else:
+                    ret += role + ":"
+        elif self.sep_style == SeparatorStyle.TWO:
+            seps = [self.sep, self.sep2]
+            ret = self.system + seps[0]
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    sep = seps[i % 2]
+                    sep = "{0} ".format(self.sep2) if sep == self.sep2 else self.sep
+                    ret += role + ": " + message.strip() + sep
+                else:
+                    ret += role + ":"
+            ret = ret.strip()
+        elif self.sep_style == SeparatorStyle.MPT:
+            ret = self.system + self.sep
+            for role, message in messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+        elif self.sep_style == SeparatorStyle.LLAMA_2:
+            wrap_sys = lambda msg: f"<<SYS>>\n{msg}\n<</SYS>>\n\n" if len(msg) > 0 else msg
+            wrap_inst = lambda msg: f"[INST] {msg} [/INST]"
+            ret = ""
+
+            for i, (role, message) in enumerate(messages):
+                if i == 0:
+                    assert message, "first message should not be none"
+                    assert role == self.roles[0], "first message should come from user"
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    if i == 0:
+                        message = wrap_sys(self.system) + message
+                    if i % 2 == 0:
+                        message = wrap_inst(message)
+                        ret += self.sep + message
+                    else:
+                        ret += " " + message + " " + self.sep2
+                else:
+                    ret += ""
+            ret = ret.lstrip(self.sep)
+        elif self.sep_style == SeparatorStyle.PLAIN:
+            seps = [self.sep, self.sep2]
+            ret = self.system
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += message + seps[i % 2]
+                else:
+                    ret += ""
+        elif self.sep_style == SeparatorStyle.MISTRAL:
+            # reference: https://docs.mistral.ai/models/
+            wrap_sys = lambda msg: f"{msg}</s>"
+            wrap_inst = lambda msg: f"[INST] {msg} [/INST]"
+            ret = ""
+            for i, (role, message) in enumerate(messages):
+                if i == 0:
+                    assert message, "first message should not be none"
+                    assert role == self.roles[0], "first message should come from user"
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    if i == 0:
+                        message = self.system + " " + message.strip()
+                    if i % 2 == 0:
+                        message = wrap_inst(message)
+                        ret += message
+                    else:
+                        ret += wrap_sys(message)
+                else:
+                    ret += ""
+            # wrap_sys = lambda msg: f"\n{msg}\n\n"
+            # wrap_inst = lambda msg: f"[INST] {msg} [/INST]"
+            # ret = ""
+            # for i, (role, message) in enumerate(messages):
+            #     if i == 0:
+            #         assert message, "first message should not be none"
+            #         assert role == self.roles[0], "first message should come from user"
+            #     if message:
+            #         if type(message) is tuple:
+            #             message, _, _ = message
+            #         if i == 0: message = wrap_sys(self.system) + message
+            #         if i % 2 == 0:
+            #             message = wrap_inst(message)
+            #             ret += message if i != 0 else self.sep + message
+            #         else:
+            #             # NOTE-JW: we need to add " " to strictly follow Mistral Instruction Format
+            #             ret += " " + message + " " + self.sep2
+            #             # ret += " " + wrap_sys(message)
+            #     else:
+            #         ret += ""
+        else:
+            raise ValueError(f"Invalid style: {self.sep_style}")
+
+        return ret
+
+    def append_message(self, role, message):
+        self.messages.append([role, message])
+
+    def get_images(self, return_pil=False):
+        images = []
+        for i, (role, msg) in enumerate(self.messages[self.offset :]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    import base64
+                    from io import BytesIO
+                    from PIL import Image
+
+                    msg, image, image_process_mode = msg
+                    if image_process_mode == "Pad":
+
+                        def expand2square(pil_img, background_color=(122, 116, 104)):
+                            width, height = pil_img.size
+                            if width == height:
+                                return pil_img
+                            elif width > height:
+                                result = Image.new(pil_img.mode, (width, width), background_color)
+                                result.paste(pil_img, (0, (width - height) // 2))
+                                return result
+                            else:
+                                result = Image.new(pil_img.mode, (height, height), background_color)
+                                result.paste(pil_img, ((height - width) // 2, 0))
+                                return result
+
+                        image = expand2square(image)
+                    elif image_process_mode in ["Default", "Crop"]:
+                        pass
+                    elif image_process_mode == "Resize":
+                        image = image.resize((336, 336))
+                    else:
+                        raise ValueError(f"Invalid image_process_mode: {image_process_mode}")
+                    max_hw, min_hw = max(image.size), min(image.size)
+                    aspect_ratio = max_hw / min_hw
+                    max_len, min_len = 800, 400
+                    shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))
+                    longest_edge = int(shortest_edge * aspect_ratio)
+                    W, H = image.size
+                    if longest_edge != max(image.size):
+                        if H > W:
+                            H, W = longest_edge, shortest_edge
+                        else:
+                            H, W = shortest_edge, longest_edge
+                        image = image.resize((W, H))
+                    if return_pil:
+                        images.append(image)
+                    else:
+                        buffered = BytesIO()
+                        image.save(buffered, format="PNG")
+                        img_b64_str = base64.b64encode(buffered.getvalue()).decode()
+                        images.append(img_b64_str)
+        return images
+
+    def to_gradio_chatbot(self):
+        ret = []
+        for i, (role, msg) in enumerate(self.messages[self.offset :]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    import base64
+                    from io import BytesIO
+
+                    msg, image, image_process_mode = msg
+                    max_hw, min_hw = max(image.size), min(image.size)
+                    aspect_ratio = max_hw / min_hw
+                    max_len, min_len = 800, 400
+                    shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))
+                    longest_edge = int(shortest_edge * aspect_ratio)
+                    W, H = image.size
+                    if H > W:
+                        H, W = longest_edge, shortest_edge
+                    else:
+                        H, W = shortest_edge, longest_edge
+                    image = image.resize((W, H))
+                    buffered = BytesIO()
+                    image.save(buffered, format="JPEG")
+                    img_b64_str = base64.b64encode(buffered.getvalue()).decode()
+                    img_str = (
+                        f'<img src="data:image/png;base64,{img_b64_str}" alt="user upload image" />'
+                    )
+                    msg = img_str + msg.replace("<image>", "").strip()
+                    ret.append([msg, None])
+                else:
+                    ret.append([msg, None])
+            else:
+                ret[-1][-1] = msg
+        return ret
+
+    def copy(self):
+        return Conversation(
+            system=self.system,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            offset=self.offset,
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2,
+            version=self.version,
+        )
+
+    def dict(self):
+        if len(self.get_images()) > 0:
+            return {
+                "system": self.system,
+                "roles": self.roles,
+                "messages": [[x, y[0] if type(y) is tuple else y] for x, y in self.messages],
+                "offset": self.offset,
+                "sep": self.sep,
+                "sep2": self.sep2,
+            }
+        return {
+            "system": self.system,
+            "roles": self.roles,
+            "messages": self.messages,
+            "offset": self.offset,
+            "sep": self.sep,
+            "sep2": self.sep2,
+        }
+
+
+conv_vicuna_v0 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("Human", "Assistant"),
+    messages=(
+        (
+            "Human",
+            "What are the key differences between renewable and non-renewable energy sources?",
+        ),
+        (
+            "Assistant",
+            "Renewable energy sources are those that can be replenished naturally in a relatively "
+            "short amount of time, such as solar, wind, hydro, geothermal, and biomass. "
+            "Non-renewable energy sources, on the other hand, are finite and will eventually be "
+            "depleted, such as coal, oil, and natural gas. Here are some key differences between "
+            "renewable and non-renewable energy sources:\n"
+            "1. Availability: Renewable energy sources are virtually inexhaustible, while non-renewable "
+            "energy sources are finite and will eventually run out.\n"
+            "2. Environmental impact: Renewable energy sources have a much lower environmental impact "
+            "than non-renewable sources, which can lead to air and water pollution, greenhouse gas emissions, "
+            "and other negative effects.\n"
+            "3. Cost: Renewable energy sources can be more expensive to initially set up, but they typically "
+            "have lower operational costs than non-renewable sources.\n"
+            "4. Reliability: Renewable energy sources are often more reliable and can be used in more remote "
+            "locations than non-renewable sources.\n"
+            "5. Flexibility: Renewable energy sources are often more flexible and can be adapted to different "
+            "situations and needs, while non-renewable sources are more rigid and inflexible.\n"
+            "6. Sustainability: Renewable energy sources are more sustainable over the long term, while "
+            "non-renewable sources are not, and their depletion can lead to economic and social instability.\n",
+        ),
+    ),
+    offset=2,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+)
+
+conv_vicuna_v1 = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="v1",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+conv_llama_2 = Conversation(
+    system="""You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
+
+If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.""",
+    roles=("USER", "ASSISTANT"),
+    version="llama_v2",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.LLAMA_2,
+    sep="<s>",
+    sep2="</s>",
+)
+
+conv_llava_llama_2 = Conversation(
+    system="You are a helpful language and vision assistant. "
+    "You are able to understand the visual content that the user provides, "
+    "and assist the user with a variety of tasks using natural language.",
+    roles=("USER", "ASSISTANT"),
+    version="llama_v2",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.LLAMA_2,
+    sep="<s>",
+    sep2="</s>",
+)
+
+conv_mpt = Conversation(
+    system="""<|im_start|>system
+A conversation between a user and an LLM-based AI assistant. The assistant gives helpful and honest answers.""",
+    roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),
+    version="mpt",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.MPT,
+    sep="<|im_end|>",
+)
+
+conv_llava_plain = Conversation(
+    system="",
+    roles=("", ""),
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.PLAIN,
+    sep="\n",
+)
+
+conv_llava_v0 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("Human", "Assistant"),
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+)
+
+conv_llava_v0_mmtag = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+    "The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language."
+    "The visual content will be provided with the following format: <Image>visual content</Image>.",
+    roles=("Human", "Assistant"),
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+    version="v0_mmtag",
+)
+
+conv_llava_v1 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="v1",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+conv_llava_v1_mmtag = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+    "The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language."
+    "The visual content will be provided with the following format: <Image>visual content</Image>.",
+    roles=("USER", "ASSISTANT"),
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+    version="v1_mmtag",
+)
+
+conv_mistral_instruct = Conversation(
+    system="",
+    roles=("USER", "ASSISTANT"),
+    version="llama_v2",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.LLAMA_2,
+    sep="",
+    sep2="</s>",
+)
+
+default_conversation = conv_vicuna_v1
+conv_templates = {
+    "default": conv_vicuna_v0,
+    "v0": conv_vicuna_v0,
+    "v1": conv_vicuna_v1,
+    "vicuna_v1": conv_vicuna_v1,
+    "llama_2": conv_llama_2,
+    "mistral_instruct": conv_mistral_instruct,
+    "plain": conv_llava_plain,
+    "v0_plain": conv_llava_plain,
+    "llava_v0": conv_llava_v0,
+    "v0_mmtag": conv_llava_v0_mmtag,
+    "llava_v1": conv_llava_v1,
+    "v1_mmtag": conv_llava_v1_mmtag,
+    "llava_llama_2": conv_llava_llama_2,
+    "mpt": conv_mpt,
+}
+
+
+if __name__ == "__main__":
+    print(default_conversation.get_prompt())

medrax/llava/eval/eval_multimodal_chat_gpt_score.py

@@ -0,0 +1,143 @@
+import os
+import json
+import argparse
+from copy import deepcopy
+from pathlib import Path
+from tqdm import tqdm
+
+import llm
+import util
+
+
+INSTRUCT_PROMPT = """We would like to request your feedback on the performance of two AI assistants in response to the user question displayed above. The user asks the question on observing an image. For your reference, the visual content in the image is represented with caption describing the same image.
+  Please rate the helpfulness, relevance, accuracy, level of details of their responses. Each assistant receives an overall score on a scale of 1 to 10, where a higher score indicates better overall performance.
+  Please first output a single line containing only two values indicating the scores for Assistant 1 and 2, respectively. The two scores are separated by a space. In the subsequent line, please provide a comprehensive explanation of your evaluation, avoiding any potential bias and ensuring that the order in which the responses were presented does not affect your judgment."""
+ROLE = "Assistant"
+
+# Generate instruction for GPT-4 to score the two answers.
+def conv_to_str(fig_label, fig_caption, fig_context, question, ans1, ans2):
+    return (
+        f"[Context]\n"
+        f"Figure Caption:\n{fig_label}: {fig_caption}\n\n"
+        f"Figure Context:\n\t- {fig_context}\n\n"
+        f"[Question]\n{question}\n\n"
+        f"[{ROLE} 1]\n{ans1}\n\n[End of {ROLE} 1]\n\n"
+        f"[{ROLE} 2]\n{ans2}\n\n[End of {ROLE} 2]\n\n"
+        f"[System]\n{INSTRUCT_PROMPT}\n\n"
+    )
+
+
+def compare_messages_gen(fig_label, fig_caption, fig_context, question, ans1, ans2):
+    messages = [
+        {
+            "role": "system",
+            "content": """'You are a helpful and precise assistant for checking the quality of the answer.""",
+        },
+    ]
+    messages.append(
+        {
+            "role": "user",
+            "content": conv_to_str(fig_label, fig_caption, fig_context, question, ans1, ans2),
+        }
+    )
+    return messages
+
+
+def sum_list_list(x):
+    return sum(item for inner_list in x for item in inner_list)
+
+
+def chunk(lst, n):
+    for i in range(0, len(lst), n):
+        if i + (1.5 * n) < len(lst):
+            end = i + n
+        else:
+            end = len(lst)
+        yield lst[i:end]
+        if end == len(lst):
+            return
+
+
+def infer(samples):
+    model_inst = llm.GPT("gpt-4-0314")
+
+    BATCH_SIZE = 1
+    batch_samples = []
+    results = []
+    batch = []
+
+    print("Starting Multimodal Chat GPT Scoring Eval")
+
+    for sample in tqdm(samples):
+        sample_copy = deepcopy(sample)
+        input_msg = compare_messages_gen(
+            sample_copy["fig_label"],
+            sample_copy["fig_caption"],
+            sample_copy["in_text_mention"],
+            sample_copy["question"],
+            sample_copy["ans1"],
+            sample_copy["ans2"],
+        )
+        batch.append(input_msg)
+        batch_samples.append(sample_copy)
+        if len(batch) >= BATCH_SIZE:
+            inference_results = [
+                x.strip()
+                for chunk_messages in chunk([x for x in batch if x], BATCH_SIZE)
+                for x in model_inst.infer(chunk_messages)
+            ]
+            for item, inference_result in zip(batch_samples, inference_results):
+                item["gpt_eval"] = inference_result
+            results.extend(batch_samples)
+            batch = []
+            batch_samples = []
+    inference_results = [
+        x.strip()
+        for chunk_messages in chunk([x for x in batch if x], BATCH_SIZE)
+        for x in model_inst.infer(chunk_messages)
+    ]
+    for item, inference_result in zip(batch_samples, inference_results):
+        item["gpt_eval"] = inference_result
+    results.extend(batch_samples)
+    print(f"Result Size: {len(results)}")
+    return results
+
+
+def main(args):
+    answer_data = util.load_file_jsonl(args.answers_file)
+    question_data = util.load_file_jsonl(args.question_file)
+
+    samples = []
+    for question, answer in zip(question_data, answer_data):
+        question_copy = deepcopy(question)
+        question["question"] = question_copy["text"]
+        question["ans1"] = question_copy.pop("gpt4_answer")
+        question["ans2"] = answer["text"]
+        samples.append(question)
+
+    results = infer(samples)
+
+    # Create parent directory of output score files if it doesn't exist
+    os.makedirs(Path(args.scores_file).parent, exist_ok=True)
+
+    with open(args.scores_file, "w") as f:
+        for row in results:
+            f.write(json.dumps(row) + "\n")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser("GPT-4 Multimodal Chat Scoring", add_help=True)
+    parser.add_argument(
+        "--answers-file", default="", metavar="FILE", help="path to model answer file"
+    )
+    parser.add_argument(
+        "--question-file",
+        default="data/questions/llava_med_eval_qa50_qa.jsonl",
+        metavar="FILE",
+        help="path to multichat questions file",
+    )
+    parser.add_argument(
+        "--scores-file", default="", metavar="FILE", help="path to save gpt-4 score file"
+    )
+    args = parser.parse_args()
+    main(args)

medrax/llava/eval/llm.py

@@ -0,0 +1,154 @@
+import abc
+import asyncio
+from abc import abstractmethod
+import math
+
+import tiktoken
+import openai
+import backoff
+
+
+class LLM(abc.ABC):
+
+    prompt_percent = 0.9
+
+    @abstractmethod
+    def __init__(self):
+        raise NotImplementedError("Subclasses should implement this!")
+
+    @abstractmethod
+    def infer(self, prompts):
+        raise NotImplementedError("Subclasses should implement this!")
+
+    @abstractmethod
+    def split_input(
+        self, fixed_instruction, few_shot_examples, splittable_input, input_header, output_header
+    ):
+        raise NotImplementedError("Subclasses should implement this!")
+
+
+class GPT(LLM):
+
+    prompt_percent = 0.8
+
+    openai_cxn_dict = {
+        "default": {
+            "endpoint": "INSERT YOUR AZURE OPENAI ENDPOINT HERE",
+            "api_key": "INSERT YOUR AZURE OPENAI API KEY HERE",
+        },
+    }
+
+    deployment_max_length_dict = {
+        "gpt-4": 8192,
+        "gpt-4-0314": 8192,
+        "gpt-4-32k": 32768,
+        "gpt-35-turbo": 4096,
+        "gpt-35-turbo-16k": 16385,
+    }
+
+    def __init__(self, model_id):
+        self.temperature = 0.0
+        self.top_k = 1
+        self.encoding = tiktoken.encoding_for_model(
+            "-".join(model_id.split("-", 2)[:2]).replace("5", ".5")
+        )
+        self.openai_api = "default"
+        self.model_id = model_id
+        self.max_length = self.deployment_max_length_dict[model_id]
+        self.client = openai.AsyncAzureOpenAI(
+            api_key=self.openai_cxn_dict[self.openai_api]["api_key"],
+            api_version="2023-12-01-preview",
+            azure_endpoint=self.openai_cxn_dict[self.openai_api]["endpoint"],
+        )
+
+    def gen_messages(
+        self, fixed_instruction, few_shot_examples, input, input_header, output_header
+    ):
+        messages = [
+            {
+                "role": "system",
+                "content": fixed_instruction,
+            },
+        ]
+        for example in few_shot_examples:
+            messages.extend(
+                [
+                    {
+                        "role": "user",
+                        "content": input_header + "\n" + example["user"] + "\n\n" + output_header,
+                    },
+                    {
+                        "role": "assistant",
+                        "content": example["assistant"],
+                    },
+                ]
+            )
+        messages.extend(
+            [
+                {
+                    "role": "user",
+                    "content": input_header + "\n" + input + "\n\n" + output_header,
+                },
+            ]
+        )
+        return messages
+
+    # Define the coroutine for making API calls to GPT
+    @backoff.on_exception(backoff.expo, openai.RateLimitError)
+    async def make_api_call_to_gpt(self, messages):
+        response = await self.client.chat.completions.create(
+            model=self.model_id,
+            messages=messages,
+            temperature=self.temperature,
+        )
+        return response.choices[0].message.content
+
+    async def dispatch_openai_requests(
+        self,
+        messages_list,
+    ):
+        # Asynchronously call the function for each prompt
+        tasks = [self.make_api_call_to_gpt(messages) for messages in messages_list]
+
+        # Gather and run the tasks concurrently
+        results = await asyncio.gather(*tasks)
+        return results
+
+    def infer(
+        self,
+        messages_list,
+    ):
+        return asyncio.run(self.dispatch_openai_requests(messages_list))
+
+    def split_input(
+        self, fixed_instruction, few_shot_examples, splittable_input, input_header, output_header
+    ):
+        # Tokenize fixed_prompt
+        fixed_token_ids = self.encoding.encode(
+            fixed_instruction
+            + " ".join([x["user"] + " " + x["assistant"] for x in few_shot_examples])
+        )
+        # Calculate remaining token length
+        remaining_token_len = math.ceil(
+            (self.prompt_percent * self.max_length) - len(fixed_token_ids)
+        )
+
+        # Tokenize splittable_input
+        split_token_ids = self.encoding.encode(splittable_input)
+
+        # Split tokenized split_prompt into list of individual inputs strings. Uses tokens to calculate length
+        split_token_ids_list = [
+            split_token_ids[i : i + remaining_token_len + 10]
+            for i in range(0, len(split_token_ids), remaining_token_len)
+        ]
+        split_input_list = [
+            self.encoding.decode(split_token_ids) for split_token_ids in split_token_ids_list
+        ]
+
+        # Take the fixed_prompt, few_shot_examples, splitted inputs, and input/output headers and generate list of prompt strings.
+        return [
+            self.gen_messages(
+                fixed_instruction, few_shot_examples, split_input, input_header, output_header
+            )
+            for split_input in split_input_list
+        ]

medrax/llava/eval/model_vqa.py

@@ -0,0 +1,133 @@
+import argparse
+import torch
+import os
+import json
+from tqdm import tqdm
+import shortuuid
+
+from medrax.llava.constants import (
+    IMAGE_TOKEN_INDEX,
+    DEFAULT_IMAGE_TOKEN,
+    DEFAULT_IM_START_TOKEN,
+    DEFAULT_IM_END_TOKEN,
+)
+from medrax.llava.conversation import conv_templates, SeparatorStyle
+from medrax.llava.model.builder import load_pretrained_model
+from medrax.llava.utils import disable_torch_init
+from medrax.llava.mm_utils import (
+    tokenizer_image_token,
+    get_model_name_from_path,
+    KeywordsStoppingCriteria,
+    process_images,
+)
+
+from PIL import Image
+import math
+from transformers import set_seed, logging
+
+logging.set_verbosity_error()
+
+
+def split_list(lst, n):
+    """Split a list into n (roughly) equal-sized chunks"""
+    chunk_size = math.ceil(len(lst) / n)  # integer division
+    return [lst[i : i + chunk_size] for i in range(0, len(lst), chunk_size)]
+
+
+def get_chunk(lst, n, k):
+    chunks = split_list(lst, n)
+    return chunks[k]
+
+
+def eval_model(args):
+    set_seed(0)
+    # Model
+    disable_torch_init()
+    model_path = os.path.expanduser(args.model_path)
+    model_name = get_model_name_from_path(model_path)
+    tokenizer, model, image_processor, context_len = load_pretrained_model(
+        model_path, args.model_base, model_name
+    )
+
+    questions = [json.loads(q) for q in open(os.path.expanduser(args.question_file), "r")]
+    questions = get_chunk(questions, args.num_chunks, args.chunk_idx)
+    answers_file = os.path.expanduser(args.answers_file)
+    os.makedirs(os.path.dirname(answers_file), exist_ok=True)
+    ans_file = open(answers_file, "w")
+    for line in tqdm(questions):
+        idx = line["question_id"]
+        image_file = line["image"]
+        qs = line["text"].replace(DEFAULT_IMAGE_TOKEN, "").strip()
+        cur_prompt = qs
+        if model.config.mm_use_im_start_end:
+            qs = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + "\n" + qs
+        else:
+            qs = DEFAULT_IMAGE_TOKEN + "\n" + qs
+
+        conv = conv_templates[args.conv_mode].copy()
+        conv.append_message(conv.roles[0], qs)
+        conv.append_message(conv.roles[1], None)
+        prompt = conv.get_prompt()
+
+        input_ids = (
+            tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt")
+            .unsqueeze(0)
+            .cuda()
+        )
+
+        image = Image.open(os.path.join(args.image_folder, image_file))
+        image_tensor = process_images([image], image_processor, model.config)[0]
+
+        stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
+        keywords = [stop_str]
+        stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
+
+        with torch.inference_mode():
+            output_ids = model.generate(
+                input_ids,
+                images=image_tensor.unsqueeze(0).half().cuda(),
+                do_sample=True if args.temperature > 0 else False,
+                temperature=args.temperature,
+                top_p=args.top_p,
+                num_beams=args.num_beams,
+                # no_repeat_ngram_size=3,
+                max_new_tokens=1024,
+                use_cache=True,
+            )
+
+        outputs = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0].strip()
+
+        ans_id = shortuuid.uuid()
+        ans_file.write(
+            json.dumps(
+                {
+                    "question_id": idx,
+                    "prompt": cur_prompt,
+                    "text": outputs,
+                    "answer_id": ans_id,
+                    "model_id": model_name,
+                    "metadata": {},
+                }
+            )
+            + "\n"
+        )
+        ans_file.flush()
+    ans_file.close()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", type=str, default="facebook/opt-350m")
+    parser.add_argument("--model-base", type=str, default=None)
+    parser.add_argument("--image-folder", type=str, default="")
+    parser.add_argument("--question-file", type=str, default="tables/question.jsonl")
+    parser.add_argument("--answers-file", type=str, default="answer.jsonl")
+    parser.add_argument("--conv-mode", type=str, default="vicuna_v1")
+    parser.add_argument("--num-chunks", type=int, default=1)
+    parser.add_argument("--chunk-idx", type=int, default=0)
+    parser.add_argument("--temperature", type=float, default=0.2)
+    parser.add_argument("--top_p", type=float, default=None)
+    parser.add_argument("--num_beams", type=int, default=1)
+    args = parser.parse_args()
+
+    eval_model(args)

medrax/llava/eval/summarize_gpt_review.py

@@ -0,0 +1,62 @@
+import argparse
+import util
+from collections import defaultdict
+import pandas as pd
+
+
+def get_domain(x):
+    for domain in ["chest_xray", "mri", "histology", "gross", "ct_scan"]:
+        in_domain = x["domain"][domain]
+        if in_domain:
+            return domain
+
+
+def main(args):
+    scores_data = util.load_file_jsonl(args.scores_file)
+    predictions = [
+        (x["question_id"], x["type"], get_domain(x), x["gpt_eval"].split("\n")[0].split(" "))
+        for x in scores_data
+    ]
+
+    score_type_dict = defaultdict(lambda: defaultdict(list))
+    for q_id, q_type, domain, (a1_score, a2_score) in predictions:
+        score_type_dict[q_type][1].append(a1_score)
+        score_type_dict[q_type][2].append(a2_score)
+        score_type_dict["overall"][1].append(a1_score)
+        score_type_dict["overall"][2].append(a2_score)
+        score_type_dict[domain][1].append(a1_score)
+        score_type_dict[domain][2].append(a2_score)
+
+    result = defaultdict(dict)
+
+    for q_type, score_dict in score_type_dict.items():
+        result[q_type]["gpt4_score"] = util.get_avg(score_dict[1])
+        result[q_type]["pred_score"] = util.get_avg(score_dict[2])
+        result[q_type]["pred_relative_score"] = (
+            util.get_avg([float(s2) / float(s1) for s1, s2 in zip(score_dict[1], score_dict[2])])
+            * 100
+        )
+        result[q_type]["data_size"] = len(score_dict[1])
+
+    df = pd.DataFrame.from_dict(result).filter(
+        [
+            "conversation",
+            "detailed_description",
+            "chest_xray",
+            "mri",
+            "histology",
+            "gross",
+            "ct_scan",
+            "overall",
+        ]
+    )
+    print(df)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser("GPT-4 Multimodal Chat Eval Postprocessing", add_help=True)
+    parser.add_argument(
+        "--scores-file", default="", metavar="FILE", help="input path to gpt-4 score file"
+    )
+    args = parser.parse_args()
+    main(args)

medrax/llava/eval/util.py

@@ -0,0 +1,10 @@
+import json
+
+
+def load_file_jsonl(path):
+    with open(path) as f:
+        return [json.loads(row) for row in f]
+
+
+def get_avg(x):
+    return sum([float(y) for y in x]) / len(x)

medrax/llava/mm_utils.py

@@ -0,0 +1,121 @@
+from PIL import Image
+from io import BytesIO
+import base64
+import random
+import torch
+from transformers import StoppingCriteria
+from medrax.llava.constants import IMAGE_TOKEN_INDEX
+
+
+def load_image_from_base64(image):
+    return Image.open(BytesIO(base64.b64decode(image)))
+
+
+def expand2square(pil_img, background_color):
+    width, height = pil_img.size
+    if width == height:
+        return pil_img
+    elif width > height:
+        result = Image.new(pil_img.mode, (width, width), background_color)
+        # sample a random between 0 and (width - height) // 2
+        y_start = random.randint((width - height) // 2, (width - height) // 2 + 1)
+        result.paste(pil_img, (0, y_start))
+        return result
+    else:
+        result = Image.new(pil_img.mode, (height, height), background_color)
+        # sample a random between 0 and (height - width) // 2
+        x_start = random.randint((height - width) // 2, (height - width) // 2 + 1)
+        result.paste(pil_img, (x_start, 0))
+        return result
+
+
+def process_images(images, image_processor, model_cfg):
+    image_aspect_ratio = getattr(model_cfg, "image_aspect_ratio", None)
+    new_images = []
+    for image in images:
+        if image_aspect_ratio == "pad":
+            if image.mode == "L":
+                background_color = int(
+                    255 * sum(image_processor.image_mean) / len(image_processor.image_mean)
+                )
+            else:
+                background_color = tuple(int(x * 255) for x in image_processor.image_mean)
+            image = expand2square(image, background_color)
+        image = image_processor.preprocess(image, return_tensors="pt")["pixel_values"][0]
+        new_images.append(image)
+    if all(x.shape == new_images[0].shape for x in new_images):
+        new_images = torch.stack(new_images, dim=0)
+    return new_images
+
+
+def tokenizer_image_token(
+    prompt, tokenizer, image_token_index=IMAGE_TOKEN_INDEX, return_tensors=None
+):
+    prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt.split("<image>")]
+
+    def insert_separator(X, sep):
+        return [ele for sublist in zip(X, [sep] * len(X)) for ele in sublist][:-1]
+
+    input_ids = []
+    offset = 0
+    if (
+        len(prompt_chunks) > 0
+        and len(prompt_chunks[0]) > 0
+        and prompt_chunks[0][0] == tokenizer.bos_token_id
+    ):
+        offset = 1
+        input_ids.append(prompt_chunks[0][0])
+
+    for x in insert_separator(prompt_chunks, [image_token_index] * (offset + 1)):
+        input_ids.extend(x[offset:])
+
+    if return_tensors is not None:
+        if return_tensors == "pt":
+            return torch.tensor(input_ids, dtype=torch.long)
+        raise ValueError(f"Unsupported tensor type: {return_tensors}")
+    return input_ids
+
+
+def get_model_name_from_path(model_path):
+    model_path = model_path.strip("/")
+    model_paths = model_path.split("/")
+    if model_paths[-1].startswith("checkpoint-"):
+        return model_paths[-2] + "_" + model_paths[-1]
+    else:
+        return model_paths[-1]
+
+
+class KeywordsStoppingCriteria(StoppingCriteria):
+    def __init__(self, keywords, tokenizer, input_ids):
+        self.keywords = keywords
+        self.keyword_ids = []
+        self.max_keyword_len = 0
+        for keyword in keywords:
+            cur_keyword_ids = tokenizer(keyword).input_ids
+            if len(cur_keyword_ids) > 1 and cur_keyword_ids[0] == tokenizer.bos_token_id:
+                cur_keyword_ids = cur_keyword_ids[1:]
+            if len(cur_keyword_ids) > self.max_keyword_len:
+                self.max_keyword_len = len(cur_keyword_ids)
+            self.keyword_ids.append(torch.tensor(cur_keyword_ids))
+        self.tokenizer = tokenizer
+        self.start_len = input_ids.shape[1]
+
+    def call_for_batch(
+        self, output_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs
+    ) -> bool:
+        offset = min(output_ids.shape[1] - self.start_len, self.max_keyword_len)
+        self.keyword_ids = [keyword_id.to(output_ids.device) for keyword_id in self.keyword_ids]
+        for keyword_id in self.keyword_ids:
+            if (output_ids[0, -keyword_id.shape[0] :] == keyword_id).all():
+                return True
+        outputs = self.tokenizer.batch_decode(output_ids[:, -offset:], skip_special_tokens=True)[0]
+        for keyword in self.keywords:
+            if keyword in outputs:
+                return True
+        return False
+
+    def __call__(self, output_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        outputs = []
+        for i in range(output_ids.shape[0]):
+            outputs.append(self.call_for_batch(output_ids[i].unsqueeze(0), scores))
+        return all(outputs)

medrax/llava/model/__init__.py

@@ -0,0 +1 @@
+from .language_model.llava_mistral import LlavaMistralForCausalLM, LlavaMistralConfig

medrax/llava/model/builder.py

@@ -0,0 +1,134 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+import torch
+from medrax.llava.model import LlavaMistralForCausalLM
+from medrax.llava.constants import (
+    DEFAULT_IMAGE_PATCH_TOKEN,
+    DEFAULT_IM_START_TOKEN,
+    DEFAULT_IM_END_TOKEN,
+)
+
+
+def load_pretrained_model(
+    model_path,
+    model_base,
+    model_name,
+    load_in_8bit=False,
+    load_in_4bit=True,
+    device="cuda",
+    cache_dir: str = "/model-weights",
+    low_cpu_mem_usage=True,
+    torch_dtype=torch.bfloat16,
+):
+
+    kwargs = {}
+
+    if device != "cuda":
+        kwargs["device_map"] = {"": device}
+    # else:
+    #     kwargs["device_map"] = "auto"
+
+    if load_in_8bit:
+        kwargs["load_in_8bit"] = True
+    elif load_in_4bit:
+        # kwargs["load_in_4bit"] = True
+        kwargs["quantization_config"] = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch_dtype,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+        )
+    # else:
+    # kwargs["torch_dtype"] = torch_dtype
+
+    if "llava" in model_name.lower():
+        # Load LLaVA model
+        if "mistral" in model_name.lower():
+            tokenizer = AutoTokenizer.from_pretrained(model_path, cache_dir=cache_dir)
+            model = LlavaMistralForCausalLM.from_pretrained(
+                model_path,
+                low_cpu_mem_usage=low_cpu_mem_usage,
+                use_flash_attention_2=False,
+                cache_dir=cache_dir,
+                torch_dtype=torch_dtype,
+                **kwargs,
+            )
+
+    else:
+        # Load language model
+        if model_base is not None:
+            # PEFT model
+            from peft import PeftModel
+
+            tokenizer = AutoTokenizer.from_pretrained(
+                model_base, use_fast=False, cache_dir=cache_dir
+            )
+            model = AutoModelForCausalLM.from_pretrained(
+                model_base,
+                low_cpu_mem_usage=True,
+                cache_dir=cache_dir,
+                torch_dtype=torch_dtype,
+                **kwargs,
+            )
+            print(f"Loading LoRA weights from {model_path}")
+            model = PeftModel.from_pretrained(model, model_path)
+            print("Merging weights")
+            model = model.merge_and_unload()
+            print("Convert to FP16...")
+            model.to(torch_dtype)
+        else:
+            use_fast = False
+            if "mpt" in model_name.lower():
+                tokenizer = AutoTokenizer.from_pretrained(
+                    model_path, use_fast=True, cache_dir=cache_dir
+                )
+                model = AutoModelForCausalLM.from_pretrained(
+                    model_path,
+                    low_cpu_mem_usage=True,
+                    trust_remote_code=True,
+                    cache_dir=cache_dir,
+                    torch_dtype=torch_dtype,
+                    **kwargs,
+                )
+            else:
+                tokenizer = AutoTokenizer.from_pretrained(
+                    model_path, use_fast=False, cache_dir=cache_dir
+                )
+                model = AutoModelForCausalLM.from_pretrained(
+                    model_path,
+                    low_cpu_mem_usage=True,
+                    cache_dir=cache_dir,
+                    torch_dtype=torch_dtype,
+                    **kwargs,
+                )
+
+    image_processor = None
+
+    if "llava" in model_name.lower():  # or 'mistral' in model_name.lower():
+        mm_use_im_start_end = getattr(model.config, "mm_use_im_start_end", False)
+        mm_use_im_patch_token = getattr(model.config, "mm_use_im_patch_token", True)
+        if mm_use_im_patch_token:
+            tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
+        if mm_use_im_start_end:
+            tokenizer.add_tokens(
+                [DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True
+            )
+        model.resize_token_embeddings(len(tokenizer))
+
+        vision_tower = model.get_vision_tower()
+        if not vision_tower.is_loaded:
+            vision_tower.load_model()
+
+        vision_tower.to(device=device, dtype=torch_dtype)
+        model.model.mm_projector.to(device=device, dtype=torch_dtype)
+
+        if not (load_in_4bit or load_in_8bit):
+            model.to(device=device, dtype=torch_dtype)
+
+        image_processor = vision_tower.image_processor
+
+    if hasattr(model.config, "max_sequence_length"):
+        context_len = model.config.max_sequence_length
+    else:
+        context_len = 2048
+
+    return tokenizer, model, image_processor, context_len

medrax/llava/model/language_model/llava_mistral.py

@@ -0,0 +1,144 @@
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    MistralConfig,
+    MistralModel,
+    MistralForCausalLM,
+)
+
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.generation.utils import GenerateOutput
+
+from ..llava_arch import LlavaMetaModel, LlavaMetaForCausalLM
+
+
+class LlavaMistralConfig(MistralConfig):
+    model_type = "llava_mistral"
+
+
+class LlavaMistralModel(LlavaMetaModel, MistralModel):
+    config_class = LlavaMistralConfig
+
+    def __init__(self, config: MistralConfig):
+        super(LlavaMistralModel, self).__init__(config)
+
+
+class LlavaMistralForCausalLM(MistralForCausalLM, LlavaMetaForCausalLM):
+    config_class = LlavaMistralConfig
+
+    def __init__(self, config):
+        super(MistralForCausalLM, self).__init__(config)
+        self.model = LlavaMistralModel(config)
+
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_model(self):
+        return self.model
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        images: Optional[torch.FloatTensor] = None,
+        image_sizes: Optional[List[List[int]]] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[str] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+
+        if inputs_embeds is None:
+            (
+                input_ids,
+                position_ids,
+                attention_mask,
+                past_key_values,
+                inputs_embeds,
+                labels,
+            ) = self.prepare_inputs_labels_for_multimodal(
+                input_ids,
+                position_ids,
+                attention_mask,
+                past_key_values,
+                labels,
+                images,
+                image_sizes,
+            )
+
+        return super().forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            labels=labels,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+    @torch.no_grad()
+    def generate(
+        self,
+        inputs: Optional[torch.Tensor] = None,
+        images: Optional[torch.Tensor] = None,
+        image_sizes: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> Union[GenerateOutput, torch.LongTensor]:
+        position_ids = kwargs.pop("position_ids", None)
+        attention_mask = kwargs.pop("attention_mask", None)
+        if "inputs_embeds" in kwargs:
+            raise NotImplementedError("`inputs_embeds` is not supported")
+
+        if images is not None:
+            (
+                inputs,
+                position_ids,
+                attention_mask,
+                _,
+                inputs_embeds,
+                _,
+            ) = self.prepare_inputs_labels_for_multimodal(
+                inputs, position_ids, attention_mask, None, None, images, image_sizes=image_sizes
+            )
+        else:
+            inputs_embeds = self.get_model().embed_tokens(inputs)
+
+        return super().generate(
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            **kwargs,
+        )
+
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs
+    ):
+        images = kwargs.pop("images", None)
+        image_sizes = kwargs.pop("image_sizes", None)
+        inputs = super().prepare_inputs_for_generation(
+            input_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, **kwargs
+        )
+        if images is not None:
+            inputs["images"] = images
+        if image_sizes is not None:
+            inputs["image_sizes"] = image_sizes
+        return inputs
+
+
+AutoConfig.register("llava_mistral", LlavaMistralConfig)
+AutoModelForCausalLM.register(LlavaMistralConfig, LlavaMistralForCausalLM)

medrax/llava/model/llava_arch.py

@@ -0,0 +1,396 @@
+#    Copyright 2023 Haotian Liu
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+
+from abc import ABC, abstractmethod
+import os
+from glob import glob
+
+import torch
+
+from .multimodal_encoder.builder import build_vision_tower
+from .multimodal_projector.builder import build_vision_projector
+
+from medrax.llava.constants import (
+    IGNORE_INDEX,
+    IMAGE_TOKEN_INDEX,
+    DEFAULT_IMAGE_PATCH_TOKEN,
+    DEFAULT_IM_START_TOKEN,
+    DEFAULT_IM_END_TOKEN,
+)
+
+
+class LlavaMetaModel:
+    def __init__(self, config):
+        super(LlavaMetaModel, self).__init__(config)
+
+        if hasattr(config, "mm_vision_tower"):
+            self.vision_tower = build_vision_tower(config, delay_load=True)
+            self.mm_projector = build_vision_projector(config)
+
+    def get_vision_tower(self):
+        vision_tower = getattr(self, "vision_tower", None)
+        if type(vision_tower) is list:
+            vision_tower = vision_tower[0]
+        return vision_tower
+
+    def initialize_vision_modules(self, model_args, fsdp=None, embed_tokens=None):
+        vision_tower = model_args.vision_tower
+        mm_vision_select_layer = model_args.mm_vision_select_layer
+        mm_vision_select_feature = model_args.mm_vision_select_feature
+        pretrain_mm_mlp_adapter = model_args.pretrain_mm_mlp_adapter
+
+        self.config.mm_vision_tower = vision_tower
+
+        if self.get_vision_tower() is None:
+            vision_tower = build_vision_tower(model_args)
+
+            if fsdp is not None and len(fsdp) > 0:
+                self.vision_tower = [vision_tower]
+            else:
+                self.vision_tower = vision_tower
+        else:
+            if fsdp is not None and len(fsdp) > 0:
+                vision_tower = self.vision_tower[0]
+            else:
+                vision_tower = self.vision_tower
+            vision_tower.load_model()
+
+        self.config.use_mm_proj = True
+        self.config.mm_projector_type = getattr(model_args, "mm_projector_type", "linear")
+        self.config.mm_hidden_size = vision_tower.hidden_size
+        self.config.mm_vision_select_layer = mm_vision_select_layer
+        self.config.mm_vision_select_feature = mm_vision_select_feature
+
+        # add additional configs for segtok
+        self.config.feature_outs = model_args.feature_outs
+        self.config.img_size = model_args.img_size
+        self.config.vision_backbone = model_args.vision_backbone
+        self.config.segtok_posembed = model_args.segtok_posembed
+
+        if getattr(self, "mm_projector", None) is None:
+            self.mm_projector = build_vision_projector(self.config)
+        else:
+            # In case it is frozen by LoRA
+            for p in self.mm_projector.parameters():
+                p.requires_grad = True
+
+        # Initialize last layer in mm_projector with weight=0 and bias=mean(embed_tokens)
+        if embed_tokens is not None:
+            embed_tokens_weight = embed_tokens.weight.data
+            self.mm_projector[-1].weight.data.zero_()
+            self.mm_projector[-1].bias.data.copy_(embed_tokens_weight.mean(dim=0))
+
+        if pretrain_mm_mlp_adapter is not None:
+
+            def get_w(weights, keyword):
+                return {k.split(keyword + ".")[1]: v for k, v in weights.items() if keyword in k}
+
+            mm_projector_weights = torch.load(pretrain_mm_mlp_adapter, map_location="cpu")
+            self.mm_projector.load_state_dict(get_w(mm_projector_weights, "mm_projector"))
+
+            # also load additional learnable parameters during feature alignment
+            checkpoint_folder = os.path.dirname(pretrain_mm_mlp_adapter)
+            ckpts = glob(f"{checkpoint_folder}/checkpoint-*", recursive=False)
+            if len(ckpts) > 0:
+                vision_module_weights = torch.load(
+                    f"{ckpts[-1]}/mm_projector.bin", map_location="cpu"
+                )
+                model_dict = get_w(vision_module_weights, "vision_tower")
+                print(f"Loading vision module weights from {ckpts[-1]}/mm_projector.bin")
+                # print keys in model_dict
+                print(f"Loaded keys: {model_dict.keys()}")
+                self.vision_tower.load_state_dict(model_dict, strict=False)
+
+
+class LlavaMetaForCausalLM(ABC):
+    @abstractmethod
+    def get_model(self):
+        pass
+
+    def get_vision_tower(self):
+        return self.get_model().get_vision_tower()
+
+    def encode_images(self, images):
+        image_features = self.get_model().get_vision_tower()(images)
+        image_features = self.get_model().mm_projector(image_features)
+        return image_features
+
+    def prepare_inputs_labels_for_multimodal(
+        self,
+        input_ids,
+        position_ids,
+        attention_mask,
+        past_key_values,
+        labels,
+        images,
+        image_sizes=None,
+    ):
+        vision_tower = self.get_vision_tower()
+        if vision_tower is None or images is None or input_ids.shape[1] == 1:
+            if (
+                past_key_values is not None
+                and vision_tower is not None
+                and images is not None
+                and input_ids.shape[1] == 1
+            ):
+                target_shape = past_key_values[-1][-1].shape[-2] + 1
+                attention_mask = torch.cat(
+                    (
+                        attention_mask,
+                        torch.ones(
+                            (attention_mask.shape[0], target_shape - attention_mask.shape[1]),
+                            dtype=attention_mask.dtype,
+                            device=attention_mask.device,
+                        ),
+                    ),
+                    dim=1,
+                )
+                position_ids = torch.sum(attention_mask, dim=1).unsqueeze(-1) - 1
+            return input_ids, position_ids, attention_mask, past_key_values, None, labels
+
+        if type(images) is list or images.ndim == 5:
+            concat_images = torch.cat([image for image in images], dim=0)
+            image_features = self.encode_images(concat_images)
+            split_sizes = [image.shape[0] for image in images]
+            image_features = torch.split(image_features, split_sizes, dim=0)
+            image_features = [x.flatten(0, 1).to(self.device) for x in image_features]
+        else:
+            image_features = self.encode_images(images).to(self.device)
+
+        # TODO: image start / end is not implemented here to support pretraining.
+        if getattr(self.config, "tune_mm_mlp_adapter", False) and getattr(
+            self.config, "mm_use_im_start_end", False
+        ):
+            raise NotImplementedError
+
+        # Let's just add dummy tensors if they do not exist,
+        # it is a headache to deal with None all the time.
+        # But it is not ideal, and if you have a better idea,
+        # please open an issue / submit a PR, thanks.
+        _labels = labels
+        _position_ids = position_ids
+        _attention_mask = attention_mask
+
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
+        else:
+            attention_mask = attention_mask.bool()
+        if position_ids is None:
+            position_ids = torch.arange(
+                0, input_ids.shape[1], dtype=torch.long, device=input_ids.device
+            )
+
+        if labels is None:
+            labels = torch.full_like(input_ids, IGNORE_INDEX)
+
+        input_ids = [
+            cur_input_ids[cur_attention_mask]
+            for cur_input_ids, cur_attention_mask in zip(input_ids, attention_mask)
+        ]
+        labels = [
+            cur_labels[cur_attention_mask]
+            for cur_labels, cur_attention_mask in zip(labels, attention_mask)
+        ]
+
+        new_input_embeds = []
+        new_labels = []
+        cur_image_idx = 0
+        for batch_idx, cur_input_ids in enumerate(input_ids):
+            num_images = (cur_input_ids == IMAGE_TOKEN_INDEX).sum()
+            if num_images == 0:
+                cur_image_features = image_features[cur_image_idx]
+                cur_input_embeds_1 = self.get_model().embed_tokens(cur_input_ids)
+                cur_input_embeds = torch.cat([cur_input_embeds_1, cur_image_features[0:0]], dim=0)
+                new_input_embeds.append(cur_input_embeds)
+                new_labels.append(labels[batch_idx])
+                cur_image_idx += 1
+                continue
+
+            image_token_indices = (
+                [-1]
+                + torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0].tolist()
+                + [cur_input_ids.shape[0]]
+            )
+            cur_input_ids_noim = []
+            cur_labels = labels[batch_idx]
+            cur_labels_noim = []
+            for i in range(len(image_token_indices) - 1):
+                cur_input_ids_noim.append(
+                    cur_input_ids[image_token_indices[i] + 1 : image_token_indices[i + 1]]
+                )
+                cur_labels_noim.append(
+                    cur_labels[image_token_indices[i] + 1 : image_token_indices[i + 1]]
+                )
+
+            split_sizes = [x.shape[0] for x in cur_labels_noim]
+            cur_input_embeds = self.get_model().embed_tokens(torch.cat(cur_input_ids_noim))
+            cur_input_embeds_no_im = torch.split(cur_input_embeds, split_sizes, dim=0)
+            cur_new_input_embeds = []
+            cur_new_labels = []
+
+            for i in range(num_images + 1):
+                cur_new_input_embeds.append(cur_input_embeds_no_im[i])
+                cur_new_labels.append(cur_labels_noim[i])
+                if i < num_images:
+                    cur_image_features = image_features[cur_image_idx]
+                    cur_image_idx += 1
+                    cur_new_input_embeds.append(cur_image_features)
+                    cur_new_labels.append(
+                        torch.full(
+                            (cur_image_features.shape[0],),
+                            IGNORE_INDEX,
+                            device=cur_labels.device,
+                            dtype=cur_labels.dtype,
+                        )
+                    )
+
+            cur_new_input_embeds = torch.cat(cur_new_input_embeds)
+            cur_new_labels = torch.cat(cur_new_labels)
+
+            new_input_embeds.append(cur_new_input_embeds)
+            new_labels.append(cur_new_labels)
+
+        # Truncate sequences to max length as image embeddings can make the sequence longer
+        tokenizer_model_max_length = getattr(self.config, "tokenizer_model_max_length", None)
+        if tokenizer_model_max_length is not None:
+            new_input_embeds = [x[:tokenizer_model_max_length] for x in new_input_embeds]
+            new_labels = [x[:tokenizer_model_max_length] for x in new_labels]
+
+        # Combine them
+        max_len = max(x.shape[0] for x in new_input_embeds)
+        batch_size = len(new_input_embeds)
+
+        new_input_embeds_padded = []
+        new_labels_padded = torch.full(
+            (batch_size, max_len),
+            IGNORE_INDEX,
+            dtype=new_labels[0].dtype,
+            device=new_labels[0].device,
+        )
+        attention_mask = torch.zeros(
+            (batch_size, max_len), dtype=attention_mask.dtype, device=attention_mask.device
+        )
+        position_ids = torch.zeros(
+            (batch_size, max_len), dtype=position_ids.dtype, device=position_ids.device
+        )
+
+        for i, (cur_new_embed, cur_new_labels) in enumerate(zip(new_input_embeds, new_labels)):
+            cur_len = cur_new_embed.shape[0]
+            if getattr(self.config, "tokenizer_padding_side", "right") == "left":
+                new_input_embeds_padded.append(
+                    torch.cat(
+                        (
+                            torch.zeros(
+                                (max_len - cur_len, cur_new_embed.shape[1]),
+                                dtype=cur_new_embed.dtype,
+                                device=cur_new_embed.device,
+                            ),
+                            cur_new_embed,
+                        ),
+                        dim=0,
+                    )
+                )
+                if cur_len > 0:
+                    new_labels_padded[i, -cur_len:] = cur_new_labels
+                    attention_mask[i, -cur_len:] = True
+                    position_ids[i, -cur_len:] = torch.arange(
+                        0, cur_len, dtype=position_ids.dtype, device=position_ids.device
+                    )
+            else:
+                new_input_embeds_padded.append(
+                    torch.cat(
+                        (
+                            cur_new_embed,
+                            torch.zeros(
+                                (max_len - cur_len, cur_new_embed.shape[1]),
+                                dtype=cur_new_embed.dtype,
+                                device=cur_new_embed.device,
+                            ),
+                        ),
+                        dim=0,
+                    )
+                )
+                if cur_len > 0:
+                    new_labels_padded[i, :cur_len] = cur_new_labels
+                    attention_mask[i, :cur_len] = True
+                    position_ids[i, :cur_len] = torch.arange(
+                        0, cur_len, dtype=position_ids.dtype, device=position_ids.device
+                    )
+
+        new_input_embeds = torch.stack(new_input_embeds_padded, dim=0)
+
+        if _labels is None:
+            new_labels = None
+        else:
+            new_labels = new_labels_padded
+
+        if _attention_mask is None:
+            attention_mask = None
+        else:
+            attention_mask = attention_mask.to(dtype=_attention_mask.dtype)
+
+        if _position_ids is None:
+            position_ids = None
+        return None, position_ids, attention_mask, past_key_values, new_input_embeds, new_labels
+
+    def initialize_vision_tokenizer(self, model_args, tokenizer):
+        if model_args.mm_use_im_patch_token:
+            tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
+            self.resize_token_embeddings(len(tokenizer))
+
+        if model_args.mm_use_im_start_end:
+            num_new_tokens = tokenizer.add_tokens(
+                [DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True
+            )
+            self.resize_token_embeddings(len(tokenizer))
+
+            if num_new_tokens > 0:
+                input_embeddings = self.get_input_embeddings().weight.data
+                output_embeddings = self.get_output_embeddings().weight.data
+
+                input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
+                output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(
+                    dim=0, keepdim=True
+                )
+
+                input_embeddings[-num_new_tokens:] = input_embeddings_avg
+                output_embeddings[-num_new_tokens:] = output_embeddings_avg
+
+            if model_args.tune_mm_mlp_adapter:
+                for p in self.get_input_embeddings().parameters():
+                    p.requires_grad = True
+                for p in self.get_output_embeddings().parameters():
+                    p.requires_grad = False
+
+            if model_args.pretrain_mm_mlp_adapter:
+                mm_projector_weights = torch.load(
+                    model_args.pretrain_mm_mlp_adapter, map_location="cpu"
+                )
+                embed_tokens_weight = mm_projector_weights["model.embed_tokens.weight"]
+                assert num_new_tokens == 2
+                if input_embeddings.shape == embed_tokens_weight.shape:
+                    input_embeddings[-num_new_tokens:] = embed_tokens_weight[-num_new_tokens:]
+                elif embed_tokens_weight.shape[0] == num_new_tokens:
+                    input_embeddings[-num_new_tokens:] = embed_tokens_weight
+                else:
+                    raise ValueError(
+                        f"Unexpected embed_tokens_weight shape. Pretrained: {embed_tokens_weight.shape}. Current: {input_embeddings.shape}. Numer of new tokens: {num_new_tokens}."
+                    )
+        elif model_args.mm_use_im_patch_token:
+            if model_args.tune_mm_mlp_adapter:
+                for p in self.get_input_embeddings().parameters():
+                    p.requires_grad = False
+                for p in self.get_output_embeddings().parameters():
+                    p.requires_grad = False

medrax/llava/model/multimodal_encoder/builder.py

@@ -0,0 +1,15 @@
+import os
+from .clip_encoder import CLIPVisionTower
+
+
+def build_vision_tower(vision_tower_cfg, **kwargs):
+    vision_tower = getattr(
+        vision_tower_cfg, "mm_vision_tower", getattr(vision_tower_cfg, "vision_tower", None)
+    )
+    is_absolute_path_exists = os.path.exists(vision_tower)
+    if (
+        is_absolute_path_exists
+        or vision_tower.startswith("openai")
+        or vision_tower.startswith("laion")
+    ):
+        return CLIPVisionTower(vision_tower, args=vision_tower_cfg, **kwargs)

medrax/llava/model/multimodal_encoder/clip_encoder.py

@@ -0,0 +1,83 @@
+import torch
+import torch.nn as nn
+
+from transformers import CLIPVisionModel, CLIPImageProcessor, CLIPVisionConfig
+
+
+class CLIPVisionTower(nn.Module):
+    def __init__(self, vision_tower, args, delay_load=False):
+        super().__init__()
+
+        self.is_loaded = False
+
+        self.vision_tower_name = vision_tower
+        self.select_layer = args.mm_vision_select_layer
+        self.select_feature = getattr(args, "mm_vision_select_feature", "patch")
+
+        if not delay_load:
+            self.load_model()
+        else:
+            self.cfg_only = CLIPVisionConfig.from_pretrained(self.vision_tower_name)
+
+    def load_model(self):
+        self.image_processor = CLIPImageProcessor.from_pretrained(self.vision_tower_name)
+        self.vision_tower = CLIPVisionModel.from_pretrained(self.vision_tower_name)
+        self.vision_tower.requires_grad_(False)
+
+        self.is_loaded = True
+
+    def feature_select(self, image_forward_outs):
+        image_features = image_forward_outs.hidden_states[self.select_layer]
+        if self.select_feature == "patch":
+            image_features = image_features[:, 1:]
+        elif self.select_feature == "cls_patch":
+            image_features = image_features
+        else:
+            raise ValueError(f"Unexpected select feature: {self.select_feature}")
+        return image_features
+
+    @torch.no_grad()
+    def forward(self, images):
+        if type(images) is list:
+            image_features = []
+            for image in images:
+                image_forward_out = self.vision_tower(
+                    image.to(device=self.device, dtype=self.dtype).unsqueeze(0),
+                    output_hidden_states=True,
+                )
+                image_feature = self.feature_select(image_forward_out).to(image.dtype)
+                image_features.append(image_feature)
+        else:
+            image_forward_outs = self.vision_tower(
+                images.to(device=self.device, dtype=self.dtype), output_hidden_states=True
+            )
+            image_features = self.feature_select(image_forward_outs).to(images.dtype)
+
+        return image_features
+
+    @property
+    def dummy_feature(self):
+        return torch.zeros(1, self.hidden_size, device=self.device, dtype=self.dtype)
+
+    @property
+    def dtype(self):
+        return self.vision_tower.dtype
+
+    @property
+    def device(self):
+        return self.vision_tower.device
+
+    @property
+    def config(self):
+        if self.is_loaded:
+            return self.vision_tower.config
+        else:
+            return self.cfg_only
+
+    @property
+    def hidden_size(self):
+        return self.config.hidden_size
+
+    @property
+    def num_patches(self):
+        return (self.config.image_size // self.config.patch_size) ** 2

medrax/llava/model/multimodal_projector/builder.py

@@ -0,0 +1,49 @@
+import torch.nn as nn
+import re
+
+
+class IdentityMap(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, *args, **kwargs):
+        return x
+
+    @property
+    def config(self):
+        return {"mm_projector_type": "identity"}
+
+
+class SimpleResBlock(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.pre_norm = nn.LayerNorm(channels)
+
+        self.proj = nn.Sequential(
+            nn.Linear(channels, channels), nn.GELU(), nn.Linear(channels, channels)
+        )
+
+    def forward(self, x):
+        x = self.pre_norm(x)
+        return x + self.proj(x)
+
+
+def build_vision_projector(config, delay_load=False, **kwargs):
+    projector_type = getattr(config, "mm_projector_type", "linear")
+
+    if projector_type == "linear":
+        return nn.Linear(config.mm_hidden_size, config.hidden_size)
+
+    mlp_gelu_match = re.match(r"^mlp(\d+)x_gelu$", projector_type)
+    if mlp_gelu_match:
+        mlp_depth = int(mlp_gelu_match.group(1))
+        modules = [nn.Linear(config.mm_hidden_size, config.hidden_size)]
+        for _ in range(1, mlp_depth):
+            modules.append(nn.GELU())
+            modules.append(nn.Linear(config.hidden_size, config.hidden_size))
+        return nn.Sequential(*modules)
+
+    if projector_type == "identity":
+        return IdentityMap()
+
+    raise ValueError(f"Unknown projector type: {projector_type}")

medrax/llava/serve/cli.py

@@ -0,0 +1,152 @@
+import argparse
+import torch
+
+from medrax.llava.constants import (
+    IMAGE_TOKEN_INDEX,
+    DEFAULT_IMAGE_TOKEN,
+    DEFAULT_IM_START_TOKEN,
+    DEFAULT_IM_END_TOKEN,
+)
+from medrax.llava.conversation import conv_templates, SeparatorStyle
+from medrax.llava.model.builder import load_pretrained_model
+from medrax.llava.utils import disable_torch_init
+from medrax.llava.mm_utils import (
+    process_images,
+    tokenizer_image_token,
+    get_model_name_from_path,
+    KeywordsStoppingCriteria,
+)
+
+from PIL import Image
+
+import requests
+from io import BytesIO
+from transformers import TextStreamer
+
+
+def load_image(image_file):
+    if image_file.startswith("http://") or image_file.startswith("https://"):
+        response = requests.get(image_file)
+        image = Image.open(BytesIO(response.content)).convert("RGB")
+    else:
+        image = Image.open(image_file).convert("RGB")
+    return image
+
+
+def main(args):
+    # Model
+    disable_torch_init()
+
+    model_name = get_model_name_from_path(args.model_path)
+    tokenizer, model, image_processor, context_len = load_pretrained_model(
+        args.model_path,
+        args.model_base,
+        model_name,
+        args.load_8bit,
+        args.load_4bit,
+        device=args.device,
+    )
+
+    if "llama-2" in model_name.lower():
+        conv_mode = "llava_llama_2"
+    elif "v1" in model_name.lower():
+        conv_mode = "llava_v1"
+    elif "mpt" in model_name.lower():
+        conv_mode = "mpt"
+    else:
+        conv_mode = "llava_v0"
+    conv_mode = "mistral_instruct"
+
+    if args.conv_mode is not None and conv_mode != args.conv_mode:
+        print(
+            "[WARNING] the auto inferred conversation mode is {}, while `--conv-mode` is {}, using {}".format(
+                conv_mode, args.conv_mode, args.conv_mode
+            )
+        )
+    else:
+        args.conv_mode = conv_mode
+
+    conv = conv_templates[args.conv_mode].copy()
+    if "mpt" in model_name.lower():
+        roles = ("user", "assistant")
+    else:
+        roles = conv.roles
+
+    image = load_image(args.image_file)
+    # Similar operation in model_worker.py
+    image_tensor = process_images([image], image_processor, model.config)
+    if type(image_tensor) is list:
+        image_tensor = [image.to(model.device, dtype=torch.float16) for image in image_tensor]
+    else:
+        image_tensor = image_tensor.to(model.device, dtype=torch.float16)
+
+    while True:
+        try:
+            inp = input(f"{roles[0]}: ")
+        except EOFError:
+            inp = ""
+        if not inp:
+            print("exit...")
+            break
+
+        print(f"{roles[1]}: ", end="")
+
+        if image is not None:
+            # first message
+            if model.config.mm_use_im_start_end:
+                inp = (
+                    DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + "\n" + inp
+                )
+            else:
+                inp = DEFAULT_IMAGE_TOKEN + "\n" + inp
+            conv.append_message(conv.roles[0], inp)
+            image = None
+        else:
+            # later messages
+            conv.append_message(conv.roles[0], inp)
+        conv.append_message(conv.roles[1], None)
+        prompt = conv.get_prompt()
+
+        input_ids = (
+            tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt")
+            .unsqueeze(0)
+            .to(model.device)
+        )
+        stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
+        keywords = [stop_str]
+        stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
+        streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+        with torch.inference_mode():
+            output_ids = model.generate(
+                input_ids,
+                images=image_tensor,
+                do_sample=True if args.temperature > 0 else False,
+                temperature=args.temperature,
+                max_new_tokens=args.max_new_tokens,
+                streamer=streamer,
+                use_cache=True,
+                stopping_criteria=[stopping_criteria],
+            )
+
+        outputs = tokenizer.decode(output_ids[0, input_ids.shape[1] :]).strip()
+        conv.messages[-1][-1] = outputs
+
+        if args.debug:
+            print("\n", {"prompt": prompt, "outputs": outputs}, "\n")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", type=str, default="facebook/opt-350m")
+    parser.add_argument("--model-base", type=str, default=None)
+    parser.add_argument("--image-file", type=str, required=True)
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--conv-mode", type=str, default=None)
+    parser.add_argument("--temperature", type=float, default=0.2)
+    parser.add_argument("--max-new-tokens", type=int, default=512)
+    parser.add_argument("--load-8bit", action="store_true")
+    parser.add_argument("--load-4bit", action="store_true")
+    parser.add_argument("--debug", action="store_true")
+    args = parser.parse_args()
+    main(args)

medrax/llava/serve/controller.py

@@ -0,0 +1,299 @@
+"""
+A controller manages distributed workers.
+It sends worker addresses to clients.
+"""
+import argparse
+import dataclasses
+from enum import Enum, auto
+import json
+import time
+from typing import List
+import threading
+
+from fastapi import FastAPI, Request
+from fastapi.responses import StreamingResponse
+import numpy as np
+import requests
+import uvicorn
+
+from medrax.llava.constants import CONTROLLER_HEART_BEAT_EXPIRATION
+from medrax.llava.utils import build_logger, server_error_msg
+
+
+logger = build_logger("controller", "controller.log")
+
+
+class DispatchMethod(Enum):
+    LOTTERY = auto()
+    SHORTEST_QUEUE = auto()
+
+    @classmethod
+    def from_str(cls, name):
+        if name == "lottery":
+            return cls.LOTTERY
+        elif name == "shortest_queue":
+            return cls.SHORTEST_QUEUE
+        else:
+            raise ValueError("Invalid dispatch method")
+
+
+@dataclasses.dataclass
+class WorkerInfo:
+    model_names: List[str]
+    speed: int
+    queue_length: int
+    check_heart_beat: bool
+    last_heart_beat: str
+
+
+def heart_beat_controller(controller):
+    while True:
+        time.sleep(CONTROLLER_HEART_BEAT_EXPIRATION)
+        controller.remove_stable_workers_by_expiration()
+
+
+class Controller:
+    def __init__(self, dispatch_method: str):
+        # Dict[str -> WorkerInfo]
+        self.worker_info = {}
+        self.dispatch_method = DispatchMethod.from_str(dispatch_method)
+
+        self.heart_beat_thread = threading.Thread(target=heart_beat_controller, args=(self,))
+        self.heart_beat_thread.start()
+
+        logger.info("Init controller")
+
+    def register_worker(self, worker_name: str, check_heart_beat: bool, worker_status: dict):
+        if worker_name not in self.worker_info:
+            logger.info(f"Register a new worker: {worker_name}")
+        else:
+            logger.info(f"Register an existing worker: {worker_name}")
+
+        if not worker_status:
+            worker_status = self.get_worker_status(worker_name)
+        if not worker_status:
+            return False
+
+        self.worker_info[worker_name] = WorkerInfo(
+            worker_status["model_names"],
+            worker_status["speed"],
+            worker_status["queue_length"],
+            check_heart_beat,
+            time.time(),
+        )
+
+        logger.info(f"Register done: {worker_name}, {worker_status}")
+        return True
+
+    def get_worker_status(self, worker_name: str):
+        try:
+            r = requests.post(worker_name + "/worker_get_status", timeout=5)
+        except requests.exceptions.RequestException as e:
+            logger.error(f"Get status fails: {worker_name}, {e}")
+            return None
+
+        if r.status_code != 200:
+            logger.error(f"Get status fails: {worker_name}, {r}")
+            return None
+
+        return r.json()
+
+    def remove_worker(self, worker_name: str):
+        del self.worker_info[worker_name]
+
+    def refresh_all_workers(self):
+        old_info = dict(self.worker_info)
+        self.worker_info = {}
+
+        for w_name, w_info in old_info.items():
+            if not self.register_worker(w_name, w_info.check_heart_beat, None):
+                logger.info(f"Remove stale worker: {w_name}")
+
+    def list_models(self):
+        model_names = set()
+
+        for w_name, w_info in self.worker_info.items():
+            model_names.update(w_info.model_names)
+
+        return list(model_names)
+
+    def get_worker_address(self, model_name: str):
+        if self.dispatch_method == DispatchMethod.LOTTERY:
+            worker_names = []
+            worker_speeds = []
+            for w_name, w_info in self.worker_info.items():
+                if model_name in w_info.model_names:
+                    worker_names.append(w_name)
+                    worker_speeds.append(w_info.speed)
+            worker_speeds = np.array(worker_speeds, dtype=np.float32)
+            norm = np.sum(worker_speeds)
+            if norm < 1e-4:
+                return ""
+            worker_speeds = worker_speeds / norm
+            if True:  # Directly return address
+                pt = np.random.choice(np.arange(len(worker_names)), p=worker_speeds)
+                worker_name = worker_names[pt]
+                return worker_name
+
+            # Check status before returning
+            while True:
+                pt = np.random.choice(np.arange(len(worker_names)), p=worker_speeds)
+                worker_name = worker_names[pt]
+
+                if self.get_worker_status(worker_name):
+                    break
+                else:
+                    self.remove_worker(worker_name)
+                    worker_speeds[pt] = 0
+                    norm = np.sum(worker_speeds)
+                    if norm < 1e-4:
+                        return ""
+                    worker_speeds = worker_speeds / norm
+                    continue
+            return worker_name
+        elif self.dispatch_method == DispatchMethod.SHORTEST_QUEUE:
+            worker_names = []
+            worker_qlen = []
+            for w_name, w_info in self.worker_info.items():
+                if model_name in w_info.model_names:
+                    worker_names.append(w_name)
+                    worker_qlen.append(w_info.queue_length / w_info.speed)
+            if len(worker_names) == 0:
+                return ""
+            min_index = np.argmin(worker_qlen)
+            w_name = worker_names[min_index]
+            self.worker_info[w_name].queue_length += 1
+            logger.info(f"names: {worker_names}, queue_lens: {worker_qlen}, ret: {w_name}")
+            return w_name
+        else:
+            raise ValueError(f"Invalid dispatch method: {self.dispatch_method}")
+
+    def receive_heart_beat(self, worker_name: str, queue_length: int):
+        if worker_name not in self.worker_info:
+            logger.info(f"Receive unknown heart beat. {worker_name}")
+            return False
+
+        self.worker_info[worker_name].queue_length = queue_length
+        self.worker_info[worker_name].last_heart_beat = time.time()
+        logger.info(f"Receive heart beat. {worker_name}")
+        return True
+
+    def remove_stable_workers_by_expiration(self):
+        expire = time.time() - CONTROLLER_HEART_BEAT_EXPIRATION
+        to_delete = []
+        for worker_name, w_info in self.worker_info.items():
+            if w_info.check_heart_beat and w_info.last_heart_beat < expire:
+                to_delete.append(worker_name)
+
+        for worker_name in to_delete:
+            self.remove_worker(worker_name)
+
+    def worker_api_generate_stream(self, params):
+        worker_addr = self.get_worker_address(params["model"])
+        if not worker_addr:
+            logger.info(f"no worker: {params['model']}")
+            ret = {
+                "text": server_error_msg,
+                "error_code": 2,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+
+        try:
+            response = requests.post(
+                worker_addr + "/worker_generate_stream", json=params, stream=True, timeout=5
+            )
+            for chunk in response.iter_lines(decode_unicode=False, delimiter=b"\0"):
+                if chunk:
+                    yield chunk + b"\0"
+        except requests.exceptions.RequestException:
+            logger.info(f"worker timeout: {worker_addr}")
+            ret = {
+                "text": server_error_msg,
+                "error_code": 3,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+
+    # Let the controller act as a worker to achieve hierarchical
+    # management. This can be used to connect isolated sub networks.
+    def worker_api_get_status(self):
+        model_names = set()
+        speed = 0
+        queue_length = 0
+
+        for w_name in self.worker_info:
+            worker_status = self.get_worker_status(w_name)
+            if worker_status is not None:
+                model_names.update(worker_status["model_names"])
+                speed += worker_status["speed"]
+                queue_length += worker_status["queue_length"]
+
+        return {
+            "model_names": list(model_names),
+            "speed": speed,
+            "queue_length": queue_length,
+        }
+
+
+app = FastAPI()
+
+
+@app.post("/register_worker")
+async def register_worker(request: Request):
+    data = await request.json()
+    controller.register_worker(
+        data["worker_name"], data["check_heart_beat"], data.get("worker_status", None)
+    )
+
+
+@app.post("/refresh_all_workers")
+async def refresh_all_workers():
+    models = controller.refresh_all_workers()
+
+
+@app.post("/list_models")
+async def list_models():
+    models = controller.list_models()
+    return {"models": models}
+
+
+@app.post("/get_worker_address")
+async def get_worker_address(request: Request):
+    data = await request.json()
+    addr = controller.get_worker_address(data["model"])
+    return {"address": addr}
+
+
+@app.post("/receive_heart_beat")
+async def receive_heart_beat(request: Request):
+    data = await request.json()
+    exist = controller.receive_heart_beat(data["worker_name"], data["queue_length"])
+    return {"exist": exist}
+
+
+@app.post("/worker_generate_stream")
+async def worker_api_generate_stream(request: Request):
+    params = await request.json()
+    generator = controller.worker_api_generate_stream(params)
+    return StreamingResponse(generator)
+
+
+@app.post("/worker_get_status")
+async def worker_api_get_status(request: Request):
+    return controller.worker_api_get_status()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="localhost")
+    parser.add_argument("--port", type=int, default=21001)
+    parser.add_argument(
+        "--dispatch-method",
+        type=str,
+        choices=["lottery", "shortest_queue"],
+        default="shortest_queue",
+    )
+    args = parser.parse_args()
+    logger.info(f"args: {args}")
+
+    controller = Controller(args.dispatch_method)
+    uvicorn.run(app, host=args.host, port=args.port, log_level="info")

medrax/llava/serve/gradio_web_server.py

@@ -0,0 +1,532 @@
+import argparse
+import datetime
+import json
+import os
+import time
+
+import gradio as gr
+import requests
+
+from medrax.llava.conversation import default_conversation, conv_templates, SeparatorStyle
+from medrax.llava.constants import LOGDIR
+from medrax.llava.utils import build_logger, server_error_msg, violates_moderation, moderation_msg
+import hashlib
+
+
+logger = build_logger("gradio_web_server", "gradio_web_server.log")
+
+headers = {"User-Agent": "LLaVA-Med Client"}
+
+no_change_btn = gr.Button.update()
+enable_btn = gr.Button.update(interactive=True)
+disable_btn = gr.Button.update(interactive=False)
+
+priority = {
+    "vicuna-13b": "aaaaaaa",
+    "koala-13b": "aaaaaab",
+}
+
+
+def get_conv_log_filename():
+    t = datetime.datetime.now()
+    name = os.path.join(LOGDIR, f"{t.year}-{t.month:02d}-{t.day:02d}-conv.json")
+    return name
+
+
+def get_model_list():
+    ret = requests.post(args.controller_url + "/refresh_all_workers")
+    assert ret.status_code == 200
+    ret = requests.post(args.controller_url + "/list_models")
+    models = ret.json()["models"]
+    models.sort(key=lambda x: priority.get(x, x))
+    logger.info(f"Models: {models}")
+    return models
+
+
+get_window_url_params = """
+function() {
+    const params = new URLSearchParams(window.location.search);
+    url_params = Object.fromEntries(params);
+    console.log(url_params);
+    return url_params;
+    }
+"""
+
+
+def load_demo(url_params, request: gr.Request):
+    logger.info(f"load_demo. ip: {request.client.host}. params: {url_params}")
+
+    dropdown_update = gr.Dropdown.update(visible=True)
+    if "model" in url_params:
+        model = url_params["model"]
+        if model in models:
+            dropdown_update = gr.Dropdown.update(value=model, visible=True)
+
+    state = default_conversation.copy()
+    return state, dropdown_update
+
+
+def load_demo_refresh_model_list(request: gr.Request):
+    logger.info(f"load_demo. ip: {request.client.host}")
+    models = get_model_list()
+    state = default_conversation.copy()
+    dropdown_update = gr.Dropdown.update(choices=models, value=models[0] if len(models) > 0 else "")
+    return state, dropdown_update
+
+
+def vote_last_response(state, vote_type, model_selector, request: gr.Request):
+    with open(get_conv_log_filename(), "a") as fout:
+        data = {
+            "tstamp": round(time.time(), 4),
+            "type": vote_type,
+            "model": model_selector,
+            "state": state.dict(),
+            "ip": request.client.host,
+        }
+        fout.write(json.dumps(data) + "\n")
+
+
+def upvote_last_response(state, model_selector, request: gr.Request):
+    logger.info(f"upvote. ip: {request.client.host}")
+    vote_last_response(state, "upvote", model_selector, request)
+    return ("",) + (disable_btn,) * 3
+
+
+def downvote_last_response(state, model_selector, request: gr.Request):
+    logger.info(f"downvote. ip: {request.client.host}")
+    vote_last_response(state, "downvote", model_selector, request)
+    return ("",) + (disable_btn,) * 3
+
+
+def flag_last_response(state, model_selector, request: gr.Request):
+    logger.info(f"flag. ip: {request.client.host}")
+    vote_last_response(state, "flag", model_selector, request)
+    return ("",) + (disable_btn,) * 3
+
+
+def regenerate(state, image_process_mode, request: gr.Request):
+    logger.info(f"regenerate. ip: {request.client.host}")
+    state.messages[-1][-1] = None
+    prev_human_msg = state.messages[-2]
+    if type(prev_human_msg[1]) in (tuple, list):
+        prev_human_msg[1] = (*prev_human_msg[1][:2], image_process_mode)
+    state.skip_next = False
+    return (state, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 5
+
+
+def clear_history(request: gr.Request):
+    logger.info(f"clear_history. ip: {request.client.host}")
+    state = default_conversation.copy()
+    return (state, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 5
+
+
+def add_text(state, text, image, image_process_mode, request: gr.Request):
+    logger.info(f"add_text. ip: {request.client.host}. len: {len(text)}")
+    if len(text) <= 0 and image is None:
+        state.skip_next = True
+        return (state, state.to_gradio_chatbot(), "", None) + (no_change_btn,) * 5
+    if args.moderate:
+        flagged = violates_moderation(text)
+        if flagged:
+            state.skip_next = True
+            return (state, state.to_gradio_chatbot(), moderation_msg, None) + (no_change_btn,) * 5
+
+    text = text[:1536]  # Hard cut-off
+    if image is not None:
+        text = text[:1200]  # Hard cut-off for images
+        if "<image>" not in text:
+            # text = '<Image><image></Image>' + text
+            text = text + "\n<image>"
+        text = (text, image, image_process_mode)
+        if len(state.get_images(return_pil=True)) > 0:
+            state = default_conversation.copy()
+    state.append_message(state.roles[0], text)
+    state.append_message(state.roles[1], None)
+    state.skip_next = False
+    return (state, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 5
+
+
+def http_bot(state, model_selector, temperature, top_p, max_new_tokens, request: gr.Request):
+    logger.info(f"http_bot. ip: {request.client.host}")
+    start_tstamp = time.time()
+    model_name = model_selector
+
+    if state.skip_next:
+        # This generate call is skipped due to invalid inputs
+        yield (state, state.to_gradio_chatbot()) + (no_change_btn,) * 5
+        return
+
+    if len(state.messages) == state.offset + 2:
+        # First round of conversation
+        if "llava" in model_name.lower():
+            if "llama-2" in model_name.lower():
+                template_name = "llava_llama_2"
+            elif "v1" in model_name.lower():
+                if "mmtag" in model_name.lower():
+                    template_name = "v1_mmtag"
+                elif "plain" in model_name.lower() and "finetune" not in model_name.lower():
+                    template_name = "v1_mmtag"
+                else:
+                    template_name = "llava_v1"
+            elif "mpt" in model_name.lower():
+                template_name = "mpt"
+            else:
+                if "mmtag" in model_name.lower():
+                    template_name = "v0_mmtag"
+                elif "plain" in model_name.lower() and "finetune" not in model_name.lower():
+                    template_name = "v0_mmtag"
+                else:
+                    template_name = "llava_v0"
+        elif "mpt" in model_name:
+            template_name = "mpt_text"
+        elif "llama-2" in model_name:
+            template_name = "llama_2"
+        else:
+            template_name = "vicuna_v1"
+        template_name = "mistral_instruct"  # FIXME: overwrite
+        new_state = conv_templates[template_name].copy()
+        new_state.append_message(new_state.roles[0], state.messages[-2][1])
+        new_state.append_message(new_state.roles[1], None)
+        state = new_state
+
+    # Query worker address
+    controller_url = args.controller_url
+    ret = requests.post(controller_url + "/get_worker_address", json={"model": model_name})
+    worker_addr = ret.json()["address"]
+    logger.info(f"model_name: {model_name}, worker_addr: {worker_addr}")
+
+    # No available worker
+    if worker_addr == "":
+        state.messages[-1][-1] = server_error_msg
+        yield (
+            state,
+            state.to_gradio_chatbot(),
+            disable_btn,
+            disable_btn,
+            disable_btn,
+            enable_btn,
+            enable_btn,
+        )
+        return
+
+    # Construct prompt
+    prompt = state.get_prompt()
+
+    all_images = state.get_images(return_pil=True)
+    all_image_hash = [hashlib.md5(image.tobytes()).hexdigest() for image in all_images]
+    for image, hash in zip(all_images, all_image_hash):
+        t = datetime.datetime.now()
+        filename = os.path.join(
+            LOGDIR, "serve_images", f"{t.year}-{t.month:02d}-{t.day:02d}", f"{hash}.jpg"
+        )
+        if not os.path.isfile(filename):
+            os.makedirs(os.path.dirname(filename), exist_ok=True)
+            image.save(filename)
+
+    # Make requests
+    pload = {
+        "model": model_name,
+        "prompt": prompt,
+        "temperature": float(temperature),
+        "top_p": float(top_p),
+        "max_new_tokens": min(int(max_new_tokens), 1536),
+        "stop": state.sep
+        if state.sep_style in [SeparatorStyle.SINGLE, SeparatorStyle.MPT]
+        else state.sep2,
+        "images": f"List of {len(state.get_images())} images: {all_image_hash}",
+    }
+    logger.info(f"==== request ====\n{pload}")
+
+    pload["images"] = state.get_images()
+
+    state.messages[-1][-1] = "▌"
+    yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 5
+
+    try:
+        # Stream output
+        response = requests.post(
+            worker_addr + "/worker_generate_stream",
+            headers=headers,
+            json=pload,
+            stream=True,
+            timeout=10,
+        )
+        for chunk in response.iter_lines(decode_unicode=False, delimiter=b"\0"):
+            if chunk:
+                data = json.loads(chunk.decode())
+                if data["error_code"] == 0:
+                    output = data["text"][len(prompt) :].strip()
+                    state.messages[-1][-1] = output + "▌"
+                    yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 5
+                else:
+                    output = data["text"] + f" (error_code: {data['error_code']})"
+                    state.messages[-1][-1] = output
+                    yield (state, state.to_gradio_chatbot()) + (
+                        disable_btn,
+                        disable_btn,
+                        disable_btn,
+                        enable_btn,
+                        enable_btn,
+                    )
+                    return
+                time.sleep(0.03)
+    except requests.exceptions.RequestException:
+        state.messages[-1][-1] = server_error_msg
+        yield (state, state.to_gradio_chatbot()) + (
+            disable_btn,
+            disable_btn,
+            disable_btn,
+            enable_btn,
+            enable_btn,
+        )
+        return
+
+    state.messages[-1][-1] = state.messages[-1][-1][:-1]
+    yield (state, state.to_gradio_chatbot()) + (enable_btn,) * 5
+
+    finish_tstamp = time.time()
+    logger.info(f"{output}")
+
+    with open(get_conv_log_filename(), "a") as fout:
+        data = {
+            "tstamp": round(finish_tstamp, 4),
+            "type": "chat",
+            "model": model_name,
+            "start": round(start_tstamp, 4),
+            "finish": round(finish_tstamp, 4),
+            "state": state.dict(),
+            "images": all_image_hash,
+            "ip": request.client.host,
+        }
+        fout.write(json.dumps(data) + "\n")
+
+
+title_markdown = """
+# 🌋 LLaVA-Med: Large Language and Vision Assistant for Medical Research
+[[Project Page]](https://llava-vl.github.io) [[Paper]](https://arxiv.org/abs/2304.08485) [[Code]](https://github.com/haotian-liu/LLaVA) [[Model]](https://huggingface.co/liuhaotian/LLaVA-13b-delta-v0)
+"""
+
+tos_markdown = """
+### Terms of use
+By using this service, users are required to agree to the following terms:
+The service is a research preview intended for non-commercial use only. It only provides limited safety measures and may generate offensive content. It must not be used for any illegal, harmful, violent, racist, or sexual purposes. The service may collect user dialogue data for future research.
+Please click the "Flag" button if you get any inappropriate answer! We will collect those to keep improving our moderator.
+For an optimal experience, please use desktop computers for this demo, as mobile devices may compromise its quality.
+"""
+
+
+learn_more_markdown = """
+### License
+The service is a research preview intended for non-commercial use only, subject to the model [License](https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md) of LLaMA, [Terms of Use](https://openai.com/policies/terms-of-use) of the data generated by OpenAI, and [Privacy Practices](https://chrome.google.com/webstore/detail/sharegpt-share-your-chatg/daiacboceoaocpibfodeljbdfacokfjb) of ShareGPT. Please contact us if you find any potential violation.
+"""
+
+block_css = """
+
+#buttons button {
+    min-width: min(120px,100%);
+}
+
+"""
+
+
+def build_demo(embed_mode):
+    textbox = gr.Textbox(
+        show_label=False, placeholder="Enter text and press ENTER", container=False
+    )
+    with gr.Blocks(title="LLaVA", theme=gr.themes.Default(), css=block_css) as demo:
+        state = gr.State()
+
+        if not embed_mode:
+            gr.Markdown(title_markdown)
+
+        with gr.Row():
+            with gr.Column(scale=3):
+                with gr.Row(elem_id="model_selector_row"):
+                    model_selector = gr.Dropdown(
+                        choices=models,
+                        value=models[0] if len(models) > 0 else "",
+                        interactive=True,
+                        show_label=False,
+                        container=False,
+                    )
+
+                imagebox = gr.Image(type="pil")
+                image_process_mode = gr.Radio(
+                    ["Crop", "Resize", "Pad", "Default"],
+                    value="Default",
+                    label="Preprocess for non-square image",
+                    visible=False,
+                )
+
+                cur_dir = os.path.dirname(os.path.abspath(__file__))
+                gr.Examples(
+                    examples=[
+                        [f"{cur_dir}/examples/bio_patch.png", "What is this image about?"],
+                        [
+                            f"{cur_dir}/examples/med_img_1.png",
+                            "Can you describe the image in details?",
+                        ],
+                        [
+                            f"{cur_dir}/examples/xy_chromosome.jpg",
+                            "Can you describe the image in details?",
+                        ],
+                        [
+                            f"{cur_dir}/examples/synpic42202.jpg",
+                            "Is there evidence of an aortic aneurysm? Please choose from the following two options: [yes, no]?",
+                        ],  # answer" yes
+                        [
+                            f"{cur_dir}/examples/synpic32933.jpg",
+                            "What is the abnormality by the right hemidiaphragm?",
+                        ],  # answer: free air
+                        [
+                            f"{cur_dir}/examples/extreme_ironing.jpg",
+                            "What is unusual about this image?",
+                        ],
+                        [
+                            f"{cur_dir}/examples/waterview.jpg",
+                            "What are the things I should be cautious about when I visit here?",
+                        ],
+                    ],
+                    inputs=[imagebox, textbox],
+                )
+
+                with gr.Accordion("Parameters", open=False) as parameter_row:
+                    temperature = gr.Slider(
+                        minimum=0.0,
+                        maximum=1.0,
+                        value=0.2,
+                        step=0.1,
+                        interactive=True,
+                        label="Temperature",
+                    )
+                    top_p = gr.Slider(
+                        minimum=0.0,
+                        maximum=1.0,
+                        value=0.7,
+                        step=0.1,
+                        interactive=True,
+                        label="Top P",
+                    )
+                    max_output_tokens = gr.Slider(
+                        minimum=0,
+                        maximum=1024,
+                        value=512,
+                        step=64,
+                        interactive=True,
+                        label="Max output tokens",
+                    )
+
+            with gr.Column(scale=8):
+                chatbot = gr.Chatbot(elem_id="chatbot", label="LLaVA-Med Chatbot", height=550)
+                with gr.Row():
+                    with gr.Column(scale=8):
+                        textbox.render()
+                    with gr.Column(scale=1, min_width=50):
+                        submit_btn = gr.Button(value="Send", variant="primary")
+                with gr.Row(elem_id="buttons") as button_row:
+                    upvote_btn = gr.Button(value="👍  Upvote", interactive=False)
+                    downvote_btn = gr.Button(value="👎  Downvote", interactive=False)
+                    flag_btn = gr.Button(value="⚠️  Flag", interactive=False)
+                    # stop_btn = gr.Button(value="⏹️  Stop Generation", interactive=False)
+                    regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False)
+                    clear_btn = gr.Button(value="🗑️  Clear", interactive=False)
+
+        if not embed_mode:
+            gr.Markdown(tos_markdown)
+            gr.Markdown(learn_more_markdown)
+        url_params = gr.JSON(visible=False)
+
+        # Register listeners
+        btn_list = [upvote_btn, downvote_btn, flag_btn, regenerate_btn, clear_btn]
+        upvote_btn.click(
+            upvote_last_response,
+            [state, model_selector],
+            [textbox, upvote_btn, downvote_btn, flag_btn],
+            queue=False,
+        )
+        downvote_btn.click(
+            downvote_last_response,
+            [state, model_selector],
+            [textbox, upvote_btn, downvote_btn, flag_btn],
+            queue=False,
+        )
+        flag_btn.click(
+            flag_last_response,
+            [state, model_selector],
+            [textbox, upvote_btn, downvote_btn, flag_btn],
+            queue=False,
+        )
+
+        regenerate_btn.click(
+            regenerate,
+            [state, image_process_mode],
+            [state, chatbot, textbox, imagebox] + btn_list,
+            queue=False,
+        ).then(
+            http_bot,
+            [state, model_selector, temperature, top_p, max_output_tokens],
+            [state, chatbot] + btn_list,
+        )
+
+        clear_btn.click(
+            clear_history, None, [state, chatbot, textbox, imagebox] + btn_list, queue=False
+        )
+
+        textbox.submit(
+            add_text,
+            [state, textbox, imagebox, image_process_mode],
+            [state, chatbot, textbox, imagebox] + btn_list,
+            queue=False,
+        ).then(
+            http_bot,
+            [state, model_selector, temperature, top_p, max_output_tokens],
+            [state, chatbot] + btn_list,
+        )
+
+        submit_btn.click(
+            add_text,
+            [state, textbox, imagebox, image_process_mode],
+            [state, chatbot, textbox, imagebox] + btn_list,
+            queue=False,
+        ).then(
+            http_bot,
+            [state, model_selector, temperature, top_p, max_output_tokens],
+            [state, chatbot] + btn_list,
+        )
+
+        if args.model_list_mode == "once":
+            demo.load(
+                load_demo,
+                [url_params],
+                [state, model_selector],
+                _js=get_window_url_params,
+                queue=False,
+            )
+        elif args.model_list_mode == "reload":
+            demo.load(load_demo_refresh_model_list, None, [state, model_selector], queue=False)
+        else:
+            raise ValueError(f"Unknown model list mode: {args.model_list_mode}")
+
+    return demo
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="0.0.0.0")
+    parser.add_argument("--port", type=int)
+    parser.add_argument("--controller-url", type=str, default="http://localhost:21001")
+    parser.add_argument("--concurrency-count", type=int, default=10)
+    parser.add_argument("--model-list-mode", type=str, default="once", choices=["once", "reload"])
+    parser.add_argument("--share", action="store_true")
+    parser.add_argument("--moderate", action="store_true")
+    parser.add_argument("--embed", action="store_true")
+    args = parser.parse_args()
+    logger.info(f"args: {args}")
+
+    models = get_model_list()
+
+    logger.info(args)
+    demo = build_demo(args.embed)
+    demo.queue(concurrency_count=args.concurrency_count, api_open=False).launch(
+        server_name=args.host, server_port=args.port, share=args.share
+    )

medrax/llava/serve/model_worker.py

@@ -0,0 +1,337 @@
+"""
+A model worker executes the model.
+"""
+import argparse
+import asyncio
+import json
+import time
+import threading
+import uuid
+
+from fastapi import FastAPI, Request, BackgroundTasks
+from fastapi.responses import StreamingResponse
+import requests
+import torch
+import uvicorn
+from functools import partial
+
+from medrax.llava.constants import WORKER_HEART_BEAT_INTERVAL
+from medrax.llava.utils import build_logger, server_error_msg, pretty_print_semaphore
+from medrax.llava.model.builder import load_pretrained_model
+from medrax.llava.mm_utils import (
+    process_images,
+    load_image_from_base64,
+    tokenizer_image_token,
+    KeywordsStoppingCriteria,
+)
+from medrax.llava.constants import (
+    IMAGE_TOKEN_INDEX,
+    DEFAULT_IMAGE_TOKEN,
+    DEFAULT_IM_START_TOKEN,
+    DEFAULT_IM_END_TOKEN,
+)
+from transformers import TextIteratorStreamer
+from threading import Thread
+
+
+GB = 1 << 30
+
+worker_id = str(uuid.uuid4())[:6]
+logger = build_logger("model_worker", f"model_worker_{worker_id}.log")
+global_counter = 0
+
+model_semaphore = None
+
+
+def heart_beat_worker(controller):
+
+    while True:
+        time.sleep(WORKER_HEART_BEAT_INTERVAL)
+        controller.send_heart_beat()
+
+
+class ModelWorker:
+    def __init__(
+        self,
+        controller_addr,
+        worker_addr,
+        worker_id,
+        no_register,
+        model_path,
+        model_base,
+        model_name,
+        load_8bit,
+        load_4bit,
+        device,
+    ):
+        self.controller_addr = controller_addr
+        self.worker_addr = worker_addr
+        self.worker_id = worker_id
+        if model_path.endswith("/"):
+            model_path = model_path[:-1]
+        if model_name is None:
+            model_paths = model_path.split("/")
+            if model_paths[-1].startswith("checkpoint-"):
+                self.model_name = model_paths[-2] + "_" + model_paths[-1]
+            else:
+                self.model_name = model_paths[-1]
+        else:
+            self.model_name = model_name
+
+        self.device = device
+        logger.info(f"Loading the model {self.model_name} on worker {worker_id} ...")
+        self.tokenizer, self.model, self.image_processor, self.context_len = load_pretrained_model(
+            model_path, model_base, self.model_name, load_8bit, load_4bit, device=self.device
+        )
+        self.is_multimodal = "llava" in self.model_name.lower()
+
+        if not no_register:
+            self.register_to_controller()
+            self.heart_beat_thread = threading.Thread(target=heart_beat_worker, args=(self,))
+            self.heart_beat_thread.start()
+
+    def register_to_controller(self):
+        logger.info("Register to controller")
+
+        url = self.controller_addr + "/register_worker"
+        data = {
+            "worker_name": self.worker_addr,
+            "check_heart_beat": True,
+            "worker_status": self.get_status(),
+        }
+        r = requests.post(url, json=data)
+        assert r.status_code == 200
+
+    def send_heart_beat(self):
+        logger.info(
+            f"Send heart beat. Models: {[self.model_name]}. "
+            f"Semaphore: {pretty_print_semaphore(model_semaphore)}. "
+            f"global_counter: {global_counter}"
+        )
+
+        url = self.controller_addr + "/receive_heart_beat"
+
+        while True:
+            try:
+                ret = requests.post(
+                    url,
+                    json={"worker_name": self.worker_addr, "queue_length": self.get_queue_length()},
+                    timeout=5,
+                )
+                exist = ret.json()["exist"]
+                break
+            except requests.exceptions.RequestException as e:
+                logger.error(f"heart beat error: {e}")
+            time.sleep(5)
+
+        if not exist:
+            self.register_to_controller()
+
+    def get_queue_length(self):
+        if model_semaphore is None:
+            return 0
+        else:
+            return (
+                args.limit_model_concurrency
+                - model_semaphore._value
+                + (len(model_semaphore._waiters) if model_semaphore._waiters is not None else 0)
+            )
+
+    def get_status(self):
+        return {
+            "model_names": [self.model_name],
+            "speed": 1,
+            "queue_length": self.get_queue_length(),
+        }
+
+    @torch.inference_mode()
+    def generate_stream(self, params):
+        tokenizer, model, image_processor = self.tokenizer, self.model, self.image_processor
+
+        prompt = params["prompt"]
+        ori_prompt = prompt
+        images = params.get("images", None)
+        num_image_tokens = 0
+        if images is not None and len(images) > 0 and self.is_multimodal:
+            if len(images) > 0:
+                if len(images) != prompt.count(DEFAULT_IMAGE_TOKEN):
+                    raise ValueError(
+                        "Number of images does not match number of <image> tokens in prompt"
+                    )
+
+                images = [load_image_from_base64(image) for image in images]
+                images = process_images(images, image_processor, model.config)
+
+                if type(images) is list:
+                    images = [image.to(self.model.device, dtype=torch.float16) for image in images]
+                else:
+                    images = images.to(self.model.device, dtype=torch.float16)
+
+                replace_token = DEFAULT_IMAGE_TOKEN
+                if getattr(self.model.config, "mm_use_im_start_end", False):
+                    replace_token = DEFAULT_IM_START_TOKEN + replace_token + DEFAULT_IM_END_TOKEN
+                prompt = prompt.replace(DEFAULT_IMAGE_TOKEN, replace_token)
+
+                num_image_tokens = (
+                    prompt.count(replace_token) * model.get_vision_tower().num_patches
+                )
+            else:
+                images = None
+            image_args = {"images": images}
+        else:
+            images = None
+            image_args = {}
+
+        temperature = float(params.get("temperature", 1.0))
+        top_p = float(params.get("top_p", 1.0))
+        max_context_length = getattr(model.config, "max_position_embeddings", 2048)
+        max_new_tokens = min(int(params.get("max_new_tokens", 256)), 1024)
+        stop_str = params.get("stop", None)
+        do_sample = True if temperature > 0.001 else False
+
+        input_ids = (
+            tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt")
+            .unsqueeze(0)
+            .to(self.device)
+        )
+        keywords = [stop_str]
+        stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
+        streamer = TextIteratorStreamer(
+            tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=15
+        )
+
+        max_new_tokens = min(
+            max_new_tokens, max_context_length - input_ids.shape[-1] - num_image_tokens
+        )
+
+        if max_new_tokens < 1:
+            yield json.dumps(
+                {
+                    "text": ori_prompt
+                    + "Exceeds max token length. Please start a new conversation, thanks.",
+                    "error_code": 0,
+                }
+            ).encode() + b"\0"
+            return
+
+        thread = Thread(
+            target=model.generate,
+            kwargs=dict(
+                inputs=input_ids,
+                do_sample=do_sample,
+                temperature=temperature,
+                top_p=top_p,
+                max_new_tokens=max_new_tokens,
+                streamer=streamer,
+                stopping_criteria=[stopping_criteria],
+                use_cache=True,
+                **image_args,
+            ),
+        )
+        thread.start()
+
+        generated_text = ori_prompt
+        for new_text in streamer:
+            generated_text += new_text
+            if generated_text.endswith(stop_str):
+                generated_text = generated_text[: -len(stop_str)]
+            yield json.dumps({"text": generated_text, "error_code": 0}).encode() + b"\0"
+
+    def generate_stream_gate(self, params):
+        try:
+            for x in self.generate_stream(params):
+                yield x
+        except ValueError as e:
+            print("Caught ValueError:", e)
+            ret = {
+                "text": server_error_msg,
+                "error_code": 1,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+        except torch.cuda.CudaError as e:
+            print("Caught torch.cuda.CudaError:", e)
+            ret = {
+                "text": server_error_msg,
+                "error_code": 1,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+        except Exception as e:
+            print("Caught Unknown Error", e)
+            ret = {
+                "text": server_error_msg,
+                "error_code": 1,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+
+
+app = FastAPI()
+
+
+def release_model_semaphore(fn=None):
+    model_semaphore.release()
+    if fn is not None:
+        fn()
+
+
+@app.post("/worker_generate_stream")
+async def generate_stream(request: Request):
+    global model_semaphore, global_counter
+    global_counter += 1
+    params = await request.json()
+
+    if model_semaphore is None:
+        model_semaphore = asyncio.Semaphore(args.limit_model_concurrency)
+    await model_semaphore.acquire()
+    worker.send_heart_beat()
+    generator = worker.generate_stream_gate(params)
+    background_tasks = BackgroundTasks()
+    background_tasks.add_task(partial(release_model_semaphore, fn=worker.send_heart_beat))
+    return StreamingResponse(generator, background=background_tasks)
+
+
+@app.post("/worker_get_status")
+async def get_status(request: Request):
+    return worker.get_status()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="localhost")
+    parser.add_argument("--port", type=int, default=21002)
+    parser.add_argument("--worker-address", type=str, default="http://localhost:21002")
+    parser.add_argument("--controller-address", type=str, default="http://localhost:21001")
+    parser.add_argument("--model-path", type=str, default="facebook/opt-350m")
+    parser.add_argument("--model-base", type=str, default=None)
+    parser.add_argument("--model-name", type=str)
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument(
+        "--multi-modal",
+        action="store_true",
+        help="Multimodal mode is automatically detected with model name, please make sure `llava` is included in the model path.",
+    )
+    parser.add_argument("--limit-model-concurrency", type=int, default=5)
+    parser.add_argument("--stream-interval", type=int, default=1)
+    parser.add_argument("--no-register", action="store_true")
+    parser.add_argument("--load-8bit", action="store_true")
+    parser.add_argument("--load-4bit", action="store_true")
+    args = parser.parse_args()
+    logger.info(f"args: {args}")
+
+    if args.multi_modal:
+        logger.warning(
+            "Multimodal mode is automatically detected with model name, please make sure `llava` is included in the model path."
+        )
+
+    worker = ModelWorker(
+        args.controller_address,
+        args.worker_address,
+        worker_id,
+        args.no_register,
+        args.model_path,
+        args.model_base,
+        args.model_name,
+        args.load_8bit,
+        args.load_4bit,
+        args.device,
+    )
+    uvicorn.run(app, host=args.host, port=args.port, log_level="info")