├── .github └── workflows │ └── python-publish.yml ├── .gitignore ├── AMIE_pytorch ├── AMIE_pytorch.py └── __init__.py ├── LICENSE ├── README.md ├── amie.png └── setup.py /.github/workflows/python-publish.yml: -------------------------------------------------------------------------------- 1 | 2 | 3 | # This workflow will upload a Python Package using Twine when a release is created 4 | # For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries 5 | 6 | # This workflow uses actions that are not certified by GitHub. 7 | # They are provided by a third-party and are governed by 8 | # separate terms of service, privacy policy, and support 9 | # documentation. 10 | 11 | name: Upload Python Package 12 | 13 | on: 14 | release: 15 | types: [published] 16 | 17 | jobs: 18 | deploy: 19 | 20 | runs-on: ubuntu-latest 21 | 22 | steps: 23 | - uses: actions/checkout@v2 24 | - name: Set up Python 25 | uses: actions/setup-python@v2 26 | with: 27 | python-version: '3.x' 28 | - name: Install dependencies 29 | run: | 30 | python -m pip install --upgrade pip 31 | pip install build 32 | - name: Build package 33 | run: python -m build 34 | - name: Publish package 35 | uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29 36 | with: 37 | user: __token__ 38 | password: ${{ secrets.PYPI_API_TOKEN }} 39 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | share/python-wheels/ 24 | *.egg-info/ 25 | .installed.cfg 26 | *.egg 27 | MANIFEST 28 | 29 | # PyInstaller 30 | # Usually these files are written by a python script from a template 31 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 32 | *.manifest 33 | *.spec 34 | 35 | # Installer logs 36 | pip-log.txt 37 | pip-delete-this-directory.txt 38 | 39 | # Unit test / coverage reports 40 | htmlcov/ 41 | .tox/ 42 | .nox/ 43 | .coverage 44 | .coverage.* 45 | .cache 46 | nosetests.xml 47 | coverage.xml 48 | *.cover 49 | *.py,cover 50 | .hypothesis/ 51 | .pytest_cache/ 52 | cover/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | .pybuilder/ 76 | target/ 77 | 78 | # Jupyter Notebook 79 | .ipynb_checkpoints 80 | 81 | # IPython 82 | profile_default/ 83 | ipython_config.py 84 | 85 | # pyenv 86 | # For a library or package, you might want to ignore these files since the code is 87 | # intended to run in multiple environments; otherwise, check them in: 88 | # .python-version 89 | 90 | # pipenv 91 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 92 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 93 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 94 | # install all needed dependencies. 95 | #Pipfile.lock 96 | 97 | # poetry 98 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. 99 | # This is especially recommended for binary packages to ensure reproducibility, and is more 100 | # commonly ignored for libraries. 101 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control 102 | #poetry.lock 103 | 104 | # pdm 105 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. 106 | #pdm.lock 107 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it 108 | # in version control. 109 | # https://pdm.fming.dev/#use-with-ide 110 | .pdm.toml 111 | 112 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm 113 | __pypackages__/ 114 | 115 | # Celery stuff 116 | celerybeat-schedule 117 | celerybeat.pid 118 | 119 | # SageMath parsed files 120 | *.sage.py 121 | 122 | # Environments 123 | .env 124 | .venv 125 | env/ 126 | venv/ 127 | ENV/ 128 | env.bak/ 129 | venv.bak/ 130 | 131 | # Spyder project settings 132 | .spyderproject 133 | .spyproject 134 | 135 | # Rope project settings 136 | .ropeproject 137 | 138 | # mkdocs documentation 139 | /site 140 | 141 | # mypy 142 | .mypy_cache/ 143 | .dmypy.json 144 | dmypy.json 145 | 146 | # Pyre type checker 147 | .pyre/ 148 | 149 | # pytype static type analyzer 150 | .pytype/ 151 | 152 | # Cython debug symbols 153 | cython_debug/ 154 | 155 | # PyCharm 156 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can 157 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore 158 | # and can be added to the global gitignore or merged into this file. For a more nuclear 159 | # option (not recommended) you can uncomment the following to ignore the entire idea folder. 160 | #.idea/ 161 | -------------------------------------------------------------------------------- /AMIE_pytorch/AMIE_pytorch.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | import torch 4 | from torch import nn, einsum 5 | from torch.nn import Module, ModuleList 6 | 7 | from einops import rearrange 8 | 9 | from beartype import beartype 10 | 11 | # functions 12 | 13 | def exists(v): 14 | return v is not None 15 | 16 | def default(v, d): 17 | return v if exists(v) else d 18 | 19 | # self critic prompt 20 | # figure A.15 in paper 21 | 22 | PROMPT_EVALUATE_EXPLANATION = """ 23 | I have a doctor-patient dialogue and the corresponding rating that quantifies its quality according to 24 | the following criterion: (e.g., maintaining patient welfare). The rating of the dialogue is 25 | on a scale of 1 to 5 where: 26 | 27 | 5: e.g., “Treats patient respectfully, and ensures comfort, safety and dignity” 28 | 1: e.g., “Causes patient physical or emotional discomfort AND jeopardises patient safety” 29 | 30 | First, describe which parts of the dialogue are good with respect to the criterion. Then, describe which parts are bad with respect to the criterion. Lastly, summarise the above to explain the 31 | provided rating, using the following format: 32 | 33 | Good: ... 34 | Bad: ... 35 | Summary: ... 36 | 37 | DIALOGUE: 38 | Rating: 39 | EVALUATION: 40 | """ 41 | 42 | # figure A.16 43 | 44 | PROMPT_EVALUATE_QUALITATIVE = """ 45 | I have a doctor-patient dialogue which I would like you to evaluate on the following criterion: 46 | (e.g., maintaining patient welfare). The dialogue should be rated on a scale of 1-5 with 47 | respect to the criterion where: 48 | 49 | 5: e.g., “Treats patient respectfully, and ensures comfort, safety and dignity” 50 | 1: e.g., “Causes patient physical or emotional discomfort AND jeopardises patient safety” 51 | 52 | Here are some example dialogues and their ratings: 53 | DIALOGUE: 54 | EVALUATION: 55 | Rating: 56 | ... 57 | 58 | Now, please rate the following dialogue as instructed below. First, describe which parts of the dialogue 59 | are good with respect to the criterion. Then, describe which parts are bad with respect to the criterion. 60 | Third, summarise the above findings. Lastly, rate the dialogue on a scale of 1-5 with respect to the 61 | criterion, according to this schema: 62 | 63 | Good: ... 64 | Bad: ... 65 | Summary: ... 66 | Rating: ... 67 | 68 | DIALOGUE: 69 | EVALUATION: 70 | """ 71 | 72 | # self play modules 73 | 74 | class OuterSelfPlay(Module): 75 | @beartype 76 | def __init__(self): 77 | super().__init__() 78 | raise NotImplementedError 79 | 80 | class InnerSelfPlay(Module): 81 | @beartype 82 | def __init__(self): 83 | super().__init__() 84 | raise NotImplementedError 85 | 86 | class PatientAgent(Module): 87 | @beartype 88 | def __init__(self): 89 | super().__init__() 90 | raise NotImplementedError 91 | 92 | class ClinicalVignetteGenerator(Module): 93 | @beartype 94 | def __init__(self): 95 | super().__init__() 96 | raise NotImplementedError 97 | 98 | class Moderator(Module): 99 | @beartype 100 | def __init__(self): 101 | super().__init__() 102 | raise NotImplementedError 103 | 104 | class DoctorAgent(Module): 105 | @beartype 106 | def __init__(self): 107 | super().__init__() 108 | raise NotImplementedError 109 | 110 | class SimulatedDialogue(Module): 111 | @beartype 112 | def __init__(self): 113 | super().__init__() 114 | raise NotImplementedError 115 | 116 | class Critic(Module): 117 | @beartype 118 | def __init__(self): 119 | super().__init__() 120 | raise NotImplementedError 121 | 122 | # main class 123 | 124 | class AMIE(Module): 125 | @beartype 126 | def __init__(self): 127 | super().__init__() 128 | 129 | def forward(self, x): 130 | return x 131 | -------------------------------------------------------------------------------- /AMIE_pytorch/__init__.py: -------------------------------------------------------------------------------- 1 | from AMIE_pytorch.AMIE_pytorch import AMIE 2 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 Phil Wang 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | 2 | 3 | ## AMIE - Pytorch (wip) 4 | 5 | Implementation of the general framework for AMIE, from the paper Towards Conversational Diagnostic AI, out of Google Deepmind 6 | 7 | Reach out to me if you are at least a 3rd year medical student, have kept up with the current state of deep learning, and interested in this project. 8 | 9 | ## Todo 10 | 11 | - [ ] allow for a DPO-like formulation. do not think google deepmind has adopted this across the org just yet. 12 | 13 | ## Citations 14 | 15 | ```bibtex 16 | @inproceedings{Tu2024TowardsCD, 17 | title = {Towards Conversational Diagnostic AI}, 18 | author = {Tao Tu and Anil Palepu and Mike Schaekermann and Khaled Saab and Jan Freyberg and Ryutaro Tanno and Amy Wang and Brenna Li and Mohamed Amin and Nenad Toma{\vs}ev and Shekoofeh Azizi and Karan Singhal and Yong Cheng and Le Hou and Albert Webson and Kavita Kulkarni and S Sara Mahdavi and Christopher Semturs and Juraj Gottweis and Joelle Barral and Katherine Chou and Greg S. Corrado and Yossi Matias and Alan Karthikesalingam and Vivek Natarajan}, 19 | year = {2024}, 20 | url = {https://api.semanticscholar.org/CorpusID:266933212} 21 | } 22 | ``` 23 | -------------------------------------------------------------------------------- /amie.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lucidrains/AMIE-pytorch/876ca7cc8a8f105592d173f9ddc6124dde536222/amie.png -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages 2 | 3 | setup( 4 | name = 'AMIE-pytorch', 5 | packages = find_packages(exclude=[]), 6 | version = '0.0.1', 7 | license='MIT', 8 | description = 'AMIE', 9 | author = 'Phil Wang', 10 | author_email = 'lucidrains@gmail.com', 11 | long_description_content_type = 'text/markdown', 12 | url = 'https://github.com/lucidrains/AMIE-pytorch', 13 | keywords = [ 14 | 'artificial intelligence', 15 | 'deep learning', 16 | 'differential diagnosis' 17 | ], 18 | install_requires=[ 19 | 'accelerate', 20 | 'beartype', 21 | 'einops>=0.7.0', 22 | 'einx>=0.1.2', 23 | 'torch>=2.0', 24 | 'tqdm' 25 | ], 26 | classifiers=[ 27 | 'Development Status :: 4 - Beta', 28 | 'Intended Audience :: Developers', 29 | 'Topic :: Scientific/Engineering :: Artificial Intelligence', 30 | 'License :: OSI Approved :: MIT License', 31 | 'Programming Language :: Python :: 3.6', 32 | ], 33 | ) 34 | --------------------------------------------------------------------------------