[
  {
    "path": "README.md",
    "content": "# XrayPULSE\n\n<!--\n**Here are some ideas to get you started:**\n🙋‍♀️ A short introduction - what is your organization all about?\n🌈 Contribution guidelines - how can the community get involved?\n👩‍💻 Useful resources - where can the community find your docs? Is there anything else the community should know?\n🍿 Fun facts - what does your team eat for breakfast?\n🧙 Remember, you can do mighty things with the power of [Markdown](https://docs.github.com/github/writing-on-github/getting-started-with-writing-and-formatting-on-github/basic-writing-and-formatting-syntax)\n-->\n\n<!-- Insert the project banner here -->\n\n<div align=\"center\">\n    <a href=\"https://\"><img width=\"1000px\" height=\"auto\" src=\"./banner.png\"></a>\n</div>\n\n\n---\n\n<!-- Select some of the point info, feel free to delete -->\n\n\n## Key Features\n\nThis repository provides the official implementation of XrayPULSE: \n\nKey feature bulletin points here\n\n- An attempt to extend [PULSE]() to a biomedical multimodal conversational assistant. \n- XrayPULSE is fintuned on Xray-Report paired datasets in Chinese\n\n\n## Details\n\nOur model is based on PULSE. We utilize [MedCLIP](https://github.com/RyanWangZf/MedCLIP)  as our medical visual encoder and Q-former ([BLIP2](https://huggingface.co/docs/transformers/main/model_doc/blip-2)) following a simple linear transformation as the adapter to inject the image to PULSE. For aligning the frozen visual encoder and the LLM by the adapter, we generate Chinese-version Xray-Report paired data from free-text radiology reports of two datasets ([MIMIC-CXR](https://physionet.org/content/mimic-cxr-jpg/2.0.0/) and [OpenI](https://openi.nlm.nih.gov/faq#collection)) with the help of chatGPT.  To facilitate research in biomedical multimodal learning, we will release the data to the public.\n\n<!-- Insert a pipeline of your algorithm here if got one -->\n\n<div align=\"center\">\n    <a href=\"https://\"><img width=\"1000px\" height=\"auto\" src=\"./framework.png\"></a>\n</div>\n\n\n\n## Get Started\n\n**Installation**\n\n```bash\ngit clone https://github.com/openmedlab/XrayPULSE.git\ncd XrayPULSE\n```\n\n**Environment**\n\n```bash\nconda env create -f env.yml\nconda activate xraypulse\n```\n\n**Prepare the pretrained weights**\n\nYou can find the pretrained model weights.\n\n- [PULSE\\_Model](https://huggingface.co/OpenMEDLab/PULSE-7bv5) \n- [Pretrained_XrayPULSE_Checkpoint](https://drive.google.com/file/d/1VsO61-3DFuK4ysGPvoD4_JZaRFKvAJR_/view?usp=drive_link)\n\nThe weights of PULSE would be in a single folder in a structure similar to the following:\n\n```\npulse_weights\n├── config.json\n├── generation_config.json\n├── tokenizer.json\n├── tokenizer_config.json\n├── special_tokens_map.json \n├── pytorch_model.bin.index.json\n├── pytorch_model-00001-of-00002.bin\n├── pytorch_model-00002-of-00002.bin \n```\n\nThen, set the path of pulse_weights to \"bloom_model\" in the model config file \"xraypulse/configs/models/xraypulse.yaml\"\n\nAnd add the path of the pretrained checkpoint  in \"demo_configs/xraypulse_demo.yaml\".\n\n**Run Demo**\n\n```bash\nbash run_demo.sh\n```\n\n\n\n## 🙏 Acknowledgement\nThis project is built upon the gaint sholders of [XrayGPT](https://github.com/mbzuai-oryx/XrayGPT). Great thanks to it!\n\nWe used medical aware image encoder from [MedCLIP](https://github.com/RyanWangZf/MedCLIP).\n\nThe model architecture of XrayGPT follows [BLIP2](https://huggingface.co/docs/transformers/main/model_doc/blip-2).\n\n\n## 🛡️ License\n\nThis project is under the CC-BY-NC 4.0 license. See [LICENSE](LICENSE) for details.\n"
  },
  {
    "path": "demo.py",
    "content": "import argparse\nimport os\nimport random\n\nimport numpy as np\nimport torch\nimport torch.backends.cudnn as cudnn\nimport gradio as gr\n\nfrom xraypulse.common.config import Config\nfrom xraypulse.common.dist_utils import get_rank\nfrom xraypulse.common.registry import registry\nfrom xraypulse.conversation.conversation import Chat, CONV_ZH\n\n# imports modules for registration\nfrom xraypulse.datasets.builders import *\nfrom xraypulse.models import *\nfrom xraypulse.processors import *\nfrom xraypulse.runners import *\nfrom xraypulse.tasks import *\n\n\ndef parse_args():\n    parser = argparse.ArgumentParser(description=\"Demo\")\n    parser.add_argument(\"--cfg-path\", required=True, help=\"path to configuration file.\")\n    parser.add_argument(\"--gpu-id\", type=int, default=0, help=\"specify the gpu to load the model.\")\n    parser.add_argument(\n        \"--options\",\n        nargs=\"+\",\n        help=\"override some settings in the used config, the key-value pair \"\n        \"in xxx=yyy format will be merged into config file (deprecate), \"\n        \"change to --cfg-options instead.\",\n    )\n    args = parser.parse_args()\n    return args\n\n\ndef setup_seeds(config):\n    seed = config.run_cfg.seed + get_rank()\n\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n\n    cudnn.benchmark = False\n    cudnn.deterministic = True\n\n\n# ========================================\n#             Model Initialization\n# ========================================\n\nprint('Initializing Chat')\nargs = parse_args()\ncfg = Config(args)\n\nmodel_config = cfg.model_cfg\nprint(model_config)\nmodel_config.device_8bit = args.gpu_id\nmodel_cls = registry.get_model_class(model_config.arch)\nprint(model_cls)\nmodel = model_cls.from_config(model_config).to('cuda:{}'.format(args.gpu_id))\n\nvis_processor_cfg = cfg.datasets_cfg.openi.vis_processor.train\nvis_processor = registry.get_processor_class(vis_processor_cfg.name).from_config(vis_processor_cfg)\nchat = Chat(model, vis_processor, device='cuda:{}'.format(args.gpu_id))\nprint('Initialization Finished')\n\n# ========================================\n#             Gradio Setting\n# ========================================\n\ndef gradio_reset(chat_state, img_list):\n    if chat_state is not None:\n        chat_state.messages = []\n    if img_list is not None:\n        img_list = []\n    return None, gr.update(value=None, interactive=True), gr.update(placeholder='请先上传图片', interactive=False),gr.update(value=\"上传图片并开始咨询\", interactive=True), chat_state, img_list\n\ndef upload_img(gr_img, text_input, chat_state):\n    if gr_img is None:\n        return None, None, gr.update(interactive=True), chat_state, None\n    chat_state = CONV_ZH.copy()\n    img_list = []\n    llm_message = chat.upload_img(gr_img, chat_state, img_list)\n    return gr.update(interactive=False), gr.update(interactive=True, placeholder='输入问题'), gr.update(value=\"开始对话\", interactive=False), chat_state, img_list\n\ndef gradio_ask(user_message, chatbot, chat_state):\n    if len(user_message) == 0:\n        return gr.update(interactive=True, placeholder='Input should not be empty!'), chatbot, chat_state\n    chat.ask(user_message, chat_state)\n    chatbot = chatbot + [[user_message, None]]\n    return '', chatbot, chat_state\n\n\ndef gradio_answer(chatbot, chat_state, img_list, num_beams, temperature):\n    llm_message = chat.answer(conv=chat_state,\n                              img_list=img_list,\n                              num_beams=num_beams,\n                              temperature=temperature,\n                              max_new_tokens=300,\n                              max_length=2000)[0]\n    chatbot[-1][1] = llm_message\n    return chatbot, chat_state, img_list\n\ntitle = \"\"\"<h1 align=\"center\"> XrayPULSE </h1>\"\"\"\ndescription = \"\"\"<h3>上传X光影像，开始诊断咨询</h3>\"\"\"\ndisclaimer = \"\"\" \n            <h1 >使用说明:</h1>\n            <ul> \n                <li>XrayPULSE为PULSE在医疗多模态领域的扩展应用之一，可以用于对X光影像进行医学诊断分析，辅助医生，并为患者提供诊断支持。</li>\n                <li>XrayPULSE尝试通过分析X光影像提供准确和有用的结果。然而，我们对所提供结果的有效性、可靠性或完整性不作任何明确的保证或陈述。我们需要不断改善和完善服务，为医疗专业人员提供最好的协助</li>\n            </ul>\n            <hr> \n            <h3 align=\"center\">OpenMedLab</h3>\n\n            \"\"\"\n\ndef set_example_xray(example: list) -> dict:\n    return gr.Image.update(value=example[0])\n\n\ndef set_example_text_input(example_text: str) -> dict:\n    return gr.Textbox.update(value=example_text[0])\n\n#TODO show examples below\n\nwith gr.Blocks() as demo:\n    gr.Markdown(title)\n    gr.Markdown(description)\n\n    with gr.Row():\n        with gr.Column(scale=0.5):\n            image = gr.Image(type=\"pil\")\n            upload_button = gr.Button(value=\"上传影像并开始咨询\", interactive=True, variant=\"primary\")\n            clear = gr.Button(\"重制\")\n            \n            num_beams = gr.Slider(\n                minimum=1,\n                maximum=10,\n                value=1,\n                step=1,\n                interactive=True,\n                label=\"beam search numbers\",\n            )\n            \n            temperature = gr.Slider(\n                minimum=0.1,\n                maximum=2.0,\n                value=1.0,\n                step=0.1,\n                interactive=True,\n                label=\"Temperature\",\n            )\n\n        with gr.Column():\n            chat_state = gr.State()\n            img_list = gr.State()\n            chatbot = gr.Chatbot(label='XrayPULSE')\n            text_input = gr.Textbox(label='用户', placeholder='请上传X光影像', interactive=False)\n\n\n    with gr.Row():\n        example_xrays = gr.Dataset(components=[image], label=\"X光影像范例\",\n                                    samples=[\n                                        [os.path.join(os.path.dirname(__file__), \"images/image1.png\")],\n                                        [os.path.join(os.path.dirname(__file__), \"images/image2.png\")],\n                                        [os.path.join(os.path.dirname(__file__), \"images/image3.png\")],\n                                        [os.path.join(os.path.dirname(__file__), \"images/image4.png\")],\n                                        [os.path.join(os.path.dirname(__file__), \"images/image5.png\")],\n                                        [os.path.join(os.path.dirname(__file__), \"images/image6.png\")],\n                                    ])\n        \n\n    with gr.Row():\n        example_texts = gr.Dataset(components=[gr.Textbox(visible=False)],\n                                    label=\"咨询问题范例\",\n                                    samples=[\n                                        [\"详细描述所给的胸部X光影像。\"],\n                                        [\"请观察这张胸部X光影像，并阐述你的发现和总结。\"],\n                                        [\"你能否对所给的胸部X光影像进行详细的描述？\"],\n                                        [\"尽可能详细地描述所给的胸部X光影像。\"],\n                                        [\"这张胸部X光影像中的关键症状是什么？\"],\n                                        [\"你能在这张胸部X光影像中，指出存在的任何异常或需要注意的地方吗\"],\n                                        [\"这张胸部X光影像中，有哪些肺部和心脏的具体特征可见？\"],\n                                        [\"在这张胸部X光影像中，最显著的特征是什么，它是如何反映出病人的健康状况？\"],\n                                        [\"根据从这张胸部X光影像中观察到的发现，给出影像的总体印象是正常还是异常？\"],\n                                    ],)\n    \n    example_xrays.click(fn=set_example_xray, inputs=example_xrays, outputs=example_xrays.components)\n\n    upload_button.click(upload_img, [image, text_input, chat_state], [image, text_input, upload_button, chat_state, img_list])\n    \n    click_response = example_texts.click(set_example_text_input, inputs=example_texts, outputs=text_input).then(\n        gradio_ask, [text_input, chatbot, chat_state], [text_input, chatbot, chat_state], queue=False)\n    click_response.then(\n        gradio_answer, [chatbot, chat_state, img_list, num_beams, temperature], [chatbot, chat_state, img_list], queue=False\n    )\n    \n    submit_response = text_input.submit(gradio_ask, [text_input, chatbot, chat_state], [text_input, chatbot, chat_state], queue=False)\n    submit_response.then(\n        gradio_answer, [chatbot, chat_state, img_list, num_beams, temperature], [chatbot, chat_state, img_list], queue=False\n    )\n    clear.click(gradio_reset, [chat_state, img_list], [chatbot, image, text_input, upload_button, chat_state, img_list], queue=False)\n    \n    gr.Markdown(disclaimer)\ndemo.launch(share=True, enable_queue=True)\n"
  },
  {
    "path": "demo_configs/xraypulse_demo.yaml",
    "content": "model:\n  arch: xray_pulse\n  model_type: pulse\n  freeze_vit: True\n  freeze_qformer: True\n  max_txt_len: 160\n  end_sym: \"</s>\"\n  low_resource: True\n  prompt_path: \"prompts/alignment.txt\"\n  prompt_template: 'Instructions: You are PULSE, a large language model trained by SHAIlab. Answer as concisely as possible.\\nKnowledge cutoff: 2021-09-01\\nCurrent date: 2022-02-01</s> User: {} </s> Helper: '\n  ckpt: './XrayPULSE_ckpt.pth'\n\ndatasets:\n  openi:\n    vis_processor:\n      train:\n        name: \"blip2_image_eval\"\n        image_size: 224\n    text_processor:\n      train:\n        name: \"blip_caption\"\n\nrun:\n  task: image_text_pretrain\n"
  },
  {
    "path": "env.yml",
    "content": "name: xraypulse\nchannels:\n  - pytorch\n  - defaults\ndependencies:\n  - _libgcc_mutex=0.1=main\n  - _openmp_mutex=5.1=1_gnu\n  - blas=1.0=mkl\n  - brotlipy=0.7.0=py39h27cfd23_1003\n  - bzip2=1.0.8=h7b6447c_0\n  - ca-certificates=2023.01.10=h06a4308_0\n  - certifi=2022.12.7=py39h06a4308_0\n  - cffi=1.15.1=py39h5eee18b_3\n  - charset-normalizer=2.0.4=pyhd3eb1b0_0\n  - cryptography=39.0.1=py39h9ce1e76_0\n  - cudatoolkit=11.3.1=h2bc3f7f_2\n  - ffmpeg=4.3=hf484d3e_0\n  - flit-core=3.8.0=py39h06a4308_0\n  - freetype=2.12.1=h4a9f257_0\n  - giflib=5.2.1=h5eee18b_3\n  - gmp=6.2.1=h295c915_3\n  - gnutls=3.6.15=he1e5248_0\n  - intel-openmp=2021.4.0=h06a4308_3561\n  - jpeg=9e=h5eee18b_1\n  - lame=3.100=h7b6447c_0\n  - lcms2=2.12=h3be6417_0\n  - ld_impl_linux-64=2.38=h1181459_1\n  - lerc=3.0=h295c915_0\n  - libdeflate=1.17=h5eee18b_0\n  - libffi=3.4.2=h6a678d5_6\n  - libgcc-ng=11.2.0=h1234567_1\n  - libgomp=11.2.0=h1234567_1\n  - libiconv=1.16=h7f8727e_2\n  - libidn2=2.3.2=h7f8727e_0\n  - libpng=1.6.39=h5eee18b_0\n  - libstdcxx-ng=11.2.0=h1234567_1\n  - libtasn1=4.19.0=h5eee18b_0\n  - libtiff=4.5.0=h6a678d5_2\n  - libunistring=0.9.10=h27cfd23_0\n  - libwebp=1.2.4=h11a3e52_1\n  - libwebp-base=1.2.4=h5eee18b_1\n  - lz4-c=1.9.4=h6a678d5_0\n  - mkl=2021.4.0=h06a4308_640\n  - mkl-service=2.4.0=py39h7f8727e_0\n  - mkl_fft=1.3.1=py39hd3c417c_0\n  - mkl_random=1.2.2=py39h51133e4_0\n  - ncurses=6.4=h6a678d5_0\n  - nettle=3.7.3=hbbd107a_1\n  - numpy=1.23.5=py39h14f4228_0\n  - numpy-base=1.23.5=py39h31eccc5_0\n  - openh264=2.1.1=h4ff587b_0\n  - openssl=1.1.1t=h7f8727e_0\n  - pillow=9.4.0=py39h6a678d5_0\n  - pip=23.0.1=py39h06a4308_0\n  - pycparser=2.21=pyhd3eb1b0_0\n  - pyopenssl=23.0.0=py39h06a4308_0\n  - pysocks=1.7.1=py39h06a4308_0\n  - python=3.9.16=h7a1cb2a_2\n  - pytorch-mutex=1.0=cuda\n  - readline=8.2=h5eee18b_0\n  - requests=2.28.1=py39h06a4308_1\n  - setuptools=66.0.0=py39h06a4308_0\n  - six=1.16.0=pyhd3eb1b0_1\n  - sqlite=3.41.2=h5eee18b_0\n  - tk=8.6.12=h1ccaba5_0\n  - torchaudio=0.12.1=py39_cu113\n  - torchvision=0.13.1=py39_cu113\n  - typing_extensions=4.4.0=py39h06a4308_0\n  - urllib3=1.26.15=py39h06a4308_0\n  - wheel=0.38.4=py39h06a4308_0\n  - xz=5.2.10=h5eee18b_1\n  - zlib=1.2.13=h5eee18b_0\n  - zstd=1.5.5=hc292b87_0\n  - pip:\n      - accelerate==0.15.0\n      - aiofiles==23.1.0\n      - aiohttp==3.8.4\n      - aiosignal==1.3.1\n      - albumentations==1.3.0\n      - altair==4.2.2\n      - antlr4-python3-runtime==4.9.3\n      - anyio==3.6.2\n      - appdirs==1.4.4\n      - argon2-cffi==21.3.0\n      - argon2-cffi-bindings==21.2.0\n      - arrow==1.2.3\n      - asttokens==2.2.1\n      - async-timeout==4.0.2\n      - attrs==22.2.0\n      - backcall==0.2.0\n      - beautifulsoup4==4.12.2\n      - bitsandbytes==0.37.0\n      - bleach==6.0.0\n      - blis==0.7.9\n      - braceexpand==0.1.7\n      - cachetools==5.3.0\n      - catalogue==2.0.8\n      - cchardet==2.1.7\n      - chardet==3.0.4\n      - click==8.1.3\n      - cmake==3.26.3\n      - comm==0.1.3\n      - confection==0.0.4\n      - contourpy==1.0.7\n      - cycler==0.11.0\n      - cymem==2.0.7\n      - dataclasses==0.6\n      - datasets==2.12.0\n      - debugpy==1.6.7\n      - decorator==5.1.1\n      - decord==0.6.0\n      - defusedxml==0.7.1\n      - dill==0.3.6\n      - docker-pycreds==0.4.0\n      - entrypoints==0.4\n      - et-xmlfile==1.1.0\n      - evaluate==0.4.0\n      - executing==1.2.0\n      - exifread-nocycle==3.0.1\n      - fairscale==0.4.13\n      - fastapi==0.95.1\n      - fastchat==0.1\n      - fastjsonschema==2.16.3\n      - ffmpy==0.3.0\n      - filelock==3.9.0\n      - fire==0.5.0\n      - fonttools==4.38.0\n      - fqdn==1.5.1\n      - frozenlist==1.3.3\n      - fschat==0.2.3\n      - fsspec==2022.11.0\n      - gensim==4.3.1\n      - gitdb==4.0.10\n      - gitpython==3.1.31\n      - googletrans==3.0.0\n      - gradio==3.23.0\n      - gradio-client==0.0.8\n      - h11==0.9.0\n      - h2==3.2.0\n      - hiq-python==1.1.12\n      - hpack==3.0.0\n      - hstspreload==2023.1.1\n      - httpcore==0.9.1\n      - httpx==0.13.3\n      - huggingface-hub==0.13.4\n      - hyperframe==5.2.0\n      - idna==2.10\n      - imageio==2.27.0\n      - img2dataset==1.25.4\n      - importlib-metadata==6.5.0\n      - importlib-resources==5.12.0\n      - iopath==0.1.10\n      - ipykernel==6.22.0\n      - ipython==8.12.0\n      - ipython-genutils==0.2.0\n      - isoduration==20.11.0\n      - jedi==0.18.2\n      - jinja2==3.1.2\n      - joblib==1.2.0\n      - jsonpointer==2.3\n      - jsonschema==4.17.3\n      - jupyter-client==8.2.0\n      - jupyter-core==5.3.0\n      - jupyter-events==0.6.3\n      - jupyter-server==2.5.0\n      - jupyter-server-terminals==0.4.4\n      - jupyterlab-pygments==0.2.2\n      - kiwisolver==1.4.4\n      - langcodes==3.3.0\n      - lazy-loader==0.2\n      - linkify-it-py==2.0.0\n      - lit==16.0.1\n      - llvmlite==0.39.1\n      - markdown-it-py==2.2.0\n      - markdown2==2.4.8\n      - markupsafe==2.1.2\n      - matplotlib==3.7.0\n      - matplotlib-inline==0.1.6\n      - mdit-py-plugins==0.3.3\n      - mdurl==0.1.2\n      - medclip==0.0.3\n      - mistune==2.0.5\n      - mpmath==1.3.0\n      - multidict==6.0.4\n      - multiprocess==0.70.14\n      - murmurhash==1.0.9\n      - nbclassic==0.5.5\n      - nbclient==0.7.3\n      - nbconvert==7.3.1\n      - nbformat==5.8.0\n      - nest-asyncio==1.5.6\n      - networkx==3.1\n      - nltk==3.8.1\n      - notebook==6.5.4\n      - notebook-shim==0.2.2\n      - numba==0.56.4\n      - nvidia-cublas-cu11==11.10.3.66\n      - nvidia-cuda-cupti-cu11==11.7.101\n      - nvidia-cuda-nvrtc-cu11==11.7.99\n      - nvidia-cuda-runtime-cu11==11.7.99\n      - nvidia-cudnn-cu11==8.5.0.96\n      - nvidia-cufft-cu11==10.9.0.58\n      - nvidia-curand-cu11==10.2.10.91\n      - nvidia-cusolver-cu11==11.4.0.1\n      - nvidia-cusparse-cu11==11.7.4.91\n      - nvidia-nccl-cu11==2.14.3\n      - nvidia-nvtx-cu11==11.7.91\n      - omegaconf==2.3.0\n      - openai==0.27.0\n      - opencv-python==4.7.0.72\n      - opencv-python-headless==4.7.0.72\n      - openpyxl==3.1.2\n      - orjson==3.8.10\n      - packaging==23.0\n      - pandas==1.5.3\n      - pandocfilters==1.5.0\n      - parso==0.8.3\n      - pathtools==0.1.2\n      - pathy==0.10.1\n      - peft==0.2.0\n      - pexpect==4.8.0\n      - pickleshare==0.7.5\n      - platformdirs==3.2.0\n      - portalocker==2.7.0\n      - preshed==3.0.8\n      - prometheus-client==0.16.0\n      - promise==2.3\n      - prompt-toolkit==3.0.38\n      - protobuf==3.20.3\n      - psutil==5.9.4\n      - ptyprocess==0.7.0\n      - pure-eval==0.2.2\n      - py-itree==0.0.19\n      - pyarrow==12.0.1\n      - pydantic==1.10.7\n      - pydub==0.25.1\n      - pygments==2.15.1\n      - pyllama==0.0.9\n      - pynndescent==0.5.9\n      - pyparsing==3.0.9\n      - pyrsistent==0.19.3\n      - python-dateutil==2.8.2\n      - python-json-logger==2.0.7\n      - python-multipart==0.0.6\n      - pytz==2023.3\n      - pywavelets==1.4.1\n      - pyyaml==6.0\n      - pyzmq==25.0.2\n      - qudida==0.0.4\n      - regex==2022.10.31\n      - responses==0.18.0\n      - rfc3339-validator==0.1.4\n      - rfc3986==1.5.0\n      - rfc3986-validator==0.1.1\n      - rich==13.3.4\n      - scikit-image==0.20.0\n      - scikit-learn==1.2.2\n      - scipy==1.9.1\n      - semantic-version==2.10.0\n      - send2trash==1.8.0\n      - sentence-transformers==2.2.2\n      - sentencepiece==0.1.97\n      - sentry-sdk==1.19.1\n      - setproctitle==1.3.2\n      - shortuuid==1.0.11\n      - smart-open==6.3.0\n      - smmap==5.0.0\n      - sniffio==1.3.0\n      - soupsieve==2.4.1\n      - spacy==3.5.1\n      - spacy-legacy==3.0.12\n      - spacy-loggers==1.0.4\n      - srsly==2.4.6\n      - stack-data==0.6.2\n      - starlette==0.26.1\n      - svgwrite==1.4.3\n      - sympy==1.11.1\n      - tenacity==8.2.2\n      - termcolor==2.2.0\n      - terminado==0.17.1\n      - textaugment==1.3.4\n      - textblob==0.17.1\n      - thinc==8.1.9\n      - threadpoolctl==3.1.0\n      - tifffile==2023.4.12\n      - timm==0.6.13\n      - tinycss2==1.2.1\n      - tokenizers==0.13.2\n      - toolz==0.12.0\n      - torch==2.0.0\n      - tornado==6.3\n      - tqdm==4.64.1\n      - traitlets==5.9.0\n      - transformers==4.29.0\n      - triton==2.0.0\n      - typer==0.7.0\n      - tzdata==2023.3\n      - uc-micro-py==1.0.1\n      - umap-learn==0.5.3\n      - uri-template==1.2.0\n      - uvicorn==0.21.1\n      - wandb==0.12.21\n      - wasabi==1.1.1\n      - wavedrom==2.0.3.post3\n      - wcwidth==0.2.6\n      - webcolors==1.13\n      - webdataset==0.2.48\n      - webencodings==0.5.1\n      - websocket-client==1.5.1\n      - websockets==11.0.2\n      - wget==3.2\n      - xxhash==3.2.0\n      - yarl==1.8.2\n      - zipp==3.14.0\n"
  },
  {
    "path": "prompts/alignment.txt",
    "content": "<Img><图片></Img> 详细描述所给的胸部X光影像。\n<Img><图片></Img> 请观察这张胸部X光影像，并阐述你的发现和总结。\n<Img><图片></Img> 你能否对所给的胸部X光影像进行详细的描述？\n<Img><图片></Img> 尽可能详细地描述所给的胸部X光影像。\n<Img><图片></Img> 这张胸部X光影像中的关键症状是什么？\n<Img><图片></Img> 你能在这张胸部X光影像中，指出存在的任何异常或需要注意的地方吗？\n<Img><图片></Img> 这张胸部X光影像中，有哪些肺部和心脏的具体特征可见？\n<Img><图片></Img> 在这张胸部X光影像中，最显著的特征是什么，它是如何反映出病人的健康状况？\n<Img><图片></Img> 这张胸部X光影像提供了哪些观察发现和总体印象？\n<Img><图片></Img> 这张胸部X光影像中，心脏的大小和形状如何？\n<Img><图片></Img> 根据从这张胸部X光影像中观察到的发现，给出影像的总体印象是正常还是异常？\n<Img><图片></Img> 在这张胸部X光影像中，有无感染或炎症的迹象？如果有，可能的原因是什么？\n<Img><图片></Img> 根据这张胸部X光影像中的发现，请你给出总体印象。\n<Img><图片></Img> 在这张胸部X光影像中，有没有患者淋巴结肿大或异常的可见迹象\n<Img><图片></Img> 这张胸部X光影像中观察到的异常有没有可能引发的并发症或风险？或者说，这张X光影像所展示的患者是正常的吗"
  },
  {
    "path": "run_demo.sh",
    "content": "CUDA_VISIBLE_DEVICES=0 python -u demo.py --cfg-path demo_configs/xraypulse_demo.yaml  --gpu-id 0"
  },
  {
    "path": "xraypulse/__init__.py",
    "content": "\"\"\"\n Copyright (c) 2022, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\n\nimport os\nimport sys\n\nfrom omegaconf import OmegaConf\n\nfrom xraypulse.common.registry import registry\n\nfrom xraypulse.datasets.builders import *\nfrom xraypulse.models import *\nfrom xraypulse.processors import *\nfrom xraypulse.tasks import *\n\n\nroot_dir = os.path.dirname(os.path.abspath(__file__))\ndefault_cfg = OmegaConf.load(os.path.join(root_dir, \"configs/default.yaml\"))\n\nregistry.register_path(\"library_root\", root_dir)\nrepo_root = os.path.join(root_dir, \"..\")\nregistry.register_path(\"repo_root\", repo_root)\ncache_root = os.path.join(repo_root, default_cfg.env.cache_root)\nregistry.register_path(\"cache_root\", cache_root)\n\nregistry.register(\"MAX_INT\", sys.maxsize)\nregistry.register(\"SPLIT_NAMES\", [\"train\", \"val\", \"test\"])\n"
  },
  {
    "path": "xraypulse/common/__init__.py",
    "content": ""
  },
  {
    "path": "xraypulse/common/config.py",
    "content": "\"\"\"\n Copyright (c) 2022, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\n\nimport logging\nimport json\nfrom typing import Dict\n\nfrom omegaconf import OmegaConf\nfrom xraypulse.common.registry import registry\n\n\nclass Config:\n    def __init__(self, args):\n        self.config = {}\n\n        self.args = args\n\n        # Register the config and configuration for setup\n        registry.register(\"configuration\", self)\n\n        user_config = self._build_opt_list(self.args.options)\n\n        config = OmegaConf.load(self.args.cfg_path)\n\n        runner_config = self.build_runner_config(config)\n        model_config = self.build_model_config(config, **user_config)\n        dataset_config = self.build_dataset_config(config)\n\n        # Validate the user-provided runner configuration\n        # model and dataset configuration are supposed to be validated by the respective classes\n        # [TODO] validate the model/dataset configuration\n        # self._validate_runner_config(runner_config)\n\n        # Override the default configuration with user options.\n        self.config = OmegaConf.merge(\n            runner_config, model_config, dataset_config, user_config\n        )\n\n    def _validate_runner_config(self, runner_config):\n        \"\"\"\n        This method validates the configuration, such that\n            1) all the user specified options are valid;\n            2) no type mismatches between the user specified options and the config.\n        \"\"\"\n        runner_config_validator = create_runner_config_validator()\n        runner_config_validator.validate(runner_config)\n\n    def _build_opt_list(self, opts):\n        opts_dot_list = self._convert_to_dot_list(opts)\n        return OmegaConf.from_dotlist(opts_dot_list)\n\n    @staticmethod\n    def build_model_config(config, **kwargs):\n        model = config.get(\"model\", None)\n        assert model is not None, \"Missing model configuration file.\"\n\n        model_cls = registry.get_model_class(model.arch)\n        assert model_cls is not None, f\"Model '{model.arch}' has not been registered.\"\n\n        model_type = kwargs.get(\"model.model_type\", None)\n        if not model_type:\n            model_type = model.get(\"model_type\", None)\n        # else use the model type selected by user.\n\n        assert model_type is not None, \"Missing model_type.\"\n\n        model_config_path = model_cls.default_config_path(model_type=model_type)\n\n        model_config = OmegaConf.create()\n        # hierarchy override, customized config > default config\n        model_config = OmegaConf.merge(\n            model_config,\n            OmegaConf.load(model_config_path),\n            {\"model\": config[\"model\"]},\n        )\n\n        return model_config\n\n    @staticmethod\n    def build_runner_config(config):\n        return {\"run\": config.run}\n\n    @staticmethod\n    def build_dataset_config(config):\n        datasets = config.get(\"datasets\", None)\n        if datasets is None:\n            raise KeyError(\n                \"Expecting 'datasets' as the root key for dataset configuration.\"\n            )\n\n        dataset_config = OmegaConf.create()\n\n        for dataset_name in datasets:\n            builder_cls = registry.get_builder_class(dataset_name)\n\n            dataset_config_type = datasets[dataset_name].get(\"type\", \"default\")\n            dataset_config_path = builder_cls.default_config_path(\n                type=dataset_config_type\n            )\n\n            # hierarchy override, customized config > default config\n            dataset_config = OmegaConf.merge(\n                dataset_config,\n                OmegaConf.load(dataset_config_path),\n                {\"datasets\": {dataset_name: config[\"datasets\"][dataset_name]}},\n            )\n\n        return dataset_config\n\n    def _convert_to_dot_list(self, opts):\n        if opts is None:\n            opts = []\n\n        if len(opts) == 0:\n            return opts\n\n        has_equal = opts[0].find(\"=\") != -1\n\n        if has_equal:\n            return opts\n\n        return [(opt + \"=\" + value) for opt, value in zip(opts[0::2], opts[1::2])]\n\n    def get_config(self):\n        return self.config\n\n    @property\n    def run_cfg(self):\n        return self.config.run\n\n    @property\n    def datasets_cfg(self):\n        return self.config.datasets\n\n    @property\n    def model_cfg(self):\n        return self.config.model\n\n    def pretty_print(self):\n        logging.info(\"\\n=====  Running Parameters    =====\")\n        logging.info(self._convert_node_to_json(self.config.run))\n\n        logging.info(\"\\n======  Dataset Attributes  ======\")\n        datasets = self.config.datasets\n\n        for dataset in datasets:\n            if dataset in self.config.datasets:\n                logging.info(f\"\\n======== {dataset} =======\")\n                dataset_config = self.config.datasets[dataset]\n                logging.info(self._convert_node_to_json(dataset_config))\n            else:\n                logging.warning(f\"No dataset named '{dataset}' in config. Skipping\")\n\n        logging.info(f\"\\n======  Model Attributes  ======\")\n        logging.info(self._convert_node_to_json(self.config.model))\n\n    def _convert_node_to_json(self, node):\n        container = OmegaConf.to_container(node, resolve=True)\n        return json.dumps(container, indent=4, sort_keys=True)\n\n    def to_dict(self):\n        return OmegaConf.to_container(self.config)\n\n\ndef node_to_dict(node):\n    return OmegaConf.to_container(node)\n\n\nclass ConfigValidator:\n    \"\"\"\n    This is a preliminary implementation to centralize and validate the configuration.\n    May be altered in the future.\n\n    A helper class to validate configurations from yaml file.\n\n    This serves the following purposes:\n        1. Ensure all the options in the yaml are defined, raise error if not.\n        2. when type mismatches are found, the validator will raise an error.\n        3. a central place to store and display helpful messages for supported configurations.\n\n    \"\"\"\n\n    class _Argument:\n        def __init__(self, name, choices=None, type=None, help=None):\n            self.name = name\n            self.val = None\n            self.choices = choices\n            self.type = type\n            self.help = help\n\n        def __str__(self):\n            s = f\"{self.name}={self.val}\"\n            if self.type is not None:\n                s += f\", ({self.type})\"\n            if self.choices is not None:\n                s += f\", choices: {self.choices}\"\n            if self.help is not None:\n                s += f\", ({self.help})\"\n            return s\n\n    def __init__(self, description):\n        self.description = description\n\n        self.arguments = dict()\n\n        self.parsed_args = None\n\n    def __getitem__(self, key):\n        assert self.parsed_args is not None, \"No arguments parsed yet.\"\n\n        return self.parsed_args[key]\n\n    def __str__(self) -> str:\n        return self.format_help()\n\n    def add_argument(self, *args, **kwargs):\n        \"\"\"\n        Assume the first argument is the name of the argument.\n        \"\"\"\n        self.arguments[args[0]] = self._Argument(*args, **kwargs)\n\n    def validate(self, config=None):\n        \"\"\"\n        Convert yaml config (dict-like) to list, required by argparse.\n        \"\"\"\n        for k, v in config.items():\n            assert (\n                k in self.arguments\n            ), f\"\"\"{k} is not a valid argument. Support arguments are {self.format_arguments()}.\"\"\"\n\n            if self.arguments[k].type is not None:\n                try:\n                    self.arguments[k].val = self.arguments[k].type(v)\n                except ValueError:\n                    raise ValueError(f\"{k} is not a valid {self.arguments[k].type}.\")\n\n            if self.arguments[k].choices is not None:\n                assert (\n                    v in self.arguments[k].choices\n                ), f\"\"\"{k} must be one of {self.arguments[k].choices}.\"\"\"\n\n        return config\n\n    def format_arguments(self):\n        return str([f\"{k}\" for k in sorted(self.arguments.keys())])\n\n    def format_help(self):\n        # description + key-value pair string for each argument\n        help_msg = str(self.description)\n        return help_msg + \", available arguments: \" + self.format_arguments()\n\n    def print_help(self):\n        # display help message\n        print(self.format_help())\n\n\ndef create_runner_config_validator():\n    validator = ConfigValidator(description=\"Runner configurations\")\n\n    validator.add_argument(\n        \"runner\",\n        type=str,\n        choices=[\"runner_base\", \"runner_iter\"],\n        help=\"\"\"Runner to use. The \"runner_base\" uses epoch-based training while iter-based\n            runner runs based on iters. Default: runner_base\"\"\",\n    )\n    # add argumetns for training dataset ratios\n    validator.add_argument(\n        \"train_dataset_ratios\",\n        type=Dict[str, float],\n        help=\"\"\"Ratios of training dataset. This is used in iteration-based runner.\n        Do not support for epoch-based runner because how to define an epoch becomes tricky.\n        Default: None\"\"\",\n    )\n    validator.add_argument(\n        \"max_iters\",\n        type=float,\n        help=\"Maximum number of iterations to run.\",\n    )\n    validator.add_argument(\n        \"max_epoch\",\n        type=int,\n        help=\"Maximum number of epochs to run.\",\n    )\n    # add arguments for iters_per_inner_epoch\n    validator.add_argument(\n        \"iters_per_inner_epoch\",\n        type=float,\n        help=\"Number of iterations per inner epoch. This is required when runner is runner_iter.\",\n    )\n    lr_scheds_choices = registry.list_lr_schedulers()\n    validator.add_argument(\n        \"lr_sched\",\n        type=str,\n        choices=lr_scheds_choices,\n        help=\"Learning rate scheduler to use, from {}\".format(lr_scheds_choices),\n    )\n    task_choices = registry.list_tasks()\n    validator.add_argument(\n        \"task\",\n        type=str,\n        choices=task_choices,\n        help=\"Task to use, from {}\".format(task_choices),\n    )\n    # add arguments for init_lr\n    validator.add_argument(\n        \"init_lr\",\n        type=float,\n        help=\"Initial learning rate. This will be the learning rate after warmup and before decay.\",\n    )\n    # add arguments for min_lr\n    validator.add_argument(\n        \"min_lr\",\n        type=float,\n        help=\"Minimum learning rate (after decay).\",\n    )\n    # add arguments for warmup_lr\n    validator.add_argument(\n        \"warmup_lr\",\n        type=float,\n        help=\"Starting learning rate for warmup.\",\n    )\n    # add arguments for learning rate decay rate\n    validator.add_argument(\n        \"lr_decay_rate\",\n        type=float,\n        help=\"Learning rate decay rate. Required if using a decaying learning rate scheduler.\",\n    )\n    # add arguments for weight decay\n    validator.add_argument(\n        \"weight_decay\",\n        type=float,\n        help=\"Weight decay rate.\",\n    )\n    # add arguments for training batch size\n    validator.add_argument(\n        \"batch_size_train\",\n        type=int,\n        help=\"Training batch size.\",\n    )\n    # add arguments for evaluation batch size\n    validator.add_argument(\n        \"batch_size_eval\",\n        type=int,\n        help=\"Evaluation batch size, including validation and testing.\",\n    )\n    # add arguments for number of workers for data loading\n    validator.add_argument(\n        \"num_workers\",\n        help=\"Number of workers for data loading.\",\n    )\n    # add arguments for warm up steps\n    validator.add_argument(\n        \"warmup_steps\",\n        type=int,\n        help=\"Number of warmup steps. Required if a warmup schedule is used.\",\n    )\n    # add arguments for random seed\n    validator.add_argument(\n        \"seed\",\n        type=int,\n        help=\"Random seed.\",\n    )\n    # add arguments for output directory\n    validator.add_argument(\n        \"output_dir\",\n        type=str,\n        help=\"Output directory to save checkpoints and logs.\",\n    )\n    # add arguments for whether only use evaluation\n    validator.add_argument(\n        \"evaluate\",\n        help=\"Whether to only evaluate the model. If true, training will not be performed.\",\n    )\n    # add arguments for splits used for training, e.g. [\"train\", \"val\"]\n    validator.add_argument(\n        \"train_splits\",\n        type=list,\n        help=\"Splits to use for training.\",\n    )\n    # add arguments for splits used for validation, e.g. [\"val\"]\n    validator.add_argument(\n        \"valid_splits\",\n        type=list,\n        help=\"Splits to use for validation. If not provided, will skip the validation.\",\n    )\n    # add arguments for splits used for testing, e.g. [\"test\"]\n    validator.add_argument(\n        \"test_splits\",\n        type=list,\n        help=\"Splits to use for testing. If not provided, will skip the testing.\",\n    )\n    # add arguments for accumulating gradient for iterations\n    validator.add_argument(\n        \"accum_grad_iters\",\n        type=int,\n        help=\"Number of iterations to accumulate gradient for.\",\n    )\n\n    # ====== distributed training ======\n    validator.add_argument(\n        \"device\",\n        type=str,\n        choices=[\"cpu\", \"cuda\"],\n        help=\"Device to use. Support 'cuda' or 'cpu' as for now.\",\n    )\n    validator.add_argument(\n        \"world_size\",\n        type=int,\n        help=\"Number of processes participating in the job.\",\n    )\n    validator.add_argument(\"dist_url\", type=str)\n    validator.add_argument(\"distributed\", type=bool)\n    # add arguments to opt using distributed sampler during evaluation or not\n    validator.add_argument(\n        \"use_dist_eval_sampler\",\n        type=bool,\n        help=\"Whether to use distributed sampler during evaluation or not.\",\n    )\n\n    # ====== task specific ======\n    # generation task specific arguments\n    # add arguments for maximal length of text output\n    validator.add_argument(\n        \"max_len\",\n        type=int,\n        help=\"Maximal length of text output.\",\n    )\n    # add arguments for minimal length of text output\n    validator.add_argument(\n        \"min_len\",\n        type=int,\n        help=\"Minimal length of text output.\",\n    )\n    # add arguments number of beams\n    validator.add_argument(\n        \"num_beams\",\n        type=int,\n        help=\"Number of beams used for beam search.\",\n    )\n\n    # vqa task specific arguments\n    # add arguments for number of answer candidates\n    validator.add_argument(\n        \"num_ans_candidates\",\n        type=int,\n        help=\"\"\"For ALBEF and BLIP, these models first rank answers according to likelihood to select answer candidates.\"\"\",\n    )\n    # add arguments for inference method\n    validator.add_argument(\n        \"inference_method\",\n        type=str,\n        choices=[\"genearte\", \"rank\"],\n        help=\"\"\"Inference method to use for question answering. If rank, requires a answer list.\"\"\",\n    )\n\n    # ====== model specific ======\n    validator.add_argument(\n        \"k_test\",\n        type=int,\n        help=\"Number of top k most similar samples from ITC/VTC selection to be tested.\",\n    )\n\n    return validator\n"
  },
  {
    "path": "xraypulse/common/dist_utils.py",
    "content": "\"\"\"\n Copyright (c) 2022, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\n\nimport datetime\nimport functools\nimport os\n\nimport torch\nimport torch.distributed as dist\nimport timm.models.hub as timm_hub\n\n\ndef setup_for_distributed(is_master):\n    \"\"\"\n    This function disables printing when not in master process\n    \"\"\"\n    import builtins as __builtin__\n\n    builtin_print = __builtin__.print\n\n    def print(*args, **kwargs):\n        force = kwargs.pop(\"force\", False)\n        if is_master or force:\n            builtin_print(*args, **kwargs)\n\n    __builtin__.print = print\n\n\ndef is_dist_avail_and_initialized():\n    if not dist.is_available():\n        return False\n    if not dist.is_initialized():\n        return False\n    return True\n\n\ndef get_world_size():\n    if not is_dist_avail_and_initialized():\n        return 1\n    return dist.get_world_size()\n\n\ndef get_rank():\n    if not is_dist_avail_and_initialized():\n        return 0\n    return dist.get_rank()\n\n\ndef is_main_process():\n    return get_rank() == 0\n\n\ndef init_distributed_mode(args):\n    if \"RANK\" in os.environ and \"WORLD_SIZE\" in os.environ:\n        args.rank = int(os.environ[\"RANK\"])\n        args.world_size = int(os.environ[\"WORLD_SIZE\"])\n        args.gpu = int(os.environ[\"LOCAL_RANK\"])\n    elif \"SLURM_PROCID\" in os.environ:\n        args.rank = int(os.environ[\"SLURM_PROCID\"])\n        args.gpu = args.rank % torch.cuda.device_count()\n    else:\n        print(\"Not using distributed mode\")\n        args.distributed = False\n        return\n\n    args.distributed = True\n\n    torch.cuda.set_device(args.gpu)\n    args.dist_backend = \"nccl\"\n    print(\n        \"| distributed init (rank {}, world {}): {}\".format(\n            args.rank, args.world_size, args.dist_url\n        ),\n        flush=True,\n    )\n    torch.distributed.init_process_group(\n        backend=args.dist_backend,\n        init_method=args.dist_url,\n        world_size=args.world_size,\n        rank=args.rank,\n        timeout=datetime.timedelta(\n            days=365\n        ),  # allow auto-downloading and de-compressing\n    )\n    torch.distributed.barrier()\n    setup_for_distributed(args.rank == 0)\n\n\ndef get_dist_info():\n    if torch.__version__ < \"1.0\":\n        initialized = dist._initialized\n    else:\n        initialized = dist.is_initialized()\n    if initialized:\n        rank = dist.get_rank()\n        world_size = dist.get_world_size()\n    else:  # non-distributed training\n        rank = 0\n        world_size = 1\n    return rank, world_size\n\n\ndef main_process(func):\n    @functools.wraps(func)\n    def wrapper(*args, **kwargs):\n        rank, _ = get_dist_info()\n        if rank == 0:\n            return func(*args, **kwargs)\n\n    return wrapper\n\n\ndef download_cached_file(url, check_hash=True, progress=False):\n    \"\"\"\n    Download a file from a URL and cache it locally. If the file already exists, it is not downloaded again.\n    If distributed, only the main process downloads the file, and the other processes wait for the file to be downloaded.\n    \"\"\"\n\n    def get_cached_file_path():\n        # a hack to sync the file path across processes\n        parts = torch.hub.urlparse(url)\n        filename = os.path.basename(parts.path)\n        cached_file = os.path.join(timm_hub.get_cache_dir(), filename)\n\n        return cached_file\n\n    if is_main_process():\n        timm_hub.download_cached_file(url, check_hash, progress)\n\n    if is_dist_avail_and_initialized():\n        dist.barrier()\n\n    return get_cached_file_path()\n"
  },
  {
    "path": "xraypulse/common/gradcam.py",
    "content": "import numpy as np\nfrom matplotlib import pyplot as plt\nfrom scipy.ndimage import filters\nfrom skimage import transform as skimage_transform\n\n\ndef getAttMap(img, attMap, blur=True, overlap=True):\n    attMap -= attMap.min()\n    if attMap.max() > 0:\n        attMap /= attMap.max()\n    attMap = skimage_transform.resize(attMap, (img.shape[:2]), order=3, mode=\"constant\")\n    if blur:\n        attMap = filters.gaussian_filter(attMap, 0.02 * max(img.shape[:2]))\n        attMap -= attMap.min()\n        attMap /= attMap.max()\n    cmap = plt.get_cmap(\"jet\")\n    attMapV = cmap(attMap)\n    attMapV = np.delete(attMapV, 3, 2)\n    if overlap:\n        attMap = (\n            1 * (1 - attMap**0.7).reshape(attMap.shape + (1,)) * img\n            + (attMap**0.7).reshape(attMap.shape + (1,)) * attMapV\n        )\n    return attMap\n"
  },
  {
    "path": "xraypulse/common/logger.py",
    "content": "\"\"\"\n Copyright (c) 2022, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\n\nimport datetime\nimport logging\nimport time\nfrom collections import defaultdict, deque\n\nimport torch\nimport torch.distributed as dist\n\nfrom xraypulse.common import dist_utils\n\n\nclass SmoothedValue(object):\n    \"\"\"Track a series of values and provide access to smoothed values over a\n    window or the global series average.\n    \"\"\"\n\n    def __init__(self, window_size=20, fmt=None):\n        if fmt is None:\n            fmt = \"{median:.4f} ({global_avg:.4f})\"\n        self.deque = deque(maxlen=window_size)\n        self.total = 0.0\n        self.count = 0\n        self.fmt = fmt\n\n    def update(self, value, n=1):\n        self.deque.append(value)\n        self.count += n\n        self.total += value * n\n\n    def synchronize_between_processes(self):\n        \"\"\"\n        Warning: does not synchronize the deque!\n        \"\"\"\n        if not dist_utils.is_dist_avail_and_initialized():\n            return\n        t = torch.tensor([self.count, self.total], dtype=torch.float64, device=\"cuda\")\n        dist.barrier()\n        dist.all_reduce(t)\n        t = t.tolist()\n        self.count = int(t[0])\n        self.total = t[1]\n\n    @property\n    def median(self):\n        d = torch.tensor(list(self.deque))\n        return d.median().item()\n\n    @property\n    def avg(self):\n        d = torch.tensor(list(self.deque), dtype=torch.float32)\n        return d.mean().item()\n\n    @property\n    def global_avg(self):\n        return self.total / self.count\n\n    @property\n    def max(self):\n        return max(self.deque)\n\n    @property\n    def value(self):\n        return self.deque[-1]\n\n    def __str__(self):\n        return self.fmt.format(\n            median=self.median,\n            avg=self.avg,\n            global_avg=self.global_avg,\n            max=self.max,\n            value=self.value,\n        )\n\n\nclass MetricLogger(object):\n    def __init__(self, delimiter=\"\\t\"):\n        self.meters = defaultdict(SmoothedValue)\n        self.delimiter = delimiter\n\n    def update(self, **kwargs):\n        for k, v in kwargs.items():\n            if isinstance(v, torch.Tensor):\n                v = v.item()\n            assert isinstance(v, (float, int))\n            self.meters[k].update(v)\n\n    def __getattr__(self, attr):\n        if attr in self.meters:\n            return self.meters[attr]\n        if attr in self.__dict__:\n            return self.__dict__[attr]\n        raise AttributeError(\n            \"'{}' object has no attribute '{}'\".format(type(self).__name__, attr)\n        )\n\n    def __str__(self):\n        loss_str = []\n        for name, meter in self.meters.items():\n            loss_str.append(\"{}: {}\".format(name, str(meter)))\n        return self.delimiter.join(loss_str)\n\n    def global_avg(self):\n        loss_str = []\n        for name, meter in self.meters.items():\n            loss_str.append(\"{}: {:.4f}\".format(name, meter.global_avg))\n        return self.delimiter.join(loss_str)\n\n    def synchronize_between_processes(self):\n        for meter in self.meters.values():\n            meter.synchronize_between_processes()\n\n    def add_meter(self, name, meter):\n        self.meters[name] = meter\n\n    def log_every(self, iterable, print_freq, header=None):\n        i = 0\n        if not header:\n            header = \"\"\n        start_time = time.time()\n        end = time.time()\n        iter_time = SmoothedValue(fmt=\"{avg:.4f}\")\n        data_time = SmoothedValue(fmt=\"{avg:.4f}\")\n        space_fmt = \":\" + str(len(str(len(iterable)))) + \"d\"\n        log_msg = [\n            header,\n            \"[{0\" + space_fmt + \"}/{1}]\",\n            \"eta: {eta}\",\n            \"{meters}\",\n            \"time: {time}\",\n            \"data: {data}\",\n        ]\n        if torch.cuda.is_available():\n            log_msg.append(\"max mem: {memory:.0f}\")\n        log_msg = self.delimiter.join(log_msg)\n        MB = 1024.0 * 1024.0\n        for obj in iterable:\n            data_time.update(time.time() - end)\n            yield obj\n            iter_time.update(time.time() - end)\n            if i % print_freq == 0 or i == len(iterable) - 1:\n                eta_seconds = iter_time.global_avg * (len(iterable) - i)\n                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))\n                if torch.cuda.is_available():\n                    print(\n                        log_msg.format(\n                            i,\n                            len(iterable),\n                            eta=eta_string,\n                            meters=str(self),\n                            time=str(iter_time),\n                            data=str(data_time),\n                            memory=torch.cuda.max_memory_allocated() / MB,\n                        )\n                    )\n                else:\n                    print(\n                        log_msg.format(\n                            i,\n                            len(iterable),\n                            eta=eta_string,\n                            meters=str(self),\n                            time=str(iter_time),\n                            data=str(data_time),\n                        )\n                    )\n            i += 1\n            end = time.time()\n        total_time = time.time() - start_time\n        total_time_str = str(datetime.timedelta(seconds=int(total_time)))\n        print(\n            \"{} Total time: {} ({:.4f} s / it)\".format(\n                header, total_time_str, total_time / len(iterable)\n            )\n        )\n\n\nclass AttrDict(dict):\n    def __init__(self, *args, **kwargs):\n        super(AttrDict, self).__init__(*args, **kwargs)\n        self.__dict__ = self\n\n\ndef setup_logger():\n    logging.basicConfig(\n        level=logging.INFO if dist_utils.is_main_process() else logging.WARN,\n        format=\"%(asctime)s [%(levelname)s] %(message)s\",\n        handlers=[logging.StreamHandler()],\n    )\n"
  },
  {
    "path": "xraypulse/common/optims.py",
    "content": "\"\"\"\n Copyright (c) 2022, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\n\nimport math\n\nfrom xraypulse.common.registry import registry\n\n\n@registry.register_lr_scheduler(\"linear_warmup_step_lr\")\nclass LinearWarmupStepLRScheduler:\n    def __init__(\n        self,\n        optimizer,\n        max_epoch,\n        min_lr,\n        init_lr,\n        decay_rate=1,\n        warmup_start_lr=-1,\n        warmup_steps=0,\n        **kwargs\n    ):\n        self.optimizer = optimizer\n\n        self.max_epoch = max_epoch\n        self.min_lr = min_lr\n\n        self.decay_rate = decay_rate\n\n        self.init_lr = init_lr\n        self.warmup_steps = warmup_steps\n        self.warmup_start_lr = warmup_start_lr if warmup_start_lr >= 0 else init_lr\n\n    def step(self, cur_epoch, cur_step):\n        if cur_epoch == 0:\n            warmup_lr_schedule(\n                step=cur_step,\n                optimizer=self.optimizer,\n                max_step=self.warmup_steps,\n                init_lr=self.warmup_start_lr,\n                max_lr=self.init_lr,\n            )\n        else:\n            step_lr_schedule(\n                epoch=cur_epoch,\n                optimizer=self.optimizer,\n                init_lr=self.init_lr,\n                min_lr=self.min_lr,\n                decay_rate=self.decay_rate,\n            )\n\n\n@registry.register_lr_scheduler(\"linear_warmup_cosine_lr\")\nclass LinearWarmupCosineLRScheduler:\n    def __init__(\n        self,\n        optimizer,\n        max_epoch,\n        iters_per_epoch,\n        min_lr,\n        init_lr,\n        warmup_steps=0,\n        warmup_start_lr=-1,\n        **kwargs\n    ):\n        self.optimizer = optimizer\n\n        self.max_epoch = max_epoch\n        self.iters_per_epoch = iters_per_epoch\n        self.min_lr = min_lr\n\n        self.init_lr = init_lr\n        self.warmup_steps = warmup_steps\n        self.warmup_start_lr = warmup_start_lr if warmup_start_lr >= 0 else init_lr\n\n    def step(self, cur_epoch, cur_step):\n        total_cur_step = cur_epoch * self.iters_per_epoch + cur_step\n        if total_cur_step < self.warmup_steps:\n            warmup_lr_schedule(\n                step=cur_step,\n                optimizer=self.optimizer,\n                max_step=self.warmup_steps,\n                init_lr=self.warmup_start_lr,\n                max_lr=self.init_lr,\n            )\n        else:\n            cosine_lr_schedule(\n                epoch=total_cur_step,\n                optimizer=self.optimizer,\n                max_epoch=self.max_epoch * self.iters_per_epoch,\n                init_lr=self.init_lr,\n                min_lr=self.min_lr,\n            )\n\n\ndef cosine_lr_schedule(optimizer, epoch, max_epoch, init_lr, min_lr):\n    \"\"\"Decay the learning rate\"\"\"\n    lr = (init_lr - min_lr) * 0.5 * (\n        1.0 + math.cos(math.pi * epoch / max_epoch)\n    ) + min_lr\n    for param_group in optimizer.param_groups:\n        param_group[\"lr\"] = lr\n\n\ndef warmup_lr_schedule(optimizer, step, max_step, init_lr, max_lr):\n    \"\"\"Warmup the learning rate\"\"\"\n    lr = min(max_lr, init_lr + (max_lr - init_lr) * step / max(max_step, 1))\n    for param_group in optimizer.param_groups:\n        param_group[\"lr\"] = lr\n\n\ndef step_lr_schedule(optimizer, epoch, init_lr, min_lr, decay_rate):\n    \"\"\"Decay the learning rate\"\"\"\n    lr = max(min_lr, init_lr * (decay_rate**epoch))\n    for param_group in optimizer.param_groups:\n        param_group[\"lr\"] = lr\n"
  },
  {
    "path": "xraypulse/common/registry.py",
    "content": "\"\"\"\n Copyright (c) 2022, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\n\n\nclass Registry:\n    mapping = {\n        \"builder_name_mapping\": {},\n        \"task_name_mapping\": {},\n        \"processor_name_mapping\": {},\n        \"model_name_mapping\": {},\n        \"lr_scheduler_name_mapping\": {},\n        \"runner_name_mapping\": {},\n        \"state\": {},\n        \"paths\": {},\n    }\n\n    @classmethod\n    def register_builder(cls, name):\n        r\"\"\"Register a dataset builder to registry with key 'name'\n\n        Args:\n            name: Key with which the builder will be registered.\n\n        Usage:\n\n            from xraypulse.common.registry import registry\n            from xraypulse.datasets.base_dataset_builder import BaseDatasetBuilder\n        \"\"\"\n\n        def wrap(builder_cls):\n            from xraypulse.datasets.builders.base_dataset_builder import BaseDatasetBuilder\n\n            assert issubclass(\n                builder_cls, BaseDatasetBuilder\n            ), \"All builders must inherit BaseDatasetBuilder class, found {}\".format(\n                builder_cls\n            )\n            if name in cls.mapping[\"builder_name_mapping\"]:\n                raise KeyError(\n                    \"Name '{}' already registered for {}.\".format(\n                        name, cls.mapping[\"builder_name_mapping\"][name]\n                    )\n                )\n            cls.mapping[\"builder_name_mapping\"][name] = builder_cls\n            return builder_cls\n\n        return wrap\n\n    @classmethod\n    def register_task(cls, name):\n        r\"\"\"Register a task to registry with key 'name'\n\n        Args:\n            name: Key with which the task will be registered.\n\n        Usage:\n\n            from minigpt4.common.registry import registry\n        \"\"\"\n\n        def wrap(task_cls):\n            from xraypulse.tasks.base_task import BaseTask\n\n            assert issubclass(\n                task_cls, BaseTask\n            ), \"All tasks must inherit BaseTask class\"\n            if name in cls.mapping[\"task_name_mapping\"]:\n                raise KeyError(\n                    \"Name '{}' already registered for {}.\".format(\n                        name, cls.mapping[\"task_name_mapping\"][name]\n                    )\n                )\n            cls.mapping[\"task_name_mapping\"][name] = task_cls\n            return task_cls\n\n        return wrap\n\n    @classmethod\n    def register_model(cls, name):\n        r\"\"\"Register a task to registry with key 'name'\n\n        Args:\n            name: Key with which the task will be registered.\n\n        Usage:\n\n            from xraypulse.common.registry import registry\n        \"\"\"\n\n        def wrap(model_cls):\n            from xraypulse.models import BaseModel\n\n            assert issubclass(\n                model_cls, BaseModel\n            ), \"All models must inherit BaseModel class\"\n            if name in cls.mapping[\"model_name_mapping\"]:\n                raise KeyError(\n                    \"Name '{}' already registered for {}.\".format(\n                        name, cls.mapping[\"model_name_mapping\"][name]\n                    )\n                )\n            cls.mapping[\"model_name_mapping\"][name] = model_cls\n            return model_cls\n\n        return wrap\n\n    @classmethod\n    def register_processor(cls, name):\n        r\"\"\"Register a processor to registry with key 'name'\n\n        Args:\n            name: Key with which the task will be registered.\n\n        Usage:\n\n            from xraypulse.common.registry import registry\n        \"\"\"\n\n        def wrap(processor_cls):\n            from xraypulse.processors import BaseProcessor\n\n            assert issubclass(\n                processor_cls, BaseProcessor\n            ), \"All processors must inherit BaseProcessor class\"\n            if name in cls.mapping[\"processor_name_mapping\"]:\n                raise KeyError(\n                    \"Name '{}' already registered for {}.\".format(\n                        name, cls.mapping[\"processor_name_mapping\"][name]\n                    )\n                )\n            cls.mapping[\"processor_name_mapping\"][name] = processor_cls\n            return processor_cls\n\n        return wrap\n\n    @classmethod\n    def register_lr_scheduler(cls, name):\n        r\"\"\"Register a model to registry with key 'name'\n\n        Args:\n            name: Key with which the task will be registered.\n\n        Usage:\n\n            from xraypulse.common.registry import registry\n        \"\"\"\n\n        def wrap(lr_sched_cls):\n            if name in cls.mapping[\"lr_scheduler_name_mapping\"]:\n                raise KeyError(\n                    \"Name '{}' already registered for {}.\".format(\n                        name, cls.mapping[\"lr_scheduler_name_mapping\"][name]\n                    )\n                )\n            cls.mapping[\"lr_scheduler_name_mapping\"][name] = lr_sched_cls\n            return lr_sched_cls\n\n        return wrap\n\n    @classmethod\n    def register_runner(cls, name):\n        r\"\"\"Register a model to registry with key 'name'\n\n        Args:\n            name: Key with which the task will be registered.\n\n        Usage:\n\n            from xraypulse.common.registry import registry\n        \"\"\"\n\n        def wrap(runner_cls):\n            if name in cls.mapping[\"runner_name_mapping\"]:\n                raise KeyError(\n                    \"Name '{}' already registered for {}.\".format(\n                        name, cls.mapping[\"runner_name_mapping\"][name]\n                    )\n                )\n            cls.mapping[\"runner_name_mapping\"][name] = runner_cls\n            return runner_cls\n\n        return wrap\n\n    @classmethod\n    def register_path(cls, name, path):\n        r\"\"\"Register a path to registry with key 'name'\n\n        Args:\n            name: Key with which the path will be registered.\n\n        Usage:\n\n            from xraypulse.common.registry import registry\n        \"\"\"\n        assert isinstance(path, str), \"All path must be str.\"\n        if name in cls.mapping[\"paths\"]:\n            raise KeyError(\"Name '{}' already registered.\".format(name))\n        cls.mapping[\"paths\"][name] = path\n\n    @classmethod\n    def register(cls, name, obj):\n        r\"\"\"Register an item to registry with key 'name'\n\n        Args:\n            name: Key with which the item will be registered.\n\n        Usage::\n\n            from minigpt4.common.registry import registry\n\n            registry.register(\"config\", {})\n        \"\"\"\n        path = name.split(\".\")\n        current = cls.mapping[\"state\"]\n\n        for part in path[:-1]:\n            if part not in current:\n                current[part] = {}\n            current = current[part]\n\n        current[path[-1]] = obj\n\n    # @classmethod\n    # def get_trainer_class(cls, name):\n    #     return cls.mapping[\"trainer_name_mapping\"].get(name, None)\n\n    @classmethod\n    def get_builder_class(cls, name):\n        return cls.mapping[\"builder_name_mapping\"].get(name, None)\n\n    @classmethod\n    def get_model_class(cls, name):\n        return cls.mapping[\"model_name_mapping\"].get(name, None)\n\n    @classmethod\n    def get_task_class(cls, name):\n        return cls.mapping[\"task_name_mapping\"].get(name, None)\n\n    @classmethod\n    def get_processor_class(cls, name):\n        return cls.mapping[\"processor_name_mapping\"].get(name, None)\n\n    @classmethod\n    def get_lr_scheduler_class(cls, name):\n        return cls.mapping[\"lr_scheduler_name_mapping\"].get(name, None)\n\n    @classmethod\n    def get_runner_class(cls, name):\n        return cls.mapping[\"runner_name_mapping\"].get(name, None)\n\n    @classmethod\n    def list_runners(cls):\n        return sorted(cls.mapping[\"runner_name_mapping\"].keys())\n\n    @classmethod\n    def list_models(cls):\n        return sorted(cls.mapping[\"model_name_mapping\"].keys())\n\n    @classmethod\n    def list_tasks(cls):\n        return sorted(cls.mapping[\"task_name_mapping\"].keys())\n\n    @classmethod\n    def list_processors(cls):\n        return sorted(cls.mapping[\"processor_name_mapping\"].keys())\n\n    @classmethod\n    def list_lr_schedulers(cls):\n        return sorted(cls.mapping[\"lr_scheduler_name_mapping\"].keys())\n\n    @classmethod\n    def list_datasets(cls):\n        return sorted(cls.mapping[\"builder_name_mapping\"].keys())\n\n    @classmethod\n    def get_path(cls, name):\n        return cls.mapping[\"paths\"].get(name, None)\n\n    @classmethod\n    def get(cls, name, default=None, no_warning=False):\n        r\"\"\"Get an item from registry with key 'name'\n\n        Args:\n            name (string): Key whose value needs to be retrieved.\n            default: If passed and key is not in registry, default value will\n                     be returned with a warning. Default: None\n            no_warning (bool): If passed as True, warning when key doesn't exist\n                               will not be generated. Useful for MMF's\n                               internal operations. Default: False\n        \"\"\"\n        original_name = name\n        name = name.split(\".\")\n        value = cls.mapping[\"state\"]\n        for subname in name:\n            value = value.get(subname, default)\n            if value is default:\n                break\n\n        if (\n            \"writer\" in cls.mapping[\"state\"]\n            and value == default\n            and no_warning is False\n        ):\n            cls.mapping[\"state\"][\"writer\"].warning(\n                \"Key {} is not present in registry, returning default value \"\n                \"of {}\".format(original_name, default)\n            )\n        return value\n\n    @classmethod\n    def unregister(cls, name):\n        r\"\"\"Remove an item from registry with key 'name'\n\n        Args:\n            name: Key which needs to be removed.\n        Usage::\n\n            from mmf.common.registry import registry\n\n            config = registry.unregister(\"config\")\n        \"\"\"\n        return cls.mapping[\"state\"].pop(name, None)\n\n\nregistry = Registry()\n"
  },
  {
    "path": "xraypulse/common/utils.py",
    "content": "\"\"\"\n Copyright (c) 2022, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\n\nimport io\nimport json\nimport logging\nimport os\nimport pickle\nimport re\nimport shutil\nimport urllib\nimport urllib.error\nimport urllib.request\nfrom typing import Optional\nfrom urllib.parse import urlparse\n\nimport numpy as np\nimport pandas as pd\nimport yaml\nfrom iopath.common.download import download\nfrom iopath.common.file_io import file_lock, g_pathmgr\nfrom xraypulse.common.registry import registry\nfrom torch.utils.model_zoo import tqdm\nfrom torchvision.datasets.utils import (\n    check_integrity,\n    download_file_from_google_drive,\n    extract_archive,\n)\n\n\ndef now():\n    from datetime import datetime\n\n    return datetime.now().strftime(\"%Y%m%d%H%M\")[:-1]\n\n\ndef is_url(url_or_filename):\n    parsed = urlparse(url_or_filename)\n    return parsed.scheme in (\"http\", \"https\")\n\n\ndef get_cache_path(rel_path):\n    return os.path.expanduser(os.path.join(registry.get_path(\"cache_root\"), rel_path))\n\n\ndef get_abs_path(rel_path):\n    return os.path.join(registry.get_path(\"library_root\"), rel_path)\n\n\ndef load_json(filename):\n    with open(filename, \"r\") as f:\n        return json.load(f)\n\n\n# The following are adapted from torchvision and vissl\n# torchvision: https://github.com/pytorch/vision\n# vissl: https://github.com/facebookresearch/vissl/blob/main/vissl/utils/download.py\n\n\ndef makedir(dir_path):\n    \"\"\"\n    Create the directory if it does not exist.\n    \"\"\"\n    is_success = False\n    try:\n        if not g_pathmgr.exists(dir_path):\n            g_pathmgr.mkdirs(dir_path)\n        is_success = True\n    except BaseException:\n        print(f\"Error creating directory: {dir_path}\")\n    return is_success\n\n\ndef get_redirected_url(url: str):\n    \"\"\"\n    Given a URL, returns the URL it redirects to or the\n    original URL in case of no indirection\n    \"\"\"\n    import requests\n\n    with requests.Session() as session:\n        with session.get(url, stream=True, allow_redirects=True) as response:\n            if response.history:\n                return response.url\n            else:\n                return url\n\n\ndef to_google_drive_download_url(view_url: str) -> str:\n    \"\"\"\n    Utility function to transform a view URL of google drive\n    to a download URL for google drive\n    Example input:\n        https://drive.google.com/file/d/137RyRjvTBkBiIfeYBNZBtViDHQ6_Ewsp/view\n    Example output:\n        https://drive.google.com/uc?export=download&id=137RyRjvTBkBiIfeYBNZBtViDHQ6_Ewsp\n    \"\"\"\n    splits = view_url.split(\"/\")\n    assert splits[-1] == \"view\"\n    file_id = splits[-2]\n    return f\"https://drive.google.com/uc?export=download&id={file_id}\"\n\n\ndef download_google_drive_url(url: str, output_path: str, output_file_name: str):\n    \"\"\"\n    Download a file from google drive\n    Downloading an URL from google drive requires confirmation when\n    the file of the size is too big (google drive notifies that\n    anti-viral checks cannot be performed on such files)\n    \"\"\"\n    import requests\n\n    with requests.Session() as session:\n\n        # First get the confirmation token and append it to the URL\n        with session.get(url, stream=True, allow_redirects=True) as response:\n            for k, v in response.cookies.items():\n                if k.startswith(\"download_warning\"):\n                    url = url + \"&confirm=\" + v\n\n        # Then download the content of the file\n        with session.get(url, stream=True, verify=True) as response:\n            makedir(output_path)\n            path = os.path.join(output_path, output_file_name)\n            total_size = int(response.headers.get(\"Content-length\", 0))\n            with open(path, \"wb\") as file:\n                from tqdm import tqdm\n\n                with tqdm(total=total_size) as progress_bar:\n                    for block in response.iter_content(\n                        chunk_size=io.DEFAULT_BUFFER_SIZE\n                    ):\n                        file.write(block)\n                        progress_bar.update(len(block))\n\n\ndef _get_google_drive_file_id(url: str) -> Optional[str]:\n    parts = urlparse(url)\n\n    if re.match(r\"(drive|docs)[.]google[.]com\", parts.netloc) is None:\n        return None\n\n    match = re.match(r\"/file/d/(?P<id>[^/]*)\", parts.path)\n    if match is None:\n        return None\n\n    return match.group(\"id\")\n\n\ndef _urlretrieve(url: str, filename: str, chunk_size: int = 1024) -> None:\n    with open(filename, \"wb\") as fh:\n        with urllib.request.urlopen(\n            urllib.request.Request(url, headers={\"User-Agent\": \"vissl\"})\n        ) as response:\n            with tqdm(total=response.length) as pbar:\n                for chunk in iter(lambda: response.read(chunk_size), \"\"):\n                    if not chunk:\n                        break\n                    pbar.update(chunk_size)\n                    fh.write(chunk)\n\n\ndef download_url(\n    url: str,\n    root: str,\n    filename: Optional[str] = None,\n    md5: Optional[str] = None,\n) -> None:\n    \"\"\"Download a file from a url and place it in root.\n    Args:\n        url (str): URL to download file from\n        root (str): Directory to place downloaded file in\n        filename (str, optional): Name to save the file under.\n                                  If None, use the basename of the URL.\n        md5 (str, optional): MD5 checksum of the download. If None, do not check\n    \"\"\"\n    root = os.path.expanduser(root)\n    if not filename:\n        filename = os.path.basename(url)\n    fpath = os.path.join(root, filename)\n\n    makedir(root)\n\n    # check if file is already present locally\n    if check_integrity(fpath, md5):\n        print(\"Using downloaded and verified file: \" + fpath)\n        return\n\n    # expand redirect chain if needed\n    url = get_redirected_url(url)\n\n    # check if file is located on Google Drive\n    file_id = _get_google_drive_file_id(url)\n    if file_id is not None:\n        return download_file_from_google_drive(file_id, root, filename, md5)\n\n    # download the file\n    try:\n        print(\"Downloading \" + url + \" to \" + fpath)\n        _urlretrieve(url, fpath)\n    except (urllib.error.URLError, IOError) as e:  # type: ignore[attr-defined]\n        if url[:5] == \"https\":\n            url = url.replace(\"https:\", \"http:\")\n            print(\n                \"Failed download. Trying https -> http instead.\"\n                \" Downloading \" + url + \" to \" + fpath\n            )\n            _urlretrieve(url, fpath)\n        else:\n            raise e\n\n    # check integrity of downloaded file\n    if not check_integrity(fpath, md5):\n        raise RuntimeError(\"File not found or corrupted.\")\n\n\ndef download_and_extract_archive(\n    url: str,\n    download_root: str,\n    extract_root: Optional[str] = None,\n    filename: Optional[str] = None,\n    md5: Optional[str] = None,\n    remove_finished: bool = False,\n) -> None:\n    download_root = os.path.expanduser(download_root)\n    if extract_root is None:\n        extract_root = download_root\n    if not filename:\n        filename = os.path.basename(url)\n\n    download_url(url, download_root, filename, md5)\n\n    archive = os.path.join(download_root, filename)\n    print(\"Extracting {} to {}\".format(archive, extract_root))\n    extract_archive(archive, extract_root, remove_finished)\n\n\ndef cache_url(url: str, cache_dir: str) -> str:\n    \"\"\"\n    This implementation downloads the remote resource and caches it locally.\n    The resource will only be downloaded if not previously requested.\n    \"\"\"\n    parsed_url = urlparse(url)\n    dirname = os.path.join(cache_dir, os.path.dirname(parsed_url.path.lstrip(\"/\")))\n    makedir(dirname)\n    filename = url.split(\"/\")[-1]\n    cached = os.path.join(dirname, filename)\n    with file_lock(cached):\n        if not os.path.isfile(cached):\n            logging.info(f\"Downloading {url} to {cached} ...\")\n            cached = download(url, dirname, filename=filename)\n    logging.info(f\"URL {url} cached in {cached}\")\n    return cached\n\n\n# TODO (prigoyal): convert this into RAII-style API\ndef create_file_symlink(file1, file2):\n    \"\"\"\n    Simply create the symlinks for a given file1 to file2.\n    Useful during model checkpointing to symlinks to the\n    latest successful checkpoint.\n    \"\"\"\n    try:\n        if g_pathmgr.exists(file2):\n            g_pathmgr.rm(file2)\n        g_pathmgr.symlink(file1, file2)\n    except Exception as e:\n        logging.info(f\"Could NOT create symlink. Error: {e}\")\n\n\ndef save_file(data, filename, append_to_json=True, verbose=True):\n    \"\"\"\n    Common i/o utility to handle saving data to various file formats.\n    Supported:\n        .pkl, .pickle, .npy, .json\n    Specifically for .json, users have the option to either append (default)\n    or rewrite by passing in Boolean value to append_to_json.\n    \"\"\"\n    if verbose:\n        logging.info(f\"Saving data to file: {filename}\")\n    file_ext = os.path.splitext(filename)[1]\n    if file_ext in [\".pkl\", \".pickle\"]:\n        with g_pathmgr.open(filename, \"wb\") as fopen:\n            pickle.dump(data, fopen, pickle.HIGHEST_PROTOCOL)\n    elif file_ext == \".npy\":\n        with g_pathmgr.open(filename, \"wb\") as fopen:\n            np.save(fopen, data)\n    elif file_ext == \".json\":\n        if append_to_json:\n            with g_pathmgr.open(filename, \"a\") as fopen:\n                fopen.write(json.dumps(data, sort_keys=True) + \"\\n\")\n                fopen.flush()\n        else:\n            with g_pathmgr.open(filename, \"w\") as fopen:\n                fopen.write(json.dumps(data, sort_keys=True) + \"\\n\")\n                fopen.flush()\n    elif file_ext == \".yaml\":\n        with g_pathmgr.open(filename, \"w\") as fopen:\n            dump = yaml.dump(data)\n            fopen.write(dump)\n            fopen.flush()\n    else:\n        raise Exception(f\"Saving {file_ext} is not supported yet\")\n\n    if verbose:\n        logging.info(f\"Saved data to file: {filename}\")\n\n\ndef load_file(filename, mmap_mode=None, verbose=True, allow_pickle=False):\n    \"\"\"\n    Common i/o utility to handle loading data from various file formats.\n    Supported:\n        .pkl, .pickle, .npy, .json\n    For the npy files, we support reading the files in mmap_mode.\n    If the mmap_mode of reading is not successful, we load data without the\n    mmap_mode.\n    \"\"\"\n    if verbose:\n        logging.info(f\"Loading data from file: {filename}\")\n\n    file_ext = os.path.splitext(filename)[1]\n    if file_ext == \".txt\":\n        with g_pathmgr.open(filename, \"r\") as fopen:\n            data = fopen.readlines()\n    elif file_ext in [\".pkl\", \".pickle\"]:\n        with g_pathmgr.open(filename, \"rb\") as fopen:\n            data = pickle.load(fopen, encoding=\"latin1\")\n    elif file_ext == \".npy\":\n        if mmap_mode:\n            try:\n                with g_pathmgr.open(filename, \"rb\") as fopen:\n                    data = np.load(\n                        fopen,\n                        allow_pickle=allow_pickle,\n                        encoding=\"latin1\",\n                        mmap_mode=mmap_mode,\n                    )\n            except ValueError as e:\n                logging.info(\n                    f\"Could not mmap {filename}: {e}. Trying without g_pathmgr\"\n                )\n                data = np.load(\n                    filename,\n                    allow_pickle=allow_pickle,\n                    encoding=\"latin1\",\n                    mmap_mode=mmap_mode,\n                )\n                logging.info(\"Successfully loaded without g_pathmgr\")\n            except Exception:\n                logging.info(\"Could not mmap without g_pathmgr. Trying without mmap\")\n                with g_pathmgr.open(filename, \"rb\") as fopen:\n                    data = np.load(fopen, allow_pickle=allow_pickle, encoding=\"latin1\")\n        else:\n            with g_pathmgr.open(filename, \"rb\") as fopen:\n                data = np.load(fopen, allow_pickle=allow_pickle, encoding=\"latin1\")\n    elif file_ext == \".json\":\n        with g_pathmgr.open(filename, \"r\") as fopen:\n            data = json.load(fopen)\n    elif file_ext == \".yaml\":\n        with g_pathmgr.open(filename, \"r\") as fopen:\n            data = yaml.load(fopen, Loader=yaml.FullLoader)\n    elif file_ext == \".csv\":\n        with g_pathmgr.open(filename, \"r\") as fopen:\n            data = pd.read_csv(fopen)\n    else:\n        raise Exception(f\"Reading from {file_ext} is not supported yet\")\n    return data\n\n\ndef abspath(resource_path: str):\n    \"\"\"\n    Make a path absolute, but take into account prefixes like\n    \"http://\" or \"manifold://\"\n    \"\"\"\n    regex = re.compile(r\"^\\w+://\")\n    if regex.match(resource_path) is None:\n        return os.path.abspath(resource_path)\n    else:\n        return resource_path\n\n\ndef makedir(dir_path):\n    \"\"\"\n    Create the directory if it does not exist.\n    \"\"\"\n    is_success = False\n    try:\n        if not g_pathmgr.exists(dir_path):\n            g_pathmgr.mkdirs(dir_path)\n        is_success = True\n    except BaseException:\n        logging.info(f\"Error creating directory: {dir_path}\")\n    return is_success\n\n\ndef is_url(input_url):\n    \"\"\"\n    Check if an input string is a url. look for http(s):// and ignoring the case\n    \"\"\"\n    is_url = re.match(r\"^(?:http)s?://\", input_url, re.IGNORECASE) is not None\n    return is_url\n\n\ndef cleanup_dir(dir):\n    \"\"\"\n    Utility for deleting a directory. Useful for cleaning the storage space\n    that contains various training artifacts like checkpoints, data etc.\n    \"\"\"\n    if os.path.exists(dir):\n        logging.info(f\"Deleting directory: {dir}\")\n        shutil.rmtree(dir)\n    logging.info(f\"Deleted contents of directory: {dir}\")\n\n\ndef get_file_size(filename):\n    \"\"\"\n    Given a file, get the size of file in MB\n    \"\"\"\n    size_in_mb = os.path.getsize(filename) / float(1024**2)\n    return size_in_mb\n"
  },
  {
    "path": "xraypulse/configs/datasets/mimic/defaults.yaml",
    "content": "datasets:\n  mimic:\n    data_type: images\n    build_info:\n      storage: /mnt/petrelfs/share_data/huangzhongzhen/multimodal_pretrain/dataset/mimic\n"
  },
  {
    "path": "xraypulse/configs/datasets/openi/defaults.yaml",
    "content": "datasets:\n  openi:\n    data_type: images\n    build_info:\n      storage: /mnt/petrelfs/share_data/huangzhongzhen/multimodal_pretrain/dataset/openi\n"
  },
  {
    "path": "xraypulse/configs/default.yaml",
    "content": "env:\n  # For default users\n  # cache_root: \"cache\"\n  # For internal use with persistent storage\n  cache_root: \"/export/home/.cache/xraypulse\"\n"
  },
  {
    "path": "xraypulse/configs/models/xraypulse.yaml",
    "content": "model:\n  arch: xray_pulse\n\n  # vit encoder\n  image_size: 224\n  drop_path_rate: 0\n  use_grad_checkpoint: False\n  vit_precision: \"fp16\"\n  freeze_vit: True\n  freeze_qformer: True\n\n  # Q-Former\n  num_query_token: 32\n\n  # Vicuna\n  bloom_model: \"OpenMEDLab/PULSE-7bv5\"\n\n  # generation configs\n  prompt: \"\"\n\npreprocess:\n    vis_processor:\n        train:\n          name: \"blip2_image_train\"\n          image_size: 224\n        eval:\n          name: \"blip2_image_eval\"\n          image_size: 224\n    text_processor:\n        train:\n          name: \"blip_caption\"\n        eval:\n          name: \"blip_caption\"\n"
  },
  {
    "path": "xraypulse/conversation/__init__.py",
    "content": ""
  },
  {
    "path": "xraypulse/conversation/conversation.py",
    "content": "import argparse\nimport time\nfrom PIL import Image\n\nimport torch\nfrom transformers import AutoTokenizer, AutoModelForCausalLM, BloomTokenizerFast\nfrom transformers import StoppingCriteria, StoppingCriteriaList\n\nimport dataclasses\nfrom enum import auto, Enum\nfrom typing import List, Tuple, Any\n\nfrom xraypulse.common.registry import registry\n\n\nclass SeparatorStyle(Enum):\n    \"\"\"Different separator style.\"\"\"\n    SINGLE = auto()\n    TWO = auto()\n\n\n@dataclasses.dataclass\nclass Conversation:\n    \"\"\"A class that keeps all conversation history.\"\"\"\n    system: str\n    roles: List[str]\n    messages: List[List[str]]\n    offset: int\n    # system_img: List[Image.Image] = []\n    sep_style: SeparatorStyle = SeparatorStyle.SINGLE\n    sep: str = \"###\"\n    sep2: str = None\n\n    skip_next: bool = False\n    conv_id: Any = None\n\n    def get_prompt(self):\n        if self.sep_style == SeparatorStyle.SINGLE:\n            ret = self.system + self.sep\n            for role, message in self.messages:\n                if message:\n                    ret += role + \": \" + message + self.sep\n                else:\n                    ret += role + \":\"\n            return ret\n        elif self.sep_style == SeparatorStyle.TWO:\n            seps = [self.sep, self.sep2]\n            ret = self.system + seps[0]\n            for i, (role, message) in enumerate(self.messages):\n                if message:\n                    ret += role + \": \" + message + seps[i % 2]\n                else:\n                    ret += role + \":\"\n            return ret\n        else:\n            raise ValueError(f\"Invalid style: {self.sep_style}\")\n\n    def append_message(self, role, message):\n        self.messages.append([role, message])\n\n    def to_gradio_chatbot(self):\n        ret = []\n        for i, (role, msg) in enumerate(self.messages[self.offset:]):\n            if i % 2 == 0:\n                ret.append([msg, None])\n            else:\n                ret[-1][-1] = msg\n        print('ret:')\n        print(ret)\n        return ret\n\n    def copy(self):\n        return Conversation(\n            system=self.system,\n            # system_img=self.system_img,\n            roles=self.roles,\n            messages=[[x, y] for x, y in self.messages],\n            offset=self.offset,\n            sep_style=self.sep_style,\n            sep=self.sep,\n            sep2=self.sep2,\n            conv_id=self.conv_id)\n\n    def dict(self):\n        return {\n            \"system\": self.system,\n            # \"system_img\": self.system_img,\n            \"roles\": self.roles,\n            \"messages\": self.messages,\n            \"offset\": self.offset,\n            \"sep\": self.sep,\n            \"sep2\": self.sep2,\n            \"conv_id\": self.conv_id,\n        }\n\n\nclass StoppingCriteriaSub(StoppingCriteria):\n\n    def __init__(self, stops=[], encounters=1):\n        super().__init__()\n        self.stops = stops\n\n    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor):\n        for stop in self.stops:\n            if torch.all((stop == input_ids[0][-len(stop):])).item():\n                return True\n\n        return False\n\nCONV_ZH = Conversation(\n    system=\"Instructions: You are PULSE, a large language model trained by SHAIlab. Answer as concisely as possible.\\nKnowledge cutoff: 2021-09-01\\nCurrent date: 2022-02-01</s> User: {} </s> Helper: \",\n        # \"Please answer the medical questions based on the patient's description. Give the following medical scan: <Img>图片</Img>.\"\n        # \"You will be able to see the medical scan once I provide it to you. Please answer the patients questions.\",\n    roles=(\"User\", \"Helper\"),\n    messages=[],\n    offset=0,\n    sep_style=SeparatorStyle.SINGLE,\n    sep=\"</s>\",\n    sep2=\"###\",  \n)\n\n\nclass Chat:\n    def __init__(self, model, vis_processor, device='cuda:0'):\n        self.device = device\n        self.model = model\n        self.vis_processor = vis_processor\n        stop_words_ids = [torch.tensor([835]).to(self.device),\n                          torch.tensor([2277, 29937]).to(self.device)]  # '###' can be encoded in two different ways.\n        self.stopping_criteria = StoppingCriteriaList([StoppingCriteriaSub(stops=stop_words_ids)])\n\n    def ask(self, text, conv):\n        if len(conv.messages) > 0 and conv.messages[-1][0] == conv.roles[0] \\\n                and conv.messages[-1][1][-6:] == '</Img>':  # last message is image.\n            conv.messages[-1][1] = ' '.join([conv.messages[-1][1], text])\n        else:\n            conv.append_message(conv.roles[0], text)\n\n    def answer(self, conv, img_list, max_new_tokens=300, num_beams=1, min_length=1, top_p=0.9,\n               repetition_penalty=1.0, length_penalty=1, temperature=1.0, max_length=2000):\n        conv.append_message(conv.roles[1], None)\n        embs = self.get_context_emb(conv, img_list)\n\n        current_max_len = embs.shape[1] + max_new_tokens\n        if current_max_len - max_length > 0:\n            print('Warning: The number of tokens in current conversation exceeds the max length. '\n                  'The model will not see the contexts outside the range.')\n        begin_idx = max(0, current_max_len - max_length)\n\n        embs = embs[:, begin_idx:]\n\n        outputs = self.model.bloom_model.generate(\n            inputs_embeds=embs,\n            max_new_tokens=max_new_tokens,\n            stopping_criteria=self.stopping_criteria,\n            num_beams=num_beams,\n            do_sample=True,\n            min_length=min_length,\n            top_p=top_p,\n            repetition_penalty=repetition_penalty,\n            length_penalty=length_penalty,\n            temperature=temperature,\n        )\n        output_token = outputs[0]\n        if output_token[0] == 0:  # the model might output a unknow token <unk> at the beginning. remove it\n            output_token = output_token[1:]\n        if output_token[0] == 1:  # some users find that there is a start token <s> at the beginning. remove it\n            output_token = output_token[1:]\n        output_text = self.model.bloom_tokenizer.decode(output_token, add_special_tokens=False)\n        output_text = output_text.split('</s>')[0]  # remove the stop sign '###'\n        output_text = output_text.split('###')[0]\n        conv.messages[-1][1] = output_text\n        return output_text, output_token.cpu().numpy()\n\n    def upload_img(self, image, conv, img_list):\n        if isinstance(image, str):  # is a image path\n            raw_image = Image.open(image).convert('RGB')\n            image = self.vis_processor(raw_image).unsqueeze(0).to(self.device)\n        elif isinstance(image, Image.Image):\n            raw_image = image\n            image = self.vis_processor(raw_image).unsqueeze(0).to(self.device)\n        elif isinstance(image, torch.Tensor):\n            if len(image.shape) == 3:\n                image = image.unsqueeze(0)\n            image = image.to(self.device)\n\n        image_emb, _ = self.model.encode_img(image)\n        img_list.append(image_emb)\n        conv.append_message(conv.roles[0], \"<Img><图片></Img>\")\n        msg = \"Received.\"\n        # self.conv.append_message(self.conv.roles[1], msg)\n        return msg\n\n    def get_context_emb(self, conv, img_list):\n        prompt = conv.get_prompt()\n        prompt_segs = prompt.split('<图片>')\n        assert len(prompt_segs) == len(img_list) + 1, \"Unmatched numbers of image placeholders and images.\"\n        seg_tokens = [\n            self.model.bloom_tokenizer(\n                seg, return_tensors=\"pt\", add_special_tokens=i == 0).to(self.device).input_ids\n            # only add bos to the first seg\n            for i, seg in enumerate(prompt_segs)\n        ]\n        seg_embs = [self.model.bloom_model.transformer.word_embeddings_layernorm(self.model.bloom_model.transformer.word_embeddings(seg_t)) for seg_t in seg_tokens]\n        mixed_embs = [emb for pair in zip(seg_embs[:-1], img_list) for emb in pair] + [seg_embs[-1]]\n        mixed_embs = torch.cat(mixed_embs, dim=1)\n        return mixed_embs\n\n\n"
  },
  {
    "path": "xraypulse/datasets/__init__.py",
    "content": ""
  },
  {
    "path": "xraypulse/datasets/builders/__init__.py",
    "content": "\"\"\"\n Copyright (c) 2022, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\n\nfrom xraypulse.datasets.builders.base_dataset_builder import load_dataset_config\nfrom xraypulse.datasets.builders.image_text_pair_builder import (\n    MIMICBuilder,\n    OpenIBuilder,\n)\nfrom xraypulse.common.registry import registry\n\n__all__ = [\n    \"MIMICBuilder\",\n    \"OpenIBuilder\",\n]\n\n\ndef load_dataset(name, cfg_path=None, vis_path=None, data_type=None):\n    \"\"\"\n    Example\n\n    >>> dataset = load_dataset(\"coco_caption\", cfg=None)\n    >>> splits = dataset.keys()\n    >>> print([len(dataset[split]) for split in splits])\n\n    \"\"\"\n    if cfg_path is None:\n        cfg = None\n    else:\n        cfg = load_dataset_config(cfg_path)\n\n    try:\n        builder = registry.get_builder_class(name)(cfg)\n    except TypeError:\n        print(\n            f\"Dataset {name} not found. Available datasets:\\n\"\n            + \", \".join([str(k) for k in dataset_zoo.get_names()])\n        )\n        exit(1)\n\n    if vis_path is not None:\n        if data_type is None:\n            # use default data type in the config\n            data_type = builder.config.data_type\n\n        assert (\n            data_type in builder.config.build_info\n        ), f\"Invalid data_type {data_type} for {name}.\"\n\n        builder.config.build_info.get(data_type).storage = vis_path\n\n    dataset = builder.build_datasets()\n    return dataset\n\n\nclass DatasetZoo:\n    def __init__(self) -> None:\n        self.dataset_zoo = {\n            k: list(v.DATASET_CONFIG_DICT.keys())\n            for k, v in sorted(registry.mapping[\"builder_name_mapping\"].items())\n        }\n\n    def get_names(self):\n        return list(self.dataset_zoo.keys())\n\n\ndataset_zoo = DatasetZoo()\n"
  },
  {
    "path": "xraypulse/datasets/builders/base_dataset_builder.py",
    "content": "\"\"\"\n This file is from\n Copyright (c) 2022, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\n\nimport logging\nimport os\nimport shutil\nimport warnings\n\nfrom omegaconf import OmegaConf\nimport torch.distributed as dist\nfrom torchvision.datasets.utils import download_url\n\nimport xraypulse.common.utils as utils\nfrom xraypulse.common.dist_utils import is_dist_avail_and_initialized, is_main_process\nfrom xraypulse.common.registry import registry\nfrom xraypulse.processors.base_processor import BaseProcessor\n\n\n\nclass BaseDatasetBuilder:\n    train_dataset_cls, eval_dataset_cls = None, None\n\n    def __init__(self, cfg=None):\n        super().__init__()\n\n        if cfg is None:\n            # help to create datasets from default config.\n            self.config = load_dataset_config(self.default_config_path())\n        elif isinstance(cfg, str):\n            self.config = load_dataset_config(cfg)\n        else:\n            # when called from task.build_dataset()\n            self.config = cfg\n\n        self.data_type = self.config.data_type\n\n        self.vis_processors = {\"train\": BaseProcessor(), \"eval\": BaseProcessor()}\n        self.text_processors = {\"train\": BaseProcessor(), \"eval\": BaseProcessor()}\n\n    def build_datasets(self):\n        # download, split, etc...\n        # only called on 1 GPU/TPU in distributed\n\n        if is_main_process():\n            self._download_data()\n\n        if is_dist_avail_and_initialized():\n            dist.barrier()\n\n        # at this point, all the annotations and image/videos should be all downloaded to the specified locations.\n        logging.info(\"Building datasets...\")\n        datasets = self.build()  # dataset['train'/'val'/'test']\n\n        return datasets\n\n    def build_processors(self):\n        vis_proc_cfg = self.config.get(\"vis_processor\")\n        txt_proc_cfg = self.config.get(\"text_processor\")\n\n        if vis_proc_cfg is not None:\n            vis_train_cfg = vis_proc_cfg.get(\"train\")\n            vis_eval_cfg = vis_proc_cfg.get(\"eval\")\n\n            self.vis_processors[\"train\"] = self._build_proc_from_cfg(vis_train_cfg)\n            self.vis_processors[\"eval\"] = self._build_proc_from_cfg(vis_eval_cfg)\n\n        if txt_proc_cfg is not None:\n            txt_train_cfg = txt_proc_cfg.get(\"train\")\n            txt_eval_cfg = txt_proc_cfg.get(\"eval\")\n\n            self.text_processors[\"train\"] = self._build_proc_from_cfg(txt_train_cfg)\n            self.text_processors[\"eval\"] = self._build_proc_from_cfg(txt_eval_cfg)\n\n    @staticmethod\n    def _build_proc_from_cfg(cfg):\n        return (\n            registry.get_processor_class(cfg.name).from_config(cfg)\n            if cfg is not None\n            else None\n        )\n\n    @classmethod\n    def default_config_path(cls, type=\"default\"):\n        return utils.get_abs_path(cls.DATASET_CONFIG_DICT[type])\n\n    def _download_data(self):\n        self._download_ann()\n        self._download_vis()\n\n    def _download_ann(self):\n        \"\"\"\n        Download annotation files if necessary.\n        All the vision-language datasets should have annotations of unified format.\n\n        storage_path can be:\n          (1) relative/absolute: will be prefixed with env.cache_root to make full path if relative.\n          (2) basename/dirname: will be suffixed with base name of URL if dirname is provided.\n\n        Local annotation paths should be relative.\n        \"\"\"\n        anns = self.config.build_info.annotations\n\n        splits = anns.keys()\n\n        cache_root = registry.get_path(\"cache_root\")\n\n        for split in splits:\n            info = anns[split]\n\n            urls, storage_paths = info.get(\"url\", None), info.storage\n\n            if isinstance(urls, str):\n                urls = [urls]\n            if isinstance(storage_paths, str):\n                storage_paths = [storage_paths]\n\n            assert len(urls) == len(storage_paths)\n\n            for url_or_filename, storage_path in zip(urls, storage_paths):\n                # if storage_path is relative, make it full by prefixing with cache_root.\n                if not os.path.isabs(storage_path):\n                    storage_path = os.path.join(cache_root, storage_path)\n\n                dirname = os.path.dirname(storage_path)\n                if not os.path.exists(dirname):\n                    os.makedirs(dirname)\n\n                if os.path.isfile(url_or_filename):\n                    src, dst = url_or_filename, storage_path\n                    if not os.path.exists(dst):\n                        shutil.copyfile(src=src, dst=dst)\n                    else:\n                        logging.info(\"Using existing file {}.\".format(dst))\n                else:\n                    if os.path.isdir(storage_path):\n                        # if only dirname is provided, suffix with basename of URL.\n                        raise ValueError(\n                            \"Expecting storage_path to be a file path, got directory {}\".format(\n                                storage_path\n                            )\n                        )\n                    else:\n                        filename = os.path.basename(storage_path)\n\n                    download_url(url=url_or_filename, root=dirname, filename=filename)\n\n    def _download_vis(self):\n\n        storage_path = self.config.build_info.get(self.data_type).storage\n        storage_path = utils.get_cache_path(storage_path)\n\n        if not os.path.exists(storage_path):\n            warnings.warn(\n                f\"\"\"\n                The specified path {storage_path} for visual inputs does not exist.\n                Please provide a correct path to the visual inputs or\n                refer to datasets/download_scripts/README.md for downloading instructions.\n                \"\"\"\n            )\n\n    def build(self):\n        \"\"\"\n        Create by split datasets inheriting torch.utils.data.Datasets.\n\n        # build() can be dataset-specific. Overwrite to customize.\n        \"\"\"\n        self.build_processors()\n\n        build_info = self.config.build_info\n\n        ann_info = build_info.annotations\n        vis_info = build_info.get(self.data_type)\n\n        datasets = dict()\n        for split in ann_info.keys():\n            if split not in [\"train\", \"val\", \"test\"]:\n                continue\n\n            is_train = split == \"train\"\n\n            # processors\n            vis_processor = (\n                self.vis_processors[\"train\"]\n                if is_train\n                else self.vis_processors[\"eval\"]\n            )\n            text_processor = (\n                self.text_processors[\"train\"]\n                if is_train\n                else self.text_processors[\"eval\"]\n            )\n\n            # annotation path\n            ann_paths = ann_info.get(split).storage\n            if isinstance(ann_paths, str):\n                ann_paths = [ann_paths]\n\n            abs_ann_paths = []\n            for ann_path in ann_paths:\n                if not os.path.isabs(ann_path):\n                    ann_path = utils.get_cache_path(ann_path)\n                abs_ann_paths.append(ann_path)\n            ann_paths = abs_ann_paths\n\n            # visual data storage path\n            vis_path = os.path.join(vis_info.storage, split)\n\n            if not os.path.isabs(vis_path):\n                # vis_path = os.path.join(utils.get_cache_path(), vis_path)\n                vis_path = utils.get_cache_path(vis_path)\n\n            if not os.path.exists(vis_path):\n                warnings.warn(\"storage path {} does not exist.\".format(vis_path))\n\n            # create datasets\n            dataset_cls = self.train_dataset_cls if is_train else self.eval_dataset_cls\n            datasets[split] = dataset_cls(\n                vis_processor=vis_processor,\n                text_processor=text_processor,\n                ann_paths=ann_paths,\n                vis_root=vis_path,\n            )\n\n        return datasets\n\n\ndef load_dataset_config(cfg_path):\n    cfg = OmegaConf.load(cfg_path).datasets\n    cfg = cfg[list(cfg.keys())[0]]\n\n    return cfg\n"
  },
  {
    "path": "xraypulse/datasets/builders/image_text_pair_builder.py",
    "content": "import os\nimport logging\nimport warnings\n\nfrom xraypulse.common.registry import registry\nfrom xraypulse.datasets.builders.base_dataset_builder import BaseDatasetBuilder\nfrom xraypulse.datasets.datasets.openi_dataset import OpenIDataset\nfrom xraypulse.datasets.datasets.mimic_dataset import MIMICDataset\n\n\n@registry.register_builder(\"mimic\")\nclass MIMICBuilder(BaseDatasetBuilder):\n    train_dataset_cls = MIMICDataset\n\n    DATASET_CONFIG_DICT = {\"default\": \"configs/datasets/mimic/defaults.yaml\"}\n\n    def _download_ann(self):\n        pass\n\n    def _download_vis(self):\n        pass\n\n    def build_datasets(self):\n        # at this point, all the annotations and image/videos should be all downloaded to the specified locations.\n        logging.info(\"Building datasets...\")\n        self.build_processors()\n\n        build_info = self.config.build_info\n        storage_path = build_info.storage\n\n        datasets = dict()\n\n        if not os.path.exists(storage_path):\n            warnings.warn(\"storage path {} does not exist.\".format(storage_path))\n\n        # create datasets\n        dataset_cls = self.train_dataset_cls\n        datasets['train'] = dataset_cls(\n            vis_processor=self.vis_processors[\"train\"],\n            text_processor=self.text_processors[\"train\"],\n            ann_paths=[os.path.join(storage_path, 'zh_filter_cap.json')],\n            vis_root=os.path.join(storage_path, 'image'),\n        )\n\n        return datasets\n\n\n@registry.register_builder(\"openi\")\nclass OpenIBuilder(BaseDatasetBuilder):\n    train_dataset_cls = OpenIDataset\n\n    DATASET_CONFIG_DICT = {\"default\": \"configs/datasets/openi/defaults.yaml\"}\n\n    def _download_ann(self):\n        pass\n\n    def _download_vis(self):\n        pass\n\n    def build(self):\n        self.build_processors()\n\n        build_info = self.config.build_info\n        storage_path = build_info.storage\n\n        datasets = dict()\n        split = \"train\"\n\n        # create datasets\n        # [NOTE] return inner_datasets (wds.DataPipeline)\n        dataset_cls = self.train_dataset_cls\n        datasets[split] = dataset_cls(\n            vis_processor=self.vis_processors[\"train\"],\n            text_processor=self.text_processors[\"train\"],\n            ann_paths=[os.path.join(storage_path, 'zh_filter_cap.json')],\n            vis_root=os.path.join(storage_path, 'image'),\n        )\n\n        return datasets\n"
  },
  {
    "path": "xraypulse/datasets/data_utils.py",
    "content": "\"\"\"\n Copyright (c) 2022, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\n\nimport gzip\nimport logging\nimport os\nimport random as rnd\nimport tarfile\nimport zipfile\nimport random\nfrom typing import List\nfrom tqdm import tqdm\n\nimport decord\nfrom decord import VideoReader\nimport webdataset as wds\nimport numpy as np\nimport torch\nfrom torch.utils.data.dataset import IterableDataset\n\nfrom xraypulse.common.registry import registry\nfrom xraypulse.datasets.datasets.base_dataset import ConcatDataset\n\n\ndecord.bridge.set_bridge(\"torch\")\nMAX_INT = registry.get(\"MAX_INT\")\n\n\nclass ChainDataset(wds.DataPipeline):\n    r\"\"\"Dataset for chaining multiple :class:`DataPipeline` s.\n\n    This class is useful to assemble different existing dataset streams. The\n    chaining operation is done on-the-fly, so concatenating large-scale\n    datasets with this class will be efficient.\n\n    Args:\n        datasets (iterable of IterableDataset): datasets to be chained together\n    \"\"\"\n    def __init__(self, datasets: List[wds.DataPipeline]) -> None:\n        super().__init__()\n        self.datasets = datasets\n        self.prob = []\n        self.names = []\n        for dataset in self.datasets:\n            if hasattr(dataset, 'name'):\n                self.names.append(dataset.name)\n            else:\n                self.names.append('Unknown')\n            if hasattr(dataset, 'sample_ratio'):\n                self.prob.append(dataset.sample_ratio)\n            else:\n                self.prob.append(1)\n                logging.info(\"One of the datapipeline doesn't define ratio and set to 1 automatically.\")\n\n    def __iter__(self):\n        datastreams = [iter(dataset) for dataset in self.datasets]\n        while True:\n            select_datastream = random.choices(datastreams, weights=self.prob, k=1)[0]\n            yield next(select_datastream)\n\n\ndef apply_to_sample(f, sample):\n    if len(sample) == 0:\n        return {}\n\n    def _apply(x):\n        if torch.is_tensor(x):\n            return f(x)\n        elif isinstance(x, dict):\n            return {key: _apply(value) for key, value in x.items()}\n        elif isinstance(x, list):\n            return [_apply(x) for x in x]\n        else:\n            return x\n\n    return _apply(sample)\n\n\ndef move_to_cuda(sample):\n    def _move_to_cuda(tensor):\n        return tensor.cuda()\n\n    return apply_to_sample(_move_to_cuda, sample)\n\n\ndef prepare_sample(samples, cuda_enabled=True):\n    if cuda_enabled:\n        samples = move_to_cuda(samples)\n\n    # TODO fp16 support\n\n    return samples\n\n\ndef reorg_datasets_by_split(datasets):\n    \"\"\"\n    Organizes datasets by split.\n\n    Args:\n        datasets: dict of torch.utils.data.Dataset objects by name.\n\n    Returns:\n        Dict of datasets by split {split_name: List[Datasets]}.\n    \"\"\"\n    # if len(datasets) == 1:\n    #     return datasets[list(datasets.keys())[0]]\n    # else:\n    reorg_datasets = dict()\n\n    # reorganize by split\n    for _, dataset in datasets.items():\n        for split_name, dataset_split in dataset.items():\n            if split_name not in reorg_datasets:\n                reorg_datasets[split_name] = [dataset_split]\n            else:\n                reorg_datasets[split_name].append(dataset_split)\n\n    return reorg_datasets\n\n\ndef concat_datasets(datasets):\n    \"\"\"\n    Concatenates multiple datasets into a single dataset.\n\n    It supports may-style datasets and DataPipeline from WebDataset. Currently, does not support\n    generic IterableDataset because it requires creating separate samplers.\n\n    Now only supports conctenating training datasets and assuming validation and testing\n    have only a single dataset. This is because metrics should not be computed on the concatenated\n    datasets.\n\n    Args:\n        datasets: dict of torch.utils.data.Dataset objects by split.\n\n    Returns:\n        Dict of concatenated datasets by split, \"train\" is the concatenation of multiple datasets,\n        \"val\" and \"test\" remain the same.\n\n        If the input training datasets contain both map-style and DataPipeline datasets, returns\n        a tuple, where the first element is a concatenated map-style dataset and the second\n        element is a chained DataPipeline dataset.\n\n    \"\"\"\n    # concatenate datasets in the same split\n    for split_name in datasets:\n        if split_name != \"train\":\n            assert (\n                len(datasets[split_name]) == 1\n            ), \"Do not support multiple {} datasets.\".format(split_name)\n            datasets[split_name] = datasets[split_name][0]\n        else:\n            iterable_datasets, map_datasets = [], []\n            for dataset in datasets[split_name]:\n                if isinstance(dataset, wds.DataPipeline):\n                    logging.info(\n                        \"Dataset {} is IterableDataset, can't be concatenated.\".format(\n                            dataset\n                        )\n                    )\n                    iterable_datasets.append(dataset)\n                elif isinstance(dataset, IterableDataset):\n                    raise NotImplementedError(\n                        \"Do not support concatenation of generic IterableDataset.\"\n                    )\n                else:\n                    map_datasets.append(dataset)\n\n            # if len(iterable_datasets) > 0:\n            # concatenate map-style datasets and iterable-style datasets separately\n            if len(iterable_datasets) > 1:\n                chained_datasets = (\n                    ChainDataset(iterable_datasets)\n                )\n            elif len(iterable_datasets) == 1:\n                chained_datasets = iterable_datasets[0]\n            else:\n                chained_datasets = None\n\n            concat_datasets = (\n                ConcatDataset(map_datasets) if len(map_datasets) > 0 else None\n            )\n\n            train_datasets = concat_datasets, chained_datasets\n            train_datasets = tuple([x for x in train_datasets if x is not None])\n            train_datasets = (\n                train_datasets[0] if len(train_datasets) == 1 else train_datasets\n            )\n\n            datasets[split_name] = train_datasets\n\n    return datasets\n\n"
  },
  {
    "path": "xraypulse/datasets/datasets/__init__.py",
    "content": ""
  },
  {
    "path": "xraypulse/datasets/datasets/base_dataset.py",
    "content": "\"\"\"\n Copyright (c) 2022, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\n\nimport json\nfrom typing import Iterable\n\nfrom torch.utils.data import Dataset, ConcatDataset\nfrom torch.utils.data.dataloader import default_collate\n\n\nclass BaseDataset(Dataset):\n    def __init__(\n        self, vis_processor=None, text_processor=None, vis_root=None, ann_paths=[]\n    ):\n        \"\"\"\n        vis_root (string): Root directory of images (e.g. coco/images/)\n        ann_root (string): directory to store the annotation file\n        \"\"\"\n        self.vis_root = vis_root\n\n        self.annotation = []\n        for ann_path in ann_paths:\n            self.annotation.extend(json.load(open(ann_path, \"r\"))['annotations'])\n\n        self.vis_processor = vis_processor\n        self.text_processor = text_processor\n\n        self._add_instance_ids()\n\n    def __len__(self):\n        return len(self.annotation)\n\n    def collater(self, samples):\n        return default_collate(samples)\n\n    def set_processors(self, vis_processor, text_processor):\n        self.vis_processor = vis_processor\n        self.text_processor = text_processor\n\n    def _add_instance_ids(self, key=\"instance_id\"):\n        for idx, ann in enumerate(self.annotation):\n            ann[key] = str(idx)\n\n\nclass ConcatDataset(ConcatDataset):\n    def __init__(self, datasets: Iterable[Dataset]) -> None:\n        super().__init__(datasets)\n\n    def collater(self, samples):\n        # TODO For now only supports datasets with same underlying collater implementations\n\n        all_keys = set()\n        for s in samples:\n            all_keys.update(s)\n\n        shared_keys = all_keys\n        for s in samples:\n            shared_keys = shared_keys & set(s.keys())\n\n        samples_shared_keys = []\n        for s in samples:\n            samples_shared_keys.append({k: s[k] for k in s.keys() if k in shared_keys})\n\n        return self.datasets[0].collater(samples_shared_keys)\n"
  },
  {
    "path": "xraypulse/datasets/datasets/caption_datasets.py",
    "content": "\"\"\"\n Copyright (c) 2022, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\n\nimport os\nfrom collections import OrderedDict\n\nfrom xraypulse.datasets.datasets.base_dataset import BaseDataset\nfrom PIL import Image\n\n\nclass __DisplMixin:\n    def displ_item(self, index):\n        sample, ann = self.__getitem__(index), self.annotation[index]\n\n        return OrderedDict(\n            {\n                \"file\": ann[\"image\"],\n                \"caption\": ann[\"caption\"],\n                \"image\": sample[\"image\"],\n            }\n        )\n\n\nclass CaptionDataset(BaseDataset, __DisplMixin):\n    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):\n        \"\"\"\n        vis_root (string): Root directory of images (e.g. coco/images/)\n        ann_root (string): directory to store the annotation file\n        \"\"\"\n        super().__init__(vis_processor, text_processor, vis_root, ann_paths)\n\n        self.img_ids = {}\n        n = 0\n        for ann in self.annotation:\n            img_id = ann[\"image_id\"]\n            if img_id not in self.img_ids.keys():\n                self.img_ids[img_id] = n\n                n += 1\n\n    def __getitem__(self, index):\n\n        # TODO this assumes image input, not general enough\n        ann = self.annotation[index]\n\n        img_file = '{:0>12}.png'.format(ann[\"image_id\"])\n        image_path = os.path.join(self.vis_root, img_file)\n        image = Image.open(image_path).convert(\"RGB\")\n\n        image = self.vis_processor(image)\n        caption = self.text_processor(ann[\"caption\"])\n\n        return {\n            \"image\": image,\n            \"text_input\": caption,\n            \"image_id\": self.img_ids[ann[\"image_id\"]],\n        }\n\n\nclass CaptionEvalDataset(BaseDataset, __DisplMixin):\n    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):\n        \"\"\"\n        vis_root (string): Root directory of images (e.g. coco/images/)\n        ann_root (string): directory to store the annotation file\n        split (string): val or test\n        \"\"\"\n        super().__init__(vis_processor, text_processor, vis_root, ann_paths)\n\n        #below lines are added during test rogue score\n        self.img_ids = {}\n        n = 0\n        for ann in self.annotation:\n            img_id = ann[\"image_id\"]\n            if img_id not in self.img_ids.keys():\n                self.img_ids[img_id] = n\n                n += 1\n\n    def __getitem__(self, index):\n\n        ann = self.annotation[index]\n\n        image_path = os.path.join(self.vis_root, ann[\"image\"])\n        image = Image.open(image_path).convert(\"RGB\")\n\n        image = self.vis_processor(image)\n\n        return {\n            \"image\": image,\n            \"image_id\": ann[\"image_id\"],\n            \"instance_id\": ann[\"instance_id\"],\n        }\n"
  },
  {
    "path": "xraypulse/datasets/datasets/dataloader_utils.py",
    "content": "\"\"\"\n Copyright (c) 2022, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\n\nimport time\nimport random\nimport torch\nfrom xraypulse.datasets.data_utils import move_to_cuda\nfrom torch.utils.data import DataLoader\n\n\nclass MultiIterLoader:\n    \"\"\"\n    A simple wrapper for iterating over multiple iterators.\n\n    Args:\n        loaders (List[Loader]): List of Iterator loaders.\n        ratios (List[float]): List of ratios to sample from each loader. If None, all loaders are sampled uniformly.\n    \"\"\"\n\n    def __init__(self, loaders, ratios=None):\n        # assert all loaders has __next__ method\n        for loader in loaders:\n            assert hasattr(\n                loader, \"__next__\"\n            ), \"Loader {} has no __next__ method.\".format(loader)\n\n        if ratios is None:\n            ratios = [1.0] * len(loaders)\n        else:\n            assert len(ratios) == len(loaders)\n            ratios = [float(ratio) / sum(ratios) for ratio in ratios]\n\n        self.loaders = loaders\n        self.ratios = ratios\n\n    def __next__(self):\n        # random sample from each loader by ratio\n        loader_idx = random.choices(range(len(self.loaders)), self.ratios, k=1)[0]\n        return next(self.loaders[loader_idx])\n\n\nclass PrefetchLoader(object):\n    \"\"\"\n    Modified from https://github.com/ChenRocks/UNITER.\n\n    overlap compute and cuda data transfer\n    (copied and then modified from nvidia apex)\n    \"\"\"\n\n    def __init__(self, loader):\n        self.loader = loader\n        self.stream = torch.cuda.Stream()\n\n    def __iter__(self):\n        loader_it = iter(self.loader)\n        self.preload(loader_it)\n        batch = self.next(loader_it)\n        while batch is not None:\n            is_tuple = isinstance(batch, tuple)\n            if is_tuple:\n                task, batch = batch\n\n            if is_tuple:\n                yield task, batch\n            else:\n                yield batch\n            batch = self.next(loader_it)\n\n    def __len__(self):\n        return len(self.loader)\n\n    def preload(self, it):\n        try:\n            self.batch = next(it)\n        except StopIteration:\n            self.batch = None\n            return\n        # if record_stream() doesn't work, another option is to make sure\n        # device inputs are created on the main stream.\n        # self.next_input_gpu = torch.empty_like(self.next_input,\n        #                                        device='cuda')\n        # self.next_target_gpu = torch.empty_like(self.next_target,\n        #                                         device='cuda')\n        # Need to make sure the memory allocated for next_* is not still in use\n        # by the main stream at the time we start copying to next_*:\n        # self.stream.wait_stream(torch.cuda.current_stream())\n        with torch.cuda.stream(self.stream):\n            self.batch = move_to_cuda(self.batch)\n            # more code for the alternative if record_stream() doesn't work:\n            # copy_ will record the use of the pinned source tensor in this\n            # side stream.\n            # self.next_input_gpu.copy_(self.next_input, non_blocking=True)\n            # self.next_target_gpu.copy_(self.next_target, non_blocking=True)\n            # self.next_input = self.next_input_gpu\n            # self.next_target = self.next_target_gpu\n\n    def next(self, it):\n        torch.cuda.current_stream().wait_stream(self.stream)\n        batch = self.batch\n        if batch is not None:\n            record_cuda_stream(batch)\n        self.preload(it)\n        return batch\n\n    def __getattr__(self, name):\n        method = self.loader.__getattribute__(name)\n        return method\n\n\ndef record_cuda_stream(batch):\n    if isinstance(batch, torch.Tensor):\n        batch.record_stream(torch.cuda.current_stream())\n    elif isinstance(batch, list) or isinstance(batch, tuple):\n        for t in batch:\n            record_cuda_stream(t)\n    elif isinstance(batch, dict):\n        for t in batch.values():\n            record_cuda_stream(t)\n    else:\n        pass\n\n\nclass IterLoader:\n    \"\"\"\n    A wrapper to convert DataLoader as an infinite iterator.\n\n    Modified from:\n        https://github.com/open-mmlab/mmcv/blob/master/mmcv/runner/iter_based_runner.py\n    \"\"\"\n\n    def __init__(self, dataloader: DataLoader, use_distributed: bool = False):\n        self._dataloader = dataloader\n        self.iter_loader = iter(self._dataloader)\n        self._use_distributed = use_distributed\n        self._epoch = 0\n\n    @property\n    def epoch(self) -> int:\n        return self._epoch\n\n    def __next__(self):\n        try:\n            data = next(self.iter_loader)\n        except StopIteration:\n            self._epoch += 1\n            if hasattr(self._dataloader.sampler, \"set_epoch\") and self._use_distributed:\n                self._dataloader.sampler.set_epoch(self._epoch)\n            time.sleep(2)  # Prevent possible deadlock during epoch transition\n            self.iter_loader = iter(self._dataloader)\n            data = next(self.iter_loader)\n\n        return data\n\n    def __iter__(self):\n        return self\n\n    def __len__(self):\n        return len(self._dataloader)\n"
  },
  {
    "path": "xraypulse/datasets/datasets/mimic_dataset.py",
    "content": "import os\nfrom PIL import Image\nimport webdataset as wds\nfrom xraypulse.datasets.datasets.base_dataset import BaseDataset\nfrom xraypulse.datasets.datasets.caption_datasets import CaptionDataset\n\nclass MIMICDataset(CaptionDataset):\n\n    def __getitem__(self, index):\n\n        # TODO this assumes image input, not general enough\n        ann = self.annotation[index]\n\n        img_file = '{}.jpg'.format(ann[\"image_id\"])\n        image_path = os.path.join(self.vis_root, img_file)\n        image = Image.open(image_path).convert(\"RGB\")\n\n        image = self.vis_processor(image)\n        caption = ann['caption']\n\n        return {\n            \"image\": image,\n            \"caption\":caption,\n            \"image_id\": self.img_ids[ann[\"image_id\"]],\n        }"
  },
  {
    "path": "xraypulse/datasets/datasets/openi_dataset.py",
    "content": "\"\"\"\n Copyright (c) 2022, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\n\nimport os\nfrom PIL import Image\nimport webdataset as wds\nfrom xraypulse.datasets.datasets.base_dataset import BaseDataset\nfrom xraypulse.datasets.datasets.caption_datasets import CaptionDataset\n    \nclass OpenIDataset(CaptionDataset):\n\n    def __getitem__(self, index):\n\n        # TODO this assumes image input, not general enough\n        ann = self.annotation[index]\n\n        img_file = '{}.png'.format(ann[\"image_id\"])\n        image_path = os.path.join(self.vis_root, img_file)\n        image = Image.open(image_path).convert(\"RGB\")\n\n        image = self.vis_processor(image)\n        caption = ann['caption']\n\n        return {\n            \"image\": image,\n            \"caption\":caption,\n            \"image_id\": self.img_ids[ann[\"image_id\"]],\n        }\n\n"
  },
  {
    "path": "xraypulse/models/Qformer.py",
    "content": "\"\"\"\n * Copyright (c) 2023, salesforce.com, inc.\n * All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n * For full license text, see LICENSE.txt file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n * By Junnan Li\n * Based on huggingface code base\n * https://github.com/huggingface/transformers/blob/v4.15.0/src/transformers/models/bert\n\"\"\"\n\nimport math\nimport os\nimport warnings\nfrom dataclasses import dataclass\nfrom typing import Optional, Tuple, Dict, Any\n\nimport torch\nfrom torch import Tensor, device, dtype, nn\nimport torch.utils.checkpoint\nfrom torch import nn\nfrom torch.nn import CrossEntropyLoss\nimport torch.nn.functional as F\n\nfrom transformers.activations import ACT2FN\nfrom transformers.file_utils import (\n    ModelOutput,\n)\nfrom transformers.modeling_outputs import (\n    BaseModelOutputWithPastAndCrossAttentions,\n    BaseModelOutputWithPoolingAndCrossAttentions,\n    CausalLMOutputWithCrossAttentions,\n    MaskedLMOutput,\n    MultipleChoiceModelOutput,\n    NextSentencePredictorOutput,\n    QuestionAnsweringModelOutput,\n    SequenceClassifierOutput,\n    TokenClassifierOutput,\n)\nfrom transformers.modeling_utils import (\n    PreTrainedModel,\n    apply_chunking_to_forward,\n    find_pruneable_heads_and_indices,\n    prune_linear_layer,\n)\nfrom transformers.utils import logging\nfrom transformers.models.bert.configuration_bert import BertConfig\n\nlogger = logging.get_logger(__name__)\n\n\nclass BertEmbeddings(nn.Module):\n    \"\"\"Construct the embeddings from word and position embeddings.\"\"\"\n\n    def __init__(self, config):\n        super().__init__()\n        self.word_embeddings = nn.Embedding(\n            config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id\n        )\n        self.position_embeddings = nn.Embedding(\n            config.max_position_embeddings, config.hidden_size\n        )\n\n        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load\n        # any TensorFlow checkpoint file\n        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)\n        self.dropout = nn.Dropout(config.hidden_dropout_prob)\n\n        # position_ids (1, len position emb) is contiguous in memory and exported when serialized\n        self.register_buffer(\n            \"position_ids\", torch.arange(config.max_position_embeddings).expand((1, -1))\n        )\n        self.position_embedding_type = getattr(\n            config, \"position_embedding_type\", \"absolute\"\n        )\n\n        self.config = config\n\n    def forward(\n        self,\n        input_ids=None,\n        position_ids=None,\n        query_embeds=None,\n        past_key_values_length=0,\n    ):\n        if input_ids is not None:\n            seq_length = input_ids.size()[1]\n        else:\n            seq_length = 0\n\n        if position_ids is None:\n            position_ids = self.position_ids[\n                :, past_key_values_length : seq_length + past_key_values_length\n            ].clone()\n\n        if input_ids is not None:\n            embeddings = self.word_embeddings(input_ids)\n            if self.position_embedding_type == \"absolute\":\n                position_embeddings = self.position_embeddings(position_ids)\n                embeddings = embeddings + position_embeddings\n\n            if query_embeds is not None:\n                embeddings = torch.cat((query_embeds, embeddings), dim=1)\n        else:\n            embeddings = query_embeds\n\n        embeddings = self.LayerNorm(embeddings)\n        embeddings = self.dropout(embeddings)\n        return embeddings\n\n\nclass BertSelfAttention(nn.Module):\n    def __init__(self, config, is_cross_attention):\n        super().__init__()\n        self.config = config\n        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(\n            config, \"embedding_size\"\n        ):\n            raise ValueError(\n                \"The hidden size (%d) is not a multiple of the number of attention \"\n                \"heads (%d)\" % (config.hidden_size, config.num_attention_heads)\n            )\n\n        self.num_attention_heads = config.num_attention_heads\n        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)\n        self.all_head_size = self.num_attention_heads * self.attention_head_size\n\n        self.query = nn.Linear(config.hidden_size, self.all_head_size)\n        if is_cross_attention:\n            self.key = nn.Linear(config.encoder_width, self.all_head_size)\n            self.value = nn.Linear(config.encoder_width, self.all_head_size)\n        else:\n            self.key = nn.Linear(config.hidden_size, self.all_head_size)\n            self.value = nn.Linear(config.hidden_size, self.all_head_size)\n\n        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)\n        self.position_embedding_type = getattr(\n            config, \"position_embedding_type\", \"absolute\"\n        )\n        if (\n            self.position_embedding_type == \"relative_key\"\n            or self.position_embedding_type == \"relative_key_query\"\n        ):\n            self.max_position_embeddings = config.max_position_embeddings\n            self.distance_embedding = nn.Embedding(\n                2 * config.max_position_embeddings - 1, self.attention_head_size\n            )\n        self.save_attention = False\n\n    def save_attn_gradients(self, attn_gradients):\n        self.attn_gradients = attn_gradients\n\n    def get_attn_gradients(self):\n        return self.attn_gradients\n\n    def save_attention_map(self, attention_map):\n        self.attention_map = attention_map\n\n    def get_attention_map(self):\n        return self.attention_map\n\n    def transpose_for_scores(self, x):\n        new_x_shape = x.size()[:-1] + (\n            self.num_attention_heads,\n            self.attention_head_size,\n        )\n        x = x.view(*new_x_shape)\n        return x.permute(0, 2, 1, 3)\n\n    def forward(\n        self,\n        hidden_states,\n        attention_mask=None,\n        head_mask=None,\n        encoder_hidden_states=None,\n        encoder_attention_mask=None,\n        past_key_value=None,\n        output_attentions=False,\n    ):\n\n        # If this is instantiated as a cross-attention module, the keys\n        # and values come from an encoder; the attention mask needs to be\n        # such that the encoder's padding tokens are not attended to.\n        is_cross_attention = encoder_hidden_states is not None\n\n        if is_cross_attention:\n            key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))\n            value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))\n            attention_mask = encoder_attention_mask\n        elif past_key_value is not None:\n            key_layer = self.transpose_for_scores(self.key(hidden_states))\n            value_layer = self.transpose_for_scores(self.value(hidden_states))\n            key_layer = torch.cat([past_key_value[0], key_layer], dim=2)\n            value_layer = torch.cat([past_key_value[1], value_layer], dim=2)\n        else:\n            key_layer = self.transpose_for_scores(self.key(hidden_states))\n            value_layer = self.transpose_for_scores(self.value(hidden_states))\n\n        mixed_query_layer = self.query(hidden_states)\n\n        query_layer = self.transpose_for_scores(mixed_query_layer)\n\n        past_key_value = (key_layer, value_layer)\n\n        # Take the dot product between \"query\" and \"key\" to get the raw attention scores.\n        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))\n\n        if (\n            self.position_embedding_type == \"relative_key\"\n            or self.position_embedding_type == \"relative_key_query\"\n        ):\n            seq_length = hidden_states.size()[1]\n            position_ids_l = torch.arange(\n                seq_length, dtype=torch.long, device=hidden_states.device\n            ).view(-1, 1)\n            position_ids_r = torch.arange(\n                seq_length, dtype=torch.long, device=hidden_states.device\n            ).view(1, -1)\n            distance = position_ids_l - position_ids_r\n            positional_embedding = self.distance_embedding(\n                distance + self.max_position_embeddings - 1\n            )\n            positional_embedding = positional_embedding.to(\n                dtype=query_layer.dtype\n            )  # fp16 compatibility\n\n            if self.position_embedding_type == \"relative_key\":\n                relative_position_scores = torch.einsum(\n                    \"bhld,lrd->bhlr\", query_layer, positional_embedding\n                )\n                attention_scores = attention_scores + relative_position_scores\n            elif self.position_embedding_type == \"relative_key_query\":\n                relative_position_scores_query = torch.einsum(\n                    \"bhld,lrd->bhlr\", query_layer, positional_embedding\n                )\n                relative_position_scores_key = torch.einsum(\n                    \"bhrd,lrd->bhlr\", key_layer, positional_embedding\n                )\n                attention_scores = (\n                    attention_scores\n                    + relative_position_scores_query\n                    + relative_position_scores_key\n                )\n\n        attention_scores = attention_scores / math.sqrt(self.attention_head_size)\n        if attention_mask is not None:\n            # Apply the attention mask is (precomputed for all layers in BertModel forward() function)\n            attention_scores = attention_scores + attention_mask\n\n        # Normalize the attention scores to probabilities.\n        attention_probs = nn.Softmax(dim=-1)(attention_scores)\n\n        if is_cross_attention and self.save_attention:\n            self.save_attention_map(attention_probs)\n            attention_probs.register_hook(self.save_attn_gradients)\n\n        # This is actually dropping out entire tokens to attend to, which might\n        # seem a bit unusual, but is taken from the original Transformer paper.\n        attention_probs_dropped = self.dropout(attention_probs)\n\n        # Mask heads if we want to\n        if head_mask is not None:\n            attention_probs_dropped = attention_probs_dropped * head_mask\n\n        context_layer = torch.matmul(attention_probs_dropped, value_layer)\n\n        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()\n        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)\n        context_layer = context_layer.view(*new_context_layer_shape)\n\n        outputs = (\n            (context_layer, attention_probs) if output_attentions else (context_layer,)\n        )\n\n        outputs = outputs + (past_key_value,)\n        return outputs\n\n\nclass BertSelfOutput(nn.Module):\n    def __init__(self, config):\n        super().__init__()\n        self.dense = nn.Linear(config.hidden_size, config.hidden_size)\n        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)\n        self.dropout = nn.Dropout(config.hidden_dropout_prob)\n\n    def forward(self, hidden_states, input_tensor):\n        hidden_states = self.dense(hidden_states)\n        hidden_states = self.dropout(hidden_states)\n        hidden_states = self.LayerNorm(hidden_states + input_tensor)\n        return hidden_states\n\n\nclass BertAttention(nn.Module):\n    def __init__(self, config, is_cross_attention=False):\n        super().__init__()\n        self.self = BertSelfAttention(config, is_cross_attention)\n        self.output = BertSelfOutput(config)\n        self.pruned_heads = set()\n\n    def prune_heads(self, heads):\n        if len(heads) == 0:\n            return\n        heads, index = find_pruneable_heads_and_indices(\n            heads,\n            self.self.num_attention_heads,\n            self.self.attention_head_size,\n            self.pruned_heads,\n        )\n\n        # Prune linear layers\n        self.self.query = prune_linear_layer(self.self.query, index)\n        self.self.key = prune_linear_layer(self.self.key, index)\n        self.self.value = prune_linear_layer(self.self.value, index)\n        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)\n\n        # Update hyper params and store pruned heads\n        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)\n        self.self.all_head_size = (\n            self.self.attention_head_size * self.self.num_attention_heads\n        )\n        self.pruned_heads = self.pruned_heads.union(heads)\n\n    def forward(\n        self,\n        hidden_states,\n        attention_mask=None,\n        head_mask=None,\n        encoder_hidden_states=None,\n        encoder_attention_mask=None,\n        past_key_value=None,\n        output_attentions=False,\n    ):\n        self_outputs = self.self(\n            hidden_states,\n            attention_mask,\n            head_mask,\n            encoder_hidden_states,\n            encoder_attention_mask,\n            past_key_value,\n            output_attentions,\n        )\n        attention_output = self.output(self_outputs[0], hidden_states)\n\n        outputs = (attention_output,) + self_outputs[\n            1:\n        ]  # add attentions if we output them\n        return outputs\n\n\nclass BertIntermediate(nn.Module):\n    def __init__(self, config):\n        super().__init__()\n        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)\n        if isinstance(config.hidden_act, str):\n            self.intermediate_act_fn = ACT2FN[config.hidden_act]\n        else:\n            self.intermediate_act_fn = config.hidden_act\n\n    def forward(self, hidden_states):\n        hidden_states = self.dense(hidden_states)\n        hidden_states = self.intermediate_act_fn(hidden_states)\n        return hidden_states\n\n\nclass BertOutput(nn.Module):\n    def __init__(self, config):\n        super().__init__()\n        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)\n        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)\n        self.dropout = nn.Dropout(config.hidden_dropout_prob)\n\n    def forward(self, hidden_states, input_tensor):\n        hidden_states = self.dense(hidden_states)\n        hidden_states = self.dropout(hidden_states)\n        hidden_states = self.LayerNorm(hidden_states + input_tensor)\n        return hidden_states\n\n\nclass BertLayer(nn.Module):\n    def __init__(self, config, layer_num):\n        super().__init__()\n        self.config = config\n        self.chunk_size_feed_forward = config.chunk_size_feed_forward\n        self.seq_len_dim = 1\n        self.attention = BertAttention(config)\n        self.layer_num = layer_num\n        if (\n            self.config.add_cross_attention\n            and layer_num % self.config.cross_attention_freq == 0\n        ):\n            self.crossattention = BertAttention(\n                config, is_cross_attention=self.config.add_cross_attention\n            )\n            self.has_cross_attention = True\n        else:\n            self.has_cross_attention = False\n        self.intermediate = BertIntermediate(config)\n        self.output = BertOutput(config)\n\n        self.intermediate_query = BertIntermediate(config)\n        self.output_query = BertOutput(config)\n\n    def forward(\n        self,\n        hidden_states,\n        attention_mask=None,\n        head_mask=None,\n        encoder_hidden_states=None,\n        encoder_attention_mask=None,\n        past_key_value=None,\n        output_attentions=False,\n        query_length=0,\n    ):\n        # decoder uni-directional self-attention cached key/values tuple is at positions 1,2\n        self_attn_past_key_value = (\n            past_key_value[:2] if past_key_value is not None else None\n        )\n        self_attention_outputs = self.attention(\n            hidden_states,\n            attention_mask,\n            head_mask,\n            output_attentions=output_attentions,\n            past_key_value=self_attn_past_key_value,\n        )\n        attention_output = self_attention_outputs[0]\n        outputs = self_attention_outputs[1:-1]\n\n        present_key_value = self_attention_outputs[-1]\n\n        if query_length > 0:\n            query_attention_output = attention_output[:, :query_length, :]\n\n            if self.has_cross_attention:\n                assert (\n                    encoder_hidden_states is not None\n                ), \"encoder_hidden_states must be given for cross-attention layers\"\n                cross_attention_outputs = self.crossattention(\n                    query_attention_output,\n                    attention_mask,\n                    head_mask,\n                    encoder_hidden_states,\n                    encoder_attention_mask,\n                    output_attentions=output_attentions,\n                )\n                query_attention_output = cross_attention_outputs[0]\n                outputs = (\n                    outputs + cross_attention_outputs[1:-1]\n                )  # add cross attentions if we output attention weights\n\n            layer_output = apply_chunking_to_forward(\n                self.feed_forward_chunk_query,\n                self.chunk_size_feed_forward,\n                self.seq_len_dim,\n                query_attention_output,\n            )\n            if attention_output.shape[1] > query_length:\n                layer_output_text = apply_chunking_to_forward(\n                    self.feed_forward_chunk,\n                    self.chunk_size_feed_forward,\n                    self.seq_len_dim,\n                    attention_output[:, query_length:, :],\n                )\n                layer_output = torch.cat([layer_output, layer_output_text], dim=1)\n        else:\n            layer_output = apply_chunking_to_forward(\n                self.feed_forward_chunk,\n                self.chunk_size_feed_forward,\n                self.seq_len_dim,\n                attention_output,\n            )\n        outputs = (layer_output,) + outputs\n\n        outputs = outputs + (present_key_value,)\n\n        return outputs\n\n    def feed_forward_chunk(self, attention_output):\n        intermediate_output = self.intermediate(attention_output)\n        layer_output = self.output(intermediate_output, attention_output)\n        return layer_output\n\n    def feed_forward_chunk_query(self, attention_output):\n        intermediate_output = self.intermediate_query(attention_output)\n        layer_output = self.output_query(intermediate_output, attention_output)\n        return layer_output\n\n\nclass BertEncoder(nn.Module):\n    def __init__(self, config):\n        super().__init__()\n        self.config = config\n        self.layer = nn.ModuleList(\n            [BertLayer(config, i) for i in range(config.num_hidden_layers)]\n        )\n\n    def forward(\n        self,\n        hidden_states,\n        attention_mask=None,\n        head_mask=None,\n        encoder_hidden_states=None,\n        encoder_attention_mask=None,\n        past_key_values=None,\n        use_cache=None,\n        output_attentions=False,\n        output_hidden_states=False,\n        return_dict=True,\n        query_length=0,\n    ):\n        all_hidden_states = () if output_hidden_states else None\n        all_self_attentions = () if output_attentions else None\n        all_cross_attentions = (\n            () if output_attentions and self.config.add_cross_attention else None\n        )\n\n        next_decoder_cache = () if use_cache else None\n\n        for i in range(self.config.num_hidden_layers):\n            layer_module = self.layer[i]\n            if output_hidden_states:\n                all_hidden_states = all_hidden_states + (hidden_states,)\n\n            layer_head_mask = head_mask[i] if head_mask is not None else None\n            past_key_value = past_key_values[i] if past_key_values is not None else None\n\n            if getattr(self.config, \"gradient_checkpointing\", False) and self.training:\n\n                if use_cache:\n                    logger.warn(\n                        \"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...\"\n                    )\n                    use_cache = False\n\n                def create_custom_forward(module):\n                    def custom_forward(*inputs):\n                        return module(\n                            *inputs, past_key_value, output_attentions, query_length\n                        )\n\n                    return custom_forward\n\n                layer_outputs = torch.utils.checkpoint.checkpoint(\n                    create_custom_forward(layer_module),\n                    hidden_states,\n                    attention_mask,\n                    layer_head_mask,\n                    encoder_hidden_states,\n                    encoder_attention_mask,\n                )\n            else:\n                layer_outputs = layer_module(\n                    hidden_states,\n                    attention_mask,\n                    layer_head_mask,\n                    encoder_hidden_states,\n                    encoder_attention_mask,\n                    past_key_value,\n                    output_attentions,\n                    query_length,\n                )\n\n            hidden_states = layer_outputs[0]\n            if use_cache:\n                next_decoder_cache += (layer_outputs[-1],)\n            if output_attentions:\n                all_self_attentions = all_self_attentions + (layer_outputs[1],)\n                all_cross_attentions = all_cross_attentions + (layer_outputs[2],)\n\n        if output_hidden_states:\n            all_hidden_states = all_hidden_states + (hidden_states,)\n\n        if not return_dict:\n            return tuple(\n                v\n                for v in [\n                    hidden_states,\n                    next_decoder_cache,\n                    all_hidden_states,\n                    all_self_attentions,\n                    all_cross_attentions,\n                ]\n                if v is not None\n            )\n        return BaseModelOutputWithPastAndCrossAttentions(\n            last_hidden_state=hidden_states,\n            past_key_values=next_decoder_cache,\n            hidden_states=all_hidden_states,\n            attentions=all_self_attentions,\n            cross_attentions=all_cross_attentions,\n        )\n\n\nclass BertPooler(nn.Module):\n    def __init__(self, config):\n        super().__init__()\n        self.dense = nn.Linear(config.hidden_size, config.hidden_size)\n        self.activation = nn.Tanh()\n\n    def forward(self, hidden_states):\n        # We \"pool\" the model by simply taking the hidden state corresponding\n        # to the first token.\n        first_token_tensor = hidden_states[:, 0]\n        pooled_output = self.dense(first_token_tensor)\n        pooled_output = self.activation(pooled_output)\n        return pooled_output\n\n\nclass BertPredictionHeadTransform(nn.Module):\n    def __init__(self, config):\n        super().__init__()\n        self.dense = nn.Linear(config.hidden_size, config.hidden_size)\n        if isinstance(config.hidden_act, str):\n            self.transform_act_fn = ACT2FN[config.hidden_act]\n        else:\n            self.transform_act_fn = config.hidden_act\n        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)\n\n    def forward(self, hidden_states):\n        hidden_states = self.dense(hidden_states)\n        hidden_states = self.transform_act_fn(hidden_states)\n        hidden_states = self.LayerNorm(hidden_states)\n        return hidden_states\n\n\nclass BertLMPredictionHead(nn.Module):\n    def __init__(self, config):\n        super().__init__()\n        self.transform = BertPredictionHeadTransform(config)\n\n        # The output weights are the same as the input embeddings, but there is\n        # an output-only bias for each token.\n        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)\n\n        self.bias = nn.Parameter(torch.zeros(config.vocab_size))\n\n        # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`\n        self.decoder.bias = self.bias\n\n    def forward(self, hidden_states):\n        hidden_states = self.transform(hidden_states)\n        hidden_states = self.decoder(hidden_states)\n        return hidden_states\n\n\nclass BertOnlyMLMHead(nn.Module):\n    def __init__(self, config):\n        super().__init__()\n        self.predictions = BertLMPredictionHead(config)\n\n    def forward(self, sequence_output):\n        prediction_scores = self.predictions(sequence_output)\n        return prediction_scores\n\n\nclass BertPreTrainedModel(PreTrainedModel):\n    \"\"\"\n    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained\n    models.\n    \"\"\"\n\n    config_class = BertConfig\n    base_model_prefix = \"bert\"\n    _keys_to_ignore_on_load_missing = [r\"position_ids\"]\n\n    def _init_weights(self, module):\n        \"\"\"Initialize the weights\"\"\"\n        if isinstance(module, (nn.Linear, nn.Embedding)):\n            # Slightly different from the TF version which uses truncated_normal for initialization\n            # cf https://github.com/pytorch/pytorch/pull/5617\n            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)\n        elif isinstance(module, nn.LayerNorm):\n            module.bias.data.zero_()\n            module.weight.data.fill_(1.0)\n        if isinstance(module, nn.Linear) and module.bias is not None:\n            module.bias.data.zero_()\n\n\nclass BertModel(BertPreTrainedModel):\n    \"\"\"\n    The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of\n    cross-attention is added between the self-attention layers, following the architecture described in `Attention is\n    all you need <https://arxiv.org/abs/1706.03762>`__ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,\n    Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.\n    argument and :obj:`add_cross_attention` set to :obj:`True`; an :obj:`encoder_hidden_states` is then expected as an\n    input to the forward pass.\n    \"\"\"\n\n    def __init__(self, config, add_pooling_layer=False):\n        super().__init__(config)\n        self.config = config\n\n        self.embeddings = BertEmbeddings(config)\n\n        self.encoder = BertEncoder(config)\n\n        self.pooler = BertPooler(config) if add_pooling_layer else None\n\n        self.init_weights()\n\n    def get_input_embeddings(self):\n        return self.embeddings.word_embeddings\n\n    def set_input_embeddings(self, value):\n        self.embeddings.word_embeddings = value\n\n    def _prune_heads(self, heads_to_prune):\n        \"\"\"\n        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base\n        class PreTrainedModel\n        \"\"\"\n        for layer, heads in heads_to_prune.items():\n            self.encoder.layer[layer].attention.prune_heads(heads)\n\n    def get_extended_attention_mask(\n        self,\n        attention_mask: Tensor,\n        input_shape: Tuple[int],\n        device: device,\n        is_decoder: bool,\n        has_query: bool = False,\n    ) -> Tensor:\n        \"\"\"\n        Makes broadcastable attention and causal masks so that future and masked tokens are ignored.\n\n        Arguments:\n            attention_mask (:obj:`torch.Tensor`):\n                Mask with ones indicating tokens to attend to, zeros for tokens to ignore.\n            input_shape (:obj:`Tuple[int]`):\n                The shape of the input to the model.\n            device: (:obj:`torch.device`):\n                The device of the input to the model.\n\n        Returns:\n            :obj:`torch.Tensor` The extended attention mask, with a the same dtype as :obj:`attention_mask.dtype`.\n        \"\"\"\n        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]\n        # ourselves in which case we just need to make it broadcastable to all heads.\n        if attention_mask.dim() == 3:\n            extended_attention_mask = attention_mask[:, None, :, :]\n        elif attention_mask.dim() == 2:\n            # Provided a padding mask of dimensions [batch_size, seq_length]\n            # - if the model is a decoder, apply a causal mask in addition to the padding mask\n            # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]\n            if is_decoder:\n                batch_size, seq_length = input_shape\n\n                seq_ids = torch.arange(seq_length, device=device)\n                causal_mask = (\n                    seq_ids[None, None, :].repeat(batch_size, seq_length, 1)\n                    <= seq_ids[None, :, None]\n                )\n\n                # add a prefix ones mask to the causal mask\n                # causal and attention masks must have same type with pytorch version < 1.3\n                causal_mask = causal_mask.to(attention_mask.dtype)\n\n                if causal_mask.shape[1] < attention_mask.shape[1]:\n                    prefix_seq_len = attention_mask.shape[1] - causal_mask.shape[1]\n                    if has_query:  # UniLM style attention mask\n                        causal_mask = torch.cat(\n                            [\n                                torch.zeros(\n                                    (batch_size, prefix_seq_len, seq_length),\n                                    device=device,\n                                    dtype=causal_mask.dtype,\n                                ),\n                                causal_mask,\n                            ],\n                            axis=1,\n                        )\n                    causal_mask = torch.cat(\n                        [\n                            torch.ones(\n                                (batch_size, causal_mask.shape[1], prefix_seq_len),\n                                device=device,\n                                dtype=causal_mask.dtype,\n                            ),\n                            causal_mask,\n                        ],\n                        axis=-1,\n                    )\n                extended_attention_mask = (\n                    causal_mask[:, None, :, :] * attention_mask[:, None, None, :]\n                )\n            else:\n                extended_attention_mask = attention_mask[:, None, None, :]\n        else:\n            raise ValueError(\n                \"Wrong shape for input_ids (shape {}) or attention_mask (shape {})\".format(\n                    input_shape, attention_mask.shape\n                )\n            )\n\n        # Since attention_mask is 1.0 for positions we want to attend and 0.0 for\n        # masked positions, this operation will create a tensor which is 0.0 for\n        # positions we want to attend and -10000.0 for masked positions.\n        # Since we are adding it to the raw scores before the softmax, this is\n        # effectively the same as removing these entirely.\n        extended_attention_mask = extended_attention_mask.to(\n            dtype=self.dtype\n        )  # fp16 compatibility\n        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0\n        return extended_attention_mask\n\n    def forward(\n        self,\n        input_ids=None,\n        attention_mask=None,\n        position_ids=None,\n        head_mask=None,\n        query_embeds=None,\n        encoder_hidden_states=None,\n        encoder_attention_mask=None,\n        past_key_values=None,\n        use_cache=None,\n        output_attentions=None,\n        output_hidden_states=None,\n        return_dict=None,\n        is_decoder=False,\n    ):\n        r\"\"\"\n        encoder_hidden_states  (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):\n            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if\n            the model is configured as a decoder.\n        encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):\n            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in\n            the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``:\n            - 1 for tokens that are **not masked**,\n            - 0 for tokens that are **masked**.\n        past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):\n            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.\n            If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids`\n            (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)`\n            instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`.\n        use_cache (:obj:`bool`, `optional`):\n            If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up\n            decoding (see :obj:`past_key_values`).\n        \"\"\"\n        output_attentions = (\n            output_attentions\n            if output_attentions is not None\n            else self.config.output_attentions\n        )\n        output_hidden_states = (\n            output_hidden_states\n            if output_hidden_states is not None\n            else self.config.output_hidden_states\n        )\n        return_dict = (\n            return_dict if return_dict is not None else self.config.use_return_dict\n        )\n\n        # use_cache = use_cache if use_cache is not None else self.config.use_cache\n\n        if input_ids is None:\n            assert (\n                query_embeds is not None\n            ), \"You have to specify query_embeds when input_ids is None\"\n\n        # past_key_values_length\n        past_key_values_length = (\n            past_key_values[0][0].shape[2] - self.config.query_length\n            if past_key_values is not None\n            else 0\n        )\n\n        query_length = query_embeds.shape[1] if query_embeds is not None else 0\n\n        embedding_output = self.embeddings(\n            input_ids=input_ids,\n            position_ids=position_ids,\n            query_embeds=query_embeds,\n            past_key_values_length=past_key_values_length,\n        )\n\n        input_shape = embedding_output.size()[:-1]\n        batch_size, seq_length = input_shape\n        device = embedding_output.device\n\n        if attention_mask is None:\n            attention_mask = torch.ones(\n                ((batch_size, seq_length + past_key_values_length)), device=device\n            )\n\n        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]\n        # ourselves in which case we just need to make it broadcastable to all heads.\n        if is_decoder:\n            extended_attention_mask = self.get_extended_attention_mask(\n                attention_mask,\n                input_ids.shape,\n                device,\n                is_decoder,\n                has_query=(query_embeds is not None),\n            )\n        else:\n            extended_attention_mask = self.get_extended_attention_mask(\n                attention_mask, input_shape, device, is_decoder\n            )\n\n        # If a 2D or 3D attention mask is provided for the cross-attention\n        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]\n        if encoder_hidden_states is not None:\n            if type(encoder_hidden_states) == list:\n                encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states[\n                    0\n                ].size()\n            else:\n                (\n                    encoder_batch_size,\n                    encoder_sequence_length,\n                    _,\n                ) = encoder_hidden_states.size()\n            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)\n\n            if type(encoder_attention_mask) == list:\n                encoder_extended_attention_mask = [\n                    self.invert_attention_mask(mask) for mask in encoder_attention_mask\n                ]\n            elif encoder_attention_mask is None:\n                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)\n                encoder_extended_attention_mask = self.invert_attention_mask(\n                    encoder_attention_mask\n                )\n            else:\n                encoder_extended_attention_mask = self.invert_attention_mask(\n                    encoder_attention_mask\n                )\n        else:\n            encoder_extended_attention_mask = None\n\n        # Prepare head mask if needed\n        # 1.0 in head_mask indicate we keep the head\n        # attention_probs has shape bsz x n_heads x N x N\n        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]\n        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]\n        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)\n\n        encoder_outputs = self.encoder(\n            embedding_output,\n            attention_mask=extended_attention_mask,\n            head_mask=head_mask,\n            encoder_hidden_states=encoder_hidden_states,\n            encoder_attention_mask=encoder_extended_attention_mask,\n            past_key_values=past_key_values,\n            use_cache=use_cache,\n            output_attentions=output_attentions,\n            output_hidden_states=output_hidden_states,\n            return_dict=return_dict,\n            query_length=query_length,\n        )\n        sequence_output = encoder_outputs[0]\n        pooled_output = (\n            self.pooler(sequence_output) if self.pooler is not None else None\n        )\n\n        if not return_dict:\n            return (sequence_output, pooled_output) + encoder_outputs[1:]\n\n        return BaseModelOutputWithPoolingAndCrossAttentions(\n            last_hidden_state=sequence_output,\n            pooler_output=pooled_output,\n            past_key_values=encoder_outputs.past_key_values,\n            hidden_states=encoder_outputs.hidden_states,\n            attentions=encoder_outputs.attentions,\n            cross_attentions=encoder_outputs.cross_attentions,\n        )\n\n\nclass BertLMHeadModel(BertPreTrainedModel):\n\n    _keys_to_ignore_on_load_unexpected = [r\"pooler\"]\n    _keys_to_ignore_on_load_missing = [r\"position_ids\", r\"predictions.decoder.bias\"]\n\n    def __init__(self, config):\n        super().__init__(config)\n\n        self.bert = BertModel(config, add_pooling_layer=False)\n        self.cls = BertOnlyMLMHead(config)\n\n        self.init_weights()\n\n    def get_output_embeddings(self):\n        return self.cls.predictions.decoder\n\n    def set_output_embeddings(self, new_embeddings):\n        self.cls.predictions.decoder = new_embeddings\n\n    def forward(\n        self,\n        input_ids=None,\n        attention_mask=None,\n        position_ids=None,\n        head_mask=None,\n        query_embeds=None,\n        encoder_hidden_states=None,\n        encoder_attention_mask=None,\n        labels=None,\n        past_key_values=None,\n        use_cache=True,\n        output_attentions=None,\n        output_hidden_states=None,\n        return_dict=None,\n        return_logits=False,\n        is_decoder=True,\n        reduction=\"mean\",\n    ):\n        r\"\"\"\n        encoder_hidden_states  (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):\n            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if\n            the model is configured as a decoder.\n        encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):\n            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in\n            the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``:\n            - 1 for tokens that are **not masked**,\n            - 0 for tokens that are **masked**.\n        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):\n            Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in\n            ``[-100, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring) Tokens with indices set to ``-100`` are\n            ignored (masked), the loss is only computed for the tokens with labels n ``[0, ..., config.vocab_size]``\n        past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):\n            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.\n            If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids`\n            (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)`\n            instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`.\n        use_cache (:obj:`bool`, `optional`):\n            If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up\n            decoding (see :obj:`past_key_values`).\n        Returns:\n        Example::\n            >>> from transformers import BertTokenizer, BertLMHeadModel, BertConfig\n            >>> import torch\n            >>> tokenizer = BertTokenizer.from_pretrained('bert-base-cased')\n            >>> config = BertConfig.from_pretrained(\"bert-base-cased\")\n            >>> model = BertLMHeadModel.from_pretrained('bert-base-cased', config=config)\n            >>> inputs = tokenizer(\"Hello, my dog is cute\", return_tensors=\"pt\")\n            >>> outputs = model(**inputs)\n            >>> prediction_logits = outputs.logits\n        \"\"\"\n        return_dict = (\n            return_dict if return_dict is not None else self.config.use_return_dict\n        )\n        if labels is not None:\n            use_cache = False\n        if past_key_values is not None:\n            query_embeds = None\n\n        outputs = self.bert(\n            input_ids,\n            attention_mask=attention_mask,\n            position_ids=position_ids,\n            head_mask=head_mask,\n            query_embeds=query_embeds,\n            encoder_hidden_states=encoder_hidden_states,\n            encoder_attention_mask=encoder_attention_mask,\n            past_key_values=past_key_values,\n            use_cache=use_cache,\n            output_attentions=output_attentions,\n            output_hidden_states=output_hidden_states,\n            return_dict=return_dict,\n            is_decoder=is_decoder,\n        )\n\n        sequence_output = outputs[0]\n        if query_embeds is not None:\n            sequence_output = outputs[0][:, query_embeds.shape[1] :, :]\n\n        prediction_scores = self.cls(sequence_output)\n\n        if return_logits:\n            return prediction_scores[:, :-1, :].contiguous()\n\n        lm_loss = None\n        if labels is not None:\n            # we are doing next-token prediction; shift prediction scores and input ids by one\n            shifted_prediction_scores = prediction_scores[:, :-1, :].contiguous()\n            labels = labels[:, 1:].contiguous()\n            loss_fct = CrossEntropyLoss(reduction=reduction, label_smoothing=0.1)\n            lm_loss = loss_fct(\n                shifted_prediction_scores.view(-1, self.config.vocab_size),\n                labels.view(-1),\n            )\n            if reduction == \"none\":\n                lm_loss = lm_loss.view(prediction_scores.size(0), -1).sum(1)\n\n        if not return_dict:\n            output = (prediction_scores,) + outputs[2:]\n            return ((lm_loss,) + output) if lm_loss is not None else output\n\n        return CausalLMOutputWithCrossAttentions(\n            loss=lm_loss,\n            logits=prediction_scores,\n            past_key_values=outputs.past_key_values,\n            hidden_states=outputs.hidden_states,\n            attentions=outputs.attentions,\n            cross_attentions=outputs.cross_attentions,\n        )\n\n    def prepare_inputs_for_generation(\n        self, input_ids, query_embeds, past=None, attention_mask=None, **model_kwargs\n    ):\n        # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly\n        if attention_mask is None:\n            attention_mask = input_ids.new_ones(input_ids.shape)\n        query_mask = input_ids.new_ones(query_embeds.shape[:-1])\n        attention_mask = torch.cat([query_mask, attention_mask], dim=-1)\n\n        # cut decoder_input_ids if past is used\n        if past is not None:\n            input_ids = input_ids[:, -1:]\n\n        return {\n            \"input_ids\": input_ids,\n            \"query_embeds\": query_embeds,\n            \"attention_mask\": attention_mask,\n            \"past_key_values\": past,\n            \"encoder_hidden_states\": model_kwargs.get(\"encoder_hidden_states\", None),\n            \"encoder_attention_mask\": model_kwargs.get(\"encoder_attention_mask\", None),\n            \"is_decoder\": True,\n        }\n\n    def _reorder_cache(self, past, beam_idx):\n        reordered_past = ()\n        for layer_past in past:\n            reordered_past += (\n                tuple(\n                    past_state.index_select(0, beam_idx) for past_state in layer_past\n                ),\n            )\n        return reordered_past\n\n\nclass BertForMaskedLM(BertPreTrainedModel):\n\n    _keys_to_ignore_on_load_unexpected = [r\"pooler\"]\n    _keys_to_ignore_on_load_missing = [r\"position_ids\", r\"predictions.decoder.bias\"]\n\n    def __init__(self, config):\n        super().__init__(config)\n\n        self.bert = BertModel(config, add_pooling_layer=False)\n        self.cls = BertOnlyMLMHead(config)\n\n        self.init_weights()\n\n    def get_output_embeddings(self):\n        return self.cls.predictions.decoder\n\n    def set_output_embeddings(self, new_embeddings):\n        self.cls.predictions.decoder = new_embeddings\n\n    def forward(\n        self,\n        input_ids=None,\n        attention_mask=None,\n        position_ids=None,\n        head_mask=None,\n        query_embeds=None,\n        encoder_hidden_states=None,\n        encoder_attention_mask=None,\n        labels=None,\n        output_attentions=None,\n        output_hidden_states=None,\n        return_dict=None,\n        return_logits=False,\n        is_decoder=False,\n    ):\n        r\"\"\"\n        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):\n            Labels for computing the masked language modeling loss. Indices should be in ``[-100, 0, ...,\n            config.vocab_size]`` (see ``input_ids`` docstring) Tokens with indices set to ``-100`` are ignored\n            (masked), the loss is only computed for the tokens with labels in ``[0, ..., config.vocab_size]``\n        \"\"\"\n\n        return_dict = (\n            return_dict if return_dict is not None else self.config.use_return_dict\n        )\n\n        outputs = self.bert(\n            input_ids,\n            attention_mask=attention_mask,\n            position_ids=position_ids,\n            head_mask=head_mask,\n            query_embeds=query_embeds,\n            encoder_hidden_states=encoder_hidden_states,\n            encoder_attention_mask=encoder_attention_mask,\n            output_attentions=output_attentions,\n            output_hidden_states=output_hidden_states,\n            return_dict=return_dict,\n            is_decoder=is_decoder,\n        )\n\n        if query_embeds is not None:\n            sequence_output = outputs[0][:, query_embeds.shape[1] :, :]\n        prediction_scores = self.cls(sequence_output)\n\n        if return_logits:\n            return prediction_scores\n\n        masked_lm_loss = None\n        if labels is not None:\n            loss_fct = CrossEntropyLoss()  # -100 index = padding token\n            masked_lm_loss = loss_fct(\n                prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)\n            )\n\n        if not return_dict:\n            output = (prediction_scores,) + outputs[2:]\n            return (\n                ((masked_lm_loss,) + output) if masked_lm_loss is not None else output\n            )\n\n        return MaskedLMOutput(\n            loss=masked_lm_loss,\n            logits=prediction_scores,\n            hidden_states=outputs.hidden_states,\n            attentions=outputs.attentions,\n        )\n"
  },
  {
    "path": "xraypulse/models/__init__.py",
    "content": "\"\"\"\n Copyright (c) 2022, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\n\nimport logging\nimport torch\nfrom omegaconf import OmegaConf\n\nfrom xraypulse.common.registry import registry\nfrom xraypulse.models.base_model import BaseModel\nfrom xraypulse.models.blip2 import Blip2Base\nfrom xraypulse.models.xray_pulse import XrayPulse\nfrom xraypulse.processors.base_processor import BaseProcessor\n\n\n__all__ = [\n    \"load_model\",\n    \"BaseModel\",\n    \"Blip2Base\",\n    \"XrayPulse\",\n]\n\n\ndef load_model(name, model_type, is_eval=False, device=\"cpu\", checkpoint=None):\n    \"\"\"\n    Load supported models.\n\n    To list all available models and types in registry:\n    >>> from minigpt4.models import model_zoo\n    >>> print(model_zoo)\n\n    Args:\n        name (str): name of the model.\n        model_type (str): type of the model.\n        is_eval (bool): whether the model is in eval mode. Default: False.\n        device (str): device to use. Default: \"cpu\".\n        checkpoint (str): path or to checkpoint. Default: None.\n            Note that expecting the checkpoint to have the same keys in state_dict as the model.\n\n    Returns:\n        model (torch.nn.Module): model.\n    \"\"\"\n\n    model = registry.get_model_class(name).from_pretrained(model_type=model_type)\n\n    if checkpoint is not None:\n        model.load_checkpoint(checkpoint)\n\n    if is_eval:\n        model.eval()\n\n    if device == \"cpu\":\n        model = model.float()\n\n    return model.to(device)\n\n\ndef load_preprocess(config):\n    \"\"\"\n    Load preprocessor configs and construct preprocessors.\n\n    If no preprocessor is specified, return BaseProcessor, which does not do any preprocessing.\n\n    Args:\n        config (dict): preprocessor configs.\n\n    Returns:\n        vis_processors (dict): preprocessors for visual inputs.\n        txt_processors (dict): preprocessors for text inputs.\n\n        Key is \"train\" or \"eval\" for processors used in training and evaluation respectively.\n    \"\"\"\n\n    def _build_proc_from_cfg(cfg):\n        return (\n            registry.get_processor_class(cfg.name).from_config(cfg)\n            if cfg is not None\n            else BaseProcessor()\n        )\n\n    vis_processors = dict()\n    txt_processors = dict()\n\n    vis_proc_cfg = config.get(\"vis_processor\")\n    txt_proc_cfg = config.get(\"text_processor\")\n\n    if vis_proc_cfg is not None:\n        vis_train_cfg = vis_proc_cfg.get(\"train\")\n        vis_eval_cfg = vis_proc_cfg.get(\"eval\")\n    else:\n        vis_train_cfg = None\n        vis_eval_cfg = None\n\n    vis_processors[\"train\"] = _build_proc_from_cfg(vis_train_cfg)\n    vis_processors[\"eval\"] = _build_proc_from_cfg(vis_eval_cfg)\n\n    if txt_proc_cfg is not None:\n        txt_train_cfg = txt_proc_cfg.get(\"train\")\n        txt_eval_cfg = txt_proc_cfg.get(\"eval\")\n    else:\n        txt_train_cfg = None\n        txt_eval_cfg = None\n\n    txt_processors[\"train\"] = _build_proc_from_cfg(txt_train_cfg)\n    txt_processors[\"eval\"] = _build_proc_from_cfg(txt_eval_cfg)\n\n    return vis_processors, txt_processors\n\n\ndef load_model_and_preprocess(name, model_type, is_eval=False, device=\"cpu\"):\n    \"\"\"\n    Load model and its related preprocessors.\n\n    List all available models and types in registry:\n    >>> from minigpt4.models import model_zoo\n    >>> print(model_zoo)\n\n    Args:\n        name (str): name of the model.\n        model_type (str): type of the model.\n        is_eval (bool): whether the model is in eval mode. Default: False.\n        device (str): device to use. Default: \"cpu\".\n\n    Returns:\n        model (torch.nn.Module): model.\n        vis_processors (dict): preprocessors for visual inputs.\n        txt_processors (dict): preprocessors for text inputs.\n    \"\"\"\n    model_cls = registry.get_model_class(name)\n\n    # load model\n    model = model_cls.from_pretrained(model_type=model_type)\n\n    if is_eval:\n        model.eval()\n\n    # load preprocess\n    cfg = OmegaConf.load(model_cls.default_config_path(model_type))\n    if cfg is not None:\n        preprocess_cfg = cfg.preprocess\n\n        vis_processors, txt_processors = load_preprocess(preprocess_cfg)\n    else:\n        vis_processors, txt_processors = None, None\n        logging.info(\n            f\"\"\"No default preprocess for model {name} ({model_type}).\n                This can happen if the model is not finetuned on downstream datasets,\n                or it is not intended for direct use without finetuning.\n            \"\"\"\n        )\n\n    if device == \"cpu\" or device == torch.device(\"cpu\"):\n        model = model.float()\n\n    return model.to(device), vis_processors, txt_processors\n\n\nclass ModelZoo:\n    \"\"\"\n    A utility class to create string representation of available model architectures and types.\n\n    >>> from minigpt4.models import model_zoo\n    >>> # list all available models\n    >>> print(model_zoo)\n    >>> # show total number of models\n    >>> print(len(model_zoo))\n    \"\"\"\n\n    def __init__(self) -> None:\n        self.model_zoo = {\n            k: list(v.PRETRAINED_MODEL_CONFIG_DICT.keys())\n            for k, v in registry.mapping[\"model_name_mapping\"].items()\n        }\n\n    def __str__(self) -> str:\n        return (\n            \"=\" * 50\n            + \"\\n\"\n            + f\"{'Architectures':<30} {'Types'}\\n\"\n            + \"=\" * 50\n            + \"\\n\"\n            + \"\\n\".join(\n                [\n                    f\"{name:<30} {', '.join(types)}\"\n                    for name, types in self.model_zoo.items()\n                ]\n            )\n        )\n\n    def __iter__(self):\n        return iter(self.model_zoo.items())\n\n    def __len__(self):\n        return sum([len(v) for v in self.model_zoo.values()])\n\n\nmodel_zoo = ModelZoo()\n"
  },
  {
    "path": "xraypulse/models/base_model.py",
    "content": "\"\"\"\n Copyright (c) 2022, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\n\nimport logging\nimport os\n\nimport numpy as np\nimport torch\nimport torch.nn as nn\nfrom xraypulse.common.dist_utils import download_cached_file, is_dist_avail_and_initialized\nfrom xraypulse.common.utils import get_abs_path, is_url\nfrom omegaconf import OmegaConf\n\n\nclass BaseModel(nn.Module):\n    \"\"\"Base class for models.\"\"\"\n\n    def __init__(self):\n        super().__init__()\n\n    @property\n    def device(self):\n        return list(self.parameters())[0].device\n\n    def load_checkpoint(self, url_or_filename):\n        \"\"\"\n        Load from a finetuned checkpoint.\n\n        This should expect no mismatch in the model keys and the checkpoint keys.\n        \"\"\"\n\n        if is_url(url_or_filename):\n            cached_file = download_cached_file(\n                url_or_filename, check_hash=False, progress=True\n            )\n            checkpoint = torch.load(cached_file, map_location=\"cpu\")\n        elif os.path.isfile(url_or_filename):\n            checkpoint = torch.load(url_or_filename, map_location=\"cpu\")\n        else:\n            raise RuntimeError(\"checkpoint url or path is invalid\")\n\n        if \"model\" in checkpoint.keys():\n            state_dict = checkpoint[\"model\"]\n        else:\n            state_dict = checkpoint\n\n        msg = self.load_state_dict(state_dict, strict=False)\n\n        logging.info(\"Missing keys {}\".format(msg.missing_keys))\n        logging.info(\"load checkpoint from %s\" % url_or_filename)\n\n        return msg\n\n    @classmethod\n    def from_pretrained(cls, model_type):\n        \"\"\"\n        Build a pretrained model from default configuration file, specified by model_type.\n\n        Args:\n            - model_type (str): model type, specifying architecture and checkpoints.\n\n        Returns:\n            - model (nn.Module): pretrained or finetuned model, depending on the configuration.\n        \"\"\"\n        model_cfg = OmegaConf.load(cls.default_config_path(model_type)).model\n        model = cls.from_config(model_cfg)\n\n        return model\n\n    @classmethod\n    def default_config_path(cls, model_type):\n        assert (\n            model_type in cls.PRETRAINED_MODEL_CONFIG_DICT\n        ), \"Unknown model type {}\".format(model_type)\n        return get_abs_path(cls.PRETRAINED_MODEL_CONFIG_DICT[model_type])\n\n    def load_checkpoint_from_config(self, cfg, **kwargs):\n        \"\"\"\n        Load checkpoint as specified in the config file.\n\n        If load_finetuned is True, load the finetuned model; otherwise, load the pretrained model.\n        When loading the pretrained model, each task-specific architecture may define their\n        own load_from_pretrained() method.\n        \"\"\"\n        load_finetuned = cfg.get(\"load_finetuned\", True)\n        if load_finetuned:\n            finetune_path = cfg.get(\"finetuned\", None)\n            assert (\n                finetune_path is not None\n            ), \"Found load_finetuned is True, but finetune_path is None.\"\n            self.load_checkpoint(url_or_filename=finetune_path)\n        else:\n            # load pre-trained weights\n            pretrain_path = cfg.get(\"pretrained\", None)\n            assert \"Found load_finetuned is False, but pretrain_path is None.\"\n            self.load_from_pretrained(url_or_filename=pretrain_path, **kwargs)\n\n    def before_evaluation(self, **kwargs):\n        pass\n\n    def show_n_params(self, return_str=True):\n        tot = 0\n        for p in self.parameters():\n            w = 1\n            for x in p.shape:\n                w *= x\n            tot += w\n        if return_str:\n            if tot >= 1e6:\n                return \"{:.1f}M\".format(tot / 1e6)\n            else:\n                return \"{:.1f}K\".format(tot / 1e3)\n        else:\n            return tot\n\n\nclass BaseEncoder(nn.Module):\n    \"\"\"\n    Base class for primitive encoders, such as ViT, TimeSformer, etc.\n    \"\"\"\n\n    def __init__(self):\n        super().__init__()\n\n    def forward_features(self, samples, **kwargs):\n        raise NotImplementedError\n\n    @property\n    def device(self):\n        return list(self.parameters())[0].device\n\n\nclass SharedQueueMixin:\n    @torch.no_grad()\n    def _dequeue_and_enqueue(self, image_feat, text_feat, idxs=None):\n        # gather keys before updating queue\n        image_feats = concat_all_gather(image_feat)\n        text_feats = concat_all_gather(text_feat)\n\n        batch_size = image_feats.shape[0]\n\n        ptr = int(self.queue_ptr)\n        assert self.queue_size % batch_size == 0  # for simplicity\n\n        # replace the keys at ptr (dequeue and enqueue)\n        self.image_queue[:, ptr : ptr + batch_size] = image_feats.T\n        self.text_queue[:, ptr : ptr + batch_size] = text_feats.T\n\n        if idxs is not None:\n            idxs = concat_all_gather(idxs)\n            self.idx_queue[:, ptr : ptr + batch_size] = idxs.T\n\n        ptr = (ptr + batch_size) % self.queue_size  # move pointer\n        self.queue_ptr[0] = ptr\n\n\nclass MomentumDistilationMixin:\n    @torch.no_grad()\n    def copy_params(self):\n        for model_pair in self.model_pairs:\n            for param, param_m in zip(\n                model_pair[0].parameters(), model_pair[1].parameters()\n            ):\n                param_m.data.copy_(param.data)  # initialize\n                param_m.requires_grad = False  # not update by gradient\n\n    @torch.no_grad()\n    def _momentum_update(self):\n        for model_pair in self.model_pairs:\n            for param, param_m in zip(\n                model_pair[0].parameters(), model_pair[1].parameters()\n            ):\n                param_m.data = param_m.data * self.momentum + param.data * (\n                    1.0 - self.momentum\n                )\n\n\nclass GatherLayer(torch.autograd.Function):\n    \"\"\"\n    Gather tensors from all workers with support for backward propagation:\n    This implementation does not cut the gradients as torch.distributed.all_gather does.\n    \"\"\"\n\n    @staticmethod\n    def forward(ctx, x):\n        output = [\n            torch.zeros_like(x) for _ in range(torch.distributed.get_world_size())\n        ]\n        torch.distributed.all_gather(output, x)\n        return tuple(output)\n\n    @staticmethod\n    def backward(ctx, *grads):\n        all_gradients = torch.stack(grads)\n        torch.distributed.all_reduce(all_gradients)\n        return all_gradients[torch.distributed.get_rank()]\n\n\ndef all_gather_with_grad(tensors):\n    \"\"\"\n    Performs all_gather operation on the provided tensors.\n    Graph remains connected for backward grad computation.\n    \"\"\"\n    # Queue the gathered tensors\n    world_size = torch.distributed.get_world_size()\n    # There is no need for reduction in the single-proc case\n    if world_size == 1:\n        return tensors\n\n    # tensor_all = GatherLayer.apply(tensors)\n    tensor_all = GatherLayer.apply(tensors)\n\n    return torch.cat(tensor_all, dim=0)\n\n\n@torch.no_grad()\ndef concat_all_gather(tensor):\n    \"\"\"\n    Performs all_gather operation on the provided tensors.\n    *** Warning ***: torch.distributed.all_gather has no gradient.\n    \"\"\"\n    # if use distributed training\n    if not is_dist_avail_and_initialized():\n        return tensor\n\n    tensors_gather = [\n        torch.ones_like(tensor) for _ in range(torch.distributed.get_world_size())\n    ]\n    torch.distributed.all_gather(tensors_gather, tensor, async_op=False)\n\n    output = torch.cat(tensors_gather, dim=0)\n    return output\n\n\ndef tile(x, dim, n_tile):\n    init_dim = x.size(dim)\n    repeat_idx = [1] * x.dim()\n    repeat_idx[dim] = n_tile\n    x = x.repeat(*(repeat_idx))\n    order_index = torch.LongTensor(\n        np.concatenate([init_dim * np.arange(n_tile) + i for i in range(init_dim)])\n    )\n    return torch.index_select(x, dim, order_index.to(x.device))\n"
  },
  {
    "path": "xraypulse/models/blip2.py",
    "content": "\"\"\"\n Copyright (c) 2023, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\nimport contextlib\nimport logging\nimport os\nimport time\nimport datetime\n\nimport torch\nimport torch.nn as nn\nimport torch.distributed as dist\nimport torch.nn.functional as F\n\nimport xraypulse.common.dist_utils as dist_utils\nfrom xraypulse.common.dist_utils import download_cached_file\nfrom xraypulse.common.utils import is_url\nfrom xraypulse.common.logger import MetricLogger\nfrom xraypulse.models.base_model import BaseModel\nfrom xraypulse.models.Qformer import BertConfig, BertLMHeadModel\nfrom xraypulse.models.eva_vit import create_eva_vit_g\nfrom transformers import BertTokenizer\n\n\nclass Blip2Base(BaseModel):\n    @classmethod\n    def init_tokenizer(cls):\n        tokenizer = BertTokenizer.from_pretrained(\"bert-base-uncased\")\n        tokenizer.add_special_tokens({\"bos_token\": \"[DEC]\"})\n        return tokenizer\n\n    def maybe_autocast(self, dtype=torch.float16):\n        # if on cpu, don't use autocast\n        # if on gpu, use autocast with dtype if provided, otherwise use torch.float16\n        enable_autocast = self.device != torch.device(\"cpu\")\n\n        if enable_autocast:\n            return torch.cuda.amp.autocast(dtype=dtype)\n        else:\n            return contextlib.nullcontext()\n\n    @classmethod\n    def init_Qformer(cls, num_query_token, vision_width, cross_attention_freq=2):\n        encoder_config = BertConfig.from_pretrained(\"bert-base-uncased\")\n        encoder_config.encoder_width = vision_width\n        # insert cross-attention layer every other block\n        encoder_config.add_cross_attention = True\n        encoder_config.cross_attention_freq = cross_attention_freq\n        encoder_config.query_length = num_query_token\n        Qformer = BertLMHeadModel(config=encoder_config)\n        query_tokens = nn.Parameter(\n            torch.zeros(1, num_query_token, encoder_config.hidden_size)\n        )\n        query_tokens.data.normal_(mean=0.0, std=encoder_config.initializer_range)\n        return Qformer, query_tokens\n\n    @classmethod\n    def init_vision_encoder(\n        cls, model_name, img_size, drop_path_rate, use_grad_checkpoint, precision\n    ):\n        assert model_name == \"eva_clip_g\", \"vit model must be eva_clip_g for current version of MiniGPT-4\"\n        visual_encoder = create_eva_vit_g(\n            img_size, drop_path_rate, use_grad_checkpoint, precision\n        )\n\n        ln_vision = LayerNorm(visual_encoder.num_features)\n        return visual_encoder, ln_vision\n\n    def load_from_pretrained(self, url_or_filename):\n        if is_url(url_or_filename):\n            cached_file = download_cached_file(\n                url_or_filename, check_hash=False, progress=True\n            )\n            checkpoint = torch.load(cached_file, map_location=\"cpu\")\n        elif os.path.isfile(url_or_filename):\n            checkpoint = torch.load(url_or_filename, map_location=\"cpu\")\n        else:\n            raise RuntimeError(\"checkpoint url or path is invalid\")\n\n        state_dict = checkpoint[\"model\"]\n\n        msg = self.load_state_dict(state_dict, strict=False)\n\n        # logging.info(\"Missing keys {}\".format(msg.missing_keys))\n        logging.info(\"load checkpoint from %s\" % url_or_filename)\n\n        return msg\n\n\ndef disabled_train(self, mode=True):\n    \"\"\"Overwrite model.train with this function to make sure train/eval mode\n    does not change anymore.\"\"\"\n    return self\n\n\nclass LayerNorm(nn.LayerNorm):\n    \"\"\"Subclass torch's LayerNorm to handle fp16.\"\"\"\n\n    def forward(self, x: torch.Tensor):\n        orig_type = x.dtype\n        ret = super().forward(x.type(torch.float32))\n        return ret.type(orig_type)\n\n\ndef compute_sim_matrix(model, data_loader, **kwargs):\n    k_test = kwargs.pop(\"k_test\")\n\n    metric_logger = MetricLogger(delimiter=\"  \")\n    header = \"Evaluation:\"\n\n    logging.info(\"Computing features for evaluation...\")\n    start_time = time.time()\n\n    texts = data_loader.dataset.text\n    num_text = len(texts)\n    text_bs = 256\n    text_ids = []\n    text_embeds = []\n    text_atts = []\n    for i in range(0, num_text, text_bs):\n        text = texts[i : min(num_text, i + text_bs)]\n        text_input = model.tokenizer(\n            text,\n            padding=\"max_length\",\n            truncation=True,\n            max_length=35,\n            return_tensors=\"pt\",\n        ).to(model.device)\n        text_feat = model.forward_text(text_input)\n        text_embed = F.normalize(model.text_proj(text_feat))\n        text_embeds.append(text_embed)\n        text_ids.append(text_input.input_ids)\n        text_atts.append(text_input.attention_mask)\n\n    text_embeds = torch.cat(text_embeds, dim=0)\n    text_ids = torch.cat(text_ids, dim=0)\n    text_atts = torch.cat(text_atts, dim=0)\n\n    vit_feats = []\n    image_embeds = []\n    for samples in data_loader:\n        image = samples[\"image\"]\n\n        image = image.to(model.device)\n        image_feat, vit_feat = model.forward_image(image)\n        image_embed = model.vision_proj(image_feat)\n        image_embed = F.normalize(image_embed, dim=-1)\n\n        vit_feats.append(vit_feat.cpu())\n        image_embeds.append(image_embed)\n\n    vit_feats = torch.cat(vit_feats, dim=0)\n    image_embeds = torch.cat(image_embeds, dim=0)\n\n    sims_matrix = []\n    for image_embed in image_embeds:\n        sim_q2t = image_embed @ text_embeds.t()\n        sim_i2t, _ = sim_q2t.max(0)\n        sims_matrix.append(sim_i2t)\n    sims_matrix = torch.stack(sims_matrix, dim=0)\n\n    score_matrix_i2t = torch.full(\n        (len(data_loader.dataset.image), len(texts)), -100.0\n    ).to(model.device)\n\n    num_tasks = dist_utils.get_world_size()\n    rank = dist_utils.get_rank()\n    step = sims_matrix.size(0) // num_tasks + 1\n    start = rank * step\n    end = min(sims_matrix.size(0), start + step)\n\n    for i, sims in enumerate(\n        metric_logger.log_every(sims_matrix[start:end], 50, header)\n    ):\n        topk_sim, topk_idx = sims.topk(k=k_test, dim=0)\n        image_inputs = vit_feats[start + i].repeat(k_test, 1, 1).to(model.device)\n        score = model.compute_itm(\n            image_inputs=image_inputs,\n            text_ids=text_ids[topk_idx],\n            text_atts=text_atts[topk_idx],\n        ).float()\n        score_matrix_i2t[start + i, topk_idx] = score + topk_sim\n\n    sims_matrix = sims_matrix.t()\n    score_matrix_t2i = torch.full(\n        (len(texts), len(data_loader.dataset.image)), -100.0\n    ).to(model.device)\n\n    step = sims_matrix.size(0) // num_tasks + 1\n    start = rank * step\n    end = min(sims_matrix.size(0), start + step)\n\n    for i, sims in enumerate(\n        metric_logger.log_every(sims_matrix[start:end], 50, header)\n    ):\n        topk_sim, topk_idx = sims.topk(k=k_test, dim=0)\n        image_inputs = vit_feats[topk_idx.cpu()].to(model.device)\n        score = model.compute_itm(\n            image_inputs=image_inputs,\n            text_ids=text_ids[start + i].repeat(k_test, 1),\n            text_atts=text_atts[start + i].repeat(k_test, 1),\n        ).float()\n        score_matrix_t2i[start + i, topk_idx] = score + topk_sim\n\n    if dist_utils.is_dist_avail_and_initialized():\n        dist.barrier()\n        torch.distributed.all_reduce(\n            score_matrix_i2t, op=torch.distributed.ReduceOp.SUM\n        )\n        torch.distributed.all_reduce(\n            score_matrix_t2i, op=torch.distributed.ReduceOp.SUM\n        )\n\n    total_time = time.time() - start_time\n    total_time_str = str(datetime.timedelta(seconds=int(total_time)))\n    logging.info(\"Evaluation time {}\".format(total_time_str))\n\n    return score_matrix_i2t.cpu().numpy(), score_matrix_t2i.cpu().numpy()\n"
  },
  {
    "path": "xraypulse/models/blip2_outputs.py",
    "content": "\"\"\"\n Copyright (c) 2022, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\n\nfrom dataclasses import dataclass\nfrom typing import Optional\n\nimport torch\nfrom transformers.modeling_outputs import (\n    ModelOutput,\n    BaseModelOutputWithPoolingAndCrossAttentions,\n    CausalLMOutputWithCrossAttentions,\n)\n\n\n@dataclass\nclass BlipSimilarity(ModelOutput):\n    sim_i2t: torch.FloatTensor = None\n    sim_t2i: torch.FloatTensor = None\n\n    sim_i2t_m: Optional[torch.FloatTensor] = None\n    sim_t2i_m: Optional[torch.FloatTensor] = None\n\n    sim_i2t_targets: Optional[torch.FloatTensor] = None\n    sim_t2i_targets: Optional[torch.FloatTensor] = None\n\n\n@dataclass\nclass BlipIntermediateOutput(ModelOutput):\n    \"\"\"\n    Data class for intermediate outputs of BLIP models.\n\n    image_embeds (torch.FloatTensor): Image embeddings, shape (batch_size, num_patches, embed_dim).\n    text_embeds (torch.FloatTensor): Text embeddings, shape (batch_size, seq_len, embed_dim).\n\n    image_embeds_m (torch.FloatTensor): Image embeddings from momentum visual encoder, shape (batch_size, num_patches, embed_dim).\n    text_embeds_m (torch.FloatTensor): Text embeddings from momentum text encoder, shape (batch_size, seq_len, embed_dim).\n\n    encoder_output (BaseModelOutputWithPoolingAndCrossAttentions): output from the image-grounded text encoder.\n    encoder_output_neg (BaseModelOutputWithPoolingAndCrossAttentions): output from the image-grounded text encoder for negative pairs.\n\n    decoder_output (CausalLMOutputWithCrossAttentions): output from the image-grounded text decoder.\n    decoder_labels (torch.LongTensor): labels for the captioning loss.\n\n    itm_logits (torch.FloatTensor): logits for the image-text matching loss, shape (batch_size * 3, 2).\n    itm_labels (torch.LongTensor): labels for the image-text matching loss, shape (batch_size * 3,)\n\n    \"\"\"\n\n    # uni-modal features\n    image_embeds: torch.FloatTensor = None\n    text_embeds: Optional[torch.FloatTensor] = None\n\n    image_embeds_m: Optional[torch.FloatTensor] = None\n    text_embeds_m: Optional[torch.FloatTensor] = None\n\n    # intermediate outputs of multimodal encoder\n    encoder_output: Optional[BaseModelOutputWithPoolingAndCrossAttentions] = None\n    encoder_output_neg: Optional[BaseModelOutputWithPoolingAndCrossAttentions] = None\n\n    itm_logits: Optional[torch.FloatTensor] = None\n    itm_labels: Optional[torch.LongTensor] = None\n\n    # intermediate outputs of multimodal decoder\n    decoder_output: Optional[CausalLMOutputWithCrossAttentions] = None\n    decoder_labels: Optional[torch.LongTensor] = None\n\n\n@dataclass\nclass BlipOutput(ModelOutput):\n    # some finetuned models (e.g. BlipVQA) do not compute similarity, thus optional.\n    sims: Optional[BlipSimilarity] = None\n\n    intermediate_output: BlipIntermediateOutput = None\n\n    loss: Optional[torch.FloatTensor] = None\n\n    loss_itc: Optional[torch.FloatTensor] = None\n\n    loss_itm: Optional[torch.FloatTensor] = None\n\n    loss_lm: Optional[torch.FloatTensor] = None\n\n\n@dataclass\nclass BlipOutputFeatures(ModelOutput):\n    \"\"\"\n    Data class of features from BlipFeatureExtractor.\n\n    Args:\n        image_embeds: (torch.FloatTensor) of shape (batch_size, num_patches+1, embed_dim), optional\n        image_features: (torch.FloatTensor) of shape (batch_size, num_patches+1, feature_dim), optional\n        text_embeds: (torch.FloatTensor) of shape (batch_size, sequence_length+1, embed_dim), optional\n        text_features: (torch.FloatTensor) of shape (batch_size, sequence_length+1, feature_dim), optional\n\n        The first embedding or feature is for the [CLS] token.\n\n        Features are obtained by projecting the corresponding embedding into a normalized low-dimensional space.\n    \"\"\"\n\n    image_embeds: Optional[torch.FloatTensor] = None\n    image_embeds_proj: Optional[torch.FloatTensor] = None\n\n    text_embeds: Optional[torch.FloatTensor] = None\n    text_embeds_proj: Optional[torch.FloatTensor] = None\n\n    multimodal_embeds: Optional[torch.FloatTensor] = None\n"
  },
  {
    "path": "xraypulse/models/eva_vit.py",
    "content": "# Based on EVA, BEIT, timm and DeiT code bases\n# https://github.com/baaivision/EVA\n# https://github.com/rwightman/pytorch-image-models/tree/master/timm\n# https://github.com/microsoft/unilm/tree/master/beit\n# https://github.com/facebookresearch/deit/\n# https://github.com/facebookresearch/dino\n# --------------------------------------------------------'\nimport math\nfrom functools import partial\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.utils.checkpoint as checkpoint\nfrom timm.models.layers import drop_path, to_2tuple, trunc_normal_\nfrom timm.models.registry import register_model\n\nfrom xraypulse.common.dist_utils import download_cached_file\n\ndef _cfg(url='', **kwargs):\n    return {\n        'url': url,\n        'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': None,\n        'crop_pct': .9, 'interpolation': 'bicubic',\n        'mean': (0.5, 0.5, 0.5), 'std': (0.5, 0.5, 0.5),\n        **kwargs\n    }\n\n\nclass DropPath(nn.Module):\n    \"\"\"Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).\n    \"\"\"\n    def __init__(self, drop_prob=None):\n        super(DropPath, self).__init__()\n        self.drop_prob = drop_prob\n\n    def forward(self, x):\n        return drop_path(x, self.drop_prob, self.training)\n    \n    def extra_repr(self) -> str:\n        return 'p={}'.format(self.drop_prob)\n\n\nclass Mlp(nn.Module):\n    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):\n        super().__init__()\n        out_features = out_features or in_features\n        hidden_features = hidden_features or in_features\n        self.fc1 = nn.Linear(in_features, hidden_features)\n        self.act = act_layer()\n        self.fc2 = nn.Linear(hidden_features, out_features)\n        self.drop = nn.Dropout(drop)\n\n    def forward(self, x):\n        x = self.fc1(x)\n        x = self.act(x)\n        # x = self.drop(x)\n        # commit this for the orignal BERT implement \n        x = self.fc2(x)\n        x = self.drop(x)\n        return x\n\n\nclass Attention(nn.Module):\n    def __init__(\n            self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.,\n            proj_drop=0., window_size=None, attn_head_dim=None):\n        super().__init__()\n        self.num_heads = num_heads\n        head_dim = dim // num_heads\n        if attn_head_dim is not None:\n            head_dim = attn_head_dim\n        all_head_dim = head_dim * self.num_heads\n        self.scale = qk_scale or head_dim ** -0.5\n\n        self.qkv = nn.Linear(dim, all_head_dim * 3, bias=False)\n        if qkv_bias:\n            self.q_bias = nn.Parameter(torch.zeros(all_head_dim))\n            self.v_bias = nn.Parameter(torch.zeros(all_head_dim))\n        else:\n            self.q_bias = None\n            self.v_bias = None\n\n        if window_size:\n            self.window_size = window_size\n            self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3\n            self.relative_position_bias_table = nn.Parameter(\n                torch.zeros(self.num_relative_distance, num_heads))  # 2*Wh-1 * 2*Ww-1, nH\n            # cls to token & token 2 cls & cls to cls\n\n            # get pair-wise relative position index for each token inside the window\n            coords_h = torch.arange(window_size[0])\n            coords_w = torch.arange(window_size[1])\n            coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww\n            coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww\n            relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww\n            relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2\n            relative_coords[:, :, 0] += window_size[0] - 1  # shift to start from 0\n            relative_coords[:, :, 1] += window_size[1] - 1\n            relative_coords[:, :, 0] *= 2 * window_size[1] - 1\n            relative_position_index = \\\n                torch.zeros(size=(window_size[0] * window_size[1] + 1, ) * 2, dtype=relative_coords.dtype)\n            relative_position_index[1:, 1:] = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww\n            relative_position_index[0, 0:] = self.num_relative_distance - 3\n            relative_position_index[0:, 0] = self.num_relative_distance - 2\n            relative_position_index[0, 0] = self.num_relative_distance - 1\n\n            self.register_buffer(\"relative_position_index\", relative_position_index)\n        else:\n            self.window_size = None\n            self.relative_position_bias_table = None\n            self.relative_position_index = None\n\n        self.attn_drop = nn.Dropout(attn_drop)\n        self.proj = nn.Linear(all_head_dim, dim)\n        self.proj_drop = nn.Dropout(proj_drop)\n\n    def forward(self, x, rel_pos_bias=None):\n        B, N, C = x.shape\n        qkv_bias = None\n        if self.q_bias is not None:\n            qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))\n        # qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)\n        qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)\n        qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)\n        q, k, v = qkv[0], qkv[1], qkv[2]   # make torchscript happy (cannot use tensor as tuple)\n\n        q = q * self.scale\n        attn = (q @ k.transpose(-2, -1))\n\n        if self.relative_position_bias_table is not None:\n            relative_position_bias = \\\n                self.relative_position_bias_table[self.relative_position_index.view(-1)].view(\n                    self.window_size[0] * self.window_size[1] + 1,\n                    self.window_size[0] * self.window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH\n            relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww\n            attn = attn + relative_position_bias.unsqueeze(0)\n\n        if rel_pos_bias is not None:\n            attn = attn + rel_pos_bias\n        \n        attn = attn.softmax(dim=-1)\n        attn = self.attn_drop(attn)\n\n        x = (attn @ v).transpose(1, 2).reshape(B, N, -1)\n        x = self.proj(x)\n        x = self.proj_drop(x)\n        return x\n\n\nclass Block(nn.Module):\n\n    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,\n                 drop_path=0., init_values=None, act_layer=nn.GELU, norm_layer=nn.LayerNorm,\n                 window_size=None, attn_head_dim=None):\n        super().__init__()\n        self.norm1 = norm_layer(dim)\n        self.attn = Attention(\n            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,\n            attn_drop=attn_drop, proj_drop=drop, window_size=window_size, attn_head_dim=attn_head_dim)\n        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here\n        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()\n        self.norm2 = norm_layer(dim)\n        mlp_hidden_dim = int(dim * mlp_ratio)\n        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)\n\n        if init_values is not None and init_values > 0:\n            self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)\n            self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)\n        else:\n            self.gamma_1, self.gamma_2 = None, None\n\n    def forward(self, x, rel_pos_bias=None):\n        if self.gamma_1 is None:\n            x = x + self.drop_path(self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias))\n            x = x + self.drop_path(self.mlp(self.norm2(x)))\n        else:\n            x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias))\n            x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))\n        return x\n\n\nclass PatchEmbed(nn.Module):\n    \"\"\" Image to Patch Embedding\n    \"\"\"\n    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):\n        super().__init__()\n        img_size = to_2tuple(img_size)\n        patch_size = to_2tuple(patch_size)\n        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])\n        self.patch_shape = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])\n        self.img_size = img_size\n        self.patch_size = patch_size\n        self.num_patches = num_patches\n\n        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)\n\n    def forward(self, x, **kwargs):\n        B, C, H, W = x.shape\n        # FIXME look at relaxing size constraints\n        assert H == self.img_size[0] and W == self.img_size[1], \\\n            f\"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]}).\"\n        x = self.proj(x).flatten(2).transpose(1, 2)\n        return x\n\n\nclass RelativePositionBias(nn.Module):\n\n    def __init__(self, window_size, num_heads):\n        super().__init__()\n        self.window_size = window_size\n        self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3\n        self.relative_position_bias_table = nn.Parameter(\n            torch.zeros(self.num_relative_distance, num_heads))  # 2*Wh-1 * 2*Ww-1, nH\n        # cls to token & token 2 cls & cls to cls\n\n        # get pair-wise relative position index for each token inside the window\n        coords_h = torch.arange(window_size[0])\n        coords_w = torch.arange(window_size[1])\n        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww\n        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww\n        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww\n        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2\n        relative_coords[:, :, 0] += window_size[0] - 1  # shift to start from 0\n        relative_coords[:, :, 1] += window_size[1] - 1\n        relative_coords[:, :, 0] *= 2 * window_size[1] - 1\n        relative_position_index = \\\n            torch.zeros(size=(window_size[0] * window_size[1] + 1,) * 2, dtype=relative_coords.dtype)\n        relative_position_index[1:, 1:] = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww\n        relative_position_index[0, 0:] = self.num_relative_distance - 3\n        relative_position_index[0:, 0] = self.num_relative_distance - 2\n        relative_position_index[0, 0] = self.num_relative_distance - 1\n\n        self.register_buffer(\"relative_position_index\", relative_position_index)\n\n        # trunc_normal_(self.relative_position_bias_table, std=.02)\n\n    def forward(self):\n        relative_position_bias = \\\n            self.relative_position_bias_table[self.relative_position_index.view(-1)].view(\n                self.window_size[0] * self.window_size[1] + 1,\n                self.window_size[0] * self.window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH\n        return relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww\n\n\nclass VisionTransformer(nn.Module):\n    \"\"\" Vision Transformer with support for patch or hybrid CNN input stage\n    \"\"\"\n    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,\n                 num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,\n                 drop_path_rate=0., norm_layer=nn.LayerNorm, init_values=None,\n                 use_abs_pos_emb=True, use_rel_pos_bias=False, use_shared_rel_pos_bias=False,\n                 use_mean_pooling=True, init_scale=0.001, use_checkpoint=False):\n        super().__init__()\n        self.image_size = img_size\n        self.num_classes = num_classes\n        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models\n\n        self.patch_embed = PatchEmbed(\n            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)\n        num_patches = self.patch_embed.num_patches\n\n        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))\n        if use_abs_pos_emb:\n            self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))\n        else:\n            self.pos_embed = None\n        self.pos_drop = nn.Dropout(p=drop_rate)\n\n        if use_shared_rel_pos_bias:\n            self.rel_pos_bias = RelativePositionBias(window_size=self.patch_embed.patch_shape, num_heads=num_heads)\n        else:\n            self.rel_pos_bias = None\n        self.use_checkpoint = use_checkpoint\n        \n        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule\n        self.use_rel_pos_bias = use_rel_pos_bias\n        self.blocks = nn.ModuleList([\n            Block(\n                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,\n                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer,\n                init_values=init_values, window_size=self.patch_embed.patch_shape if use_rel_pos_bias else None)\n            for i in range(depth)])\n#         self.norm = nn.Identity() if use_mean_pooling else norm_layer(embed_dim)\n#         self.fc_norm = norm_layer(embed_dim) if use_mean_pooling else None\n#         self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()\n\n        if self.pos_embed is not None:\n            trunc_normal_(self.pos_embed, std=.02)\n        trunc_normal_(self.cls_token, std=.02)\n        # trunc_normal_(self.mask_token, std=.02)\n#         if isinstance(self.head, nn.Linear):\n#             trunc_normal_(self.head.weight, std=.02)\n        self.apply(self._init_weights)\n        self.fix_init_weight()\n#         if isinstance(self.head, nn.Linear):\n#             self.head.weight.data.mul_(init_scale)\n#             self.head.bias.data.mul_(init_scale)\n\n    def fix_init_weight(self):\n        def rescale(param, layer_id):\n            param.div_(math.sqrt(2.0 * layer_id))\n\n        for layer_id, layer in enumerate(self.blocks):\n            rescale(layer.attn.proj.weight.data, layer_id + 1)\n            rescale(layer.mlp.fc2.weight.data, layer_id + 1)\n\n    def _init_weights(self, m):\n        if isinstance(m, nn.Linear):\n            trunc_normal_(m.weight, std=.02)\n            if isinstance(m, nn.Linear) and m.bias is not None:\n                nn.init.constant_(m.bias, 0)\n        elif isinstance(m, nn.LayerNorm):\n            nn.init.constant_(m.bias, 0)\n            nn.init.constant_(m.weight, 1.0)\n\n    def get_classifier(self):\n        return self.head\n\n    def reset_classifier(self, num_classes, global_pool=''):\n        self.num_classes = num_classes\n        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()\n\n    def forward_features(self, x):\n        x = self.patch_embed(x)\n        batch_size, seq_len, _ = x.size()\n\n        cls_tokens = self.cls_token.expand(batch_size, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks\n        x = torch.cat((cls_tokens, x), dim=1)\n        if self.pos_embed is not None:\n            x = x + self.pos_embed\n        x = self.pos_drop(x)\n\n        rel_pos_bias = self.rel_pos_bias() if self.rel_pos_bias is not None else None\n        for blk in self.blocks:\n            if self.use_checkpoint:\n                x = checkpoint.checkpoint(blk, x, rel_pos_bias)\n            else:\n                x = blk(x, rel_pos_bias)\n        return x\n#         x = self.norm(x)\n\n#         if self.fc_norm is not None:\n#             t = x[:, 1:, :]\n#             return self.fc_norm(t.mean(1))\n#         else:\n#             return x[:, 0]\n\n    def forward(self, x):\n        x = self.forward_features(x)\n#         x = self.head(x)\n        return x\n\n    def get_intermediate_layers(self, x):\n        x = self.patch_embed(x)\n        batch_size, seq_len, _ = x.size()\n\n        cls_tokens = self.cls_token.expand(batch_size, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks\n        x = torch.cat((cls_tokens, x), dim=1)\n        if self.pos_embed is not None:\n            x = x + self.pos_embed\n        x = self.pos_drop(x)\n\n        features = []\n        rel_pos_bias = self.rel_pos_bias() if self.rel_pos_bias is not None else None\n        for blk in self.blocks:\n            x = blk(x, rel_pos_bias)\n            features.append(x)\n\n        return features\n    \n    \ndef interpolate_pos_embed(model, checkpoint_model):\n    if 'pos_embed' in checkpoint_model:\n        pos_embed_checkpoint = checkpoint_model['pos_embed'].float()\n        embedding_size = pos_embed_checkpoint.shape[-1]\n        num_patches = model.patch_embed.num_patches\n        num_extra_tokens = model.pos_embed.shape[-2] - num_patches\n        # height (== width) for the checkpoint position embedding\n        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)\n        # height (== width) for the new position embedding\n        new_size = int(num_patches ** 0.5)\n        # class_token and dist_token are kept unchanged\n        if orig_size != new_size:\n            print(\"Position interpolate from %dx%d to %dx%d\" % (orig_size, orig_size, new_size, new_size))\n            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]\n            # only the position tokens are interpolated\n            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]\n            pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)\n            pos_tokens = torch.nn.functional.interpolate(\n                pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)\n            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)\n            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)\n            checkpoint_model['pos_embed'] = new_pos_embed\n            \n            \ndef convert_weights_to_fp16(model: nn.Module):\n    \"\"\"Convert applicable model parameters to fp16\"\"\"\n\n    def _convert_weights_to_fp16(l):\n        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):\n            l.weight.data = l.weight.data.half()\n            if l.bias is not None:\n                l.bias.data = l.bias.data.half()\n\n#         if isinstance(l, (nn.MultiheadAttention, Attention)):\n#             for attr in [*[f\"{s}_proj_weight\" for s in [\"in\", \"q\", \"k\", \"v\"]], \"in_proj_bias\", \"bias_k\", \"bias_v\"]:\n#                 tensor = getattr(l, attr)\n#                 if tensor is not None:\n#                     tensor.data = tensor.data.half()\n\n    model.apply(_convert_weights_to_fp16)\n    \n    \ndef create_eva_vit_g(img_size=224,drop_path_rate=0.4,use_checkpoint=False,precision=\"fp16\"):\n    model = VisionTransformer(\n        img_size=img_size,\n        patch_size=14,\n        use_mean_pooling=False,\n        embed_dim=1408,\n        depth=39,\n        num_heads=1408//88,\n        mlp_ratio=4.3637,\n        qkv_bias=True,\n        drop_path_rate=drop_path_rate,\n        norm_layer=partial(nn.LayerNorm, eps=1e-6),\n        use_checkpoint=use_checkpoint,\n    )  \n    url = \"https://storage.googleapis.com/sfr-vision-language-research/LAVIS/models/BLIP2/eva_vit_g.pth\"\n    cached_file = download_cached_file(\n        url, check_hash=False, progress=True\n    )\n    state_dict = torch.load(cached_file, map_location=\"cpu\")    \n    interpolate_pos_embed(model,state_dict)\n    \n    incompatible_keys = model.load_state_dict(state_dict, strict=False)\n#     print(incompatible_keys)\n    \n    if precision == \"fp16\":\n#         model.to(\"cuda\") \n        convert_weights_to_fp16(model)\n    return model"
  },
  {
    "path": "xraypulse/models/pos_embed.py",
    "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n# --------------------------------------------------------\n# Position embedding utils\n# --------------------------------------------------------\n\nimport numpy as np\n\nimport torch\n\n# --------------------------------------------------------\n# 2D sine-cosine position embedding\n# References:\n# Transformer: https://github.com/tensorflow/models/blob/master/official/nlp/transformer/model_utils.py\n# MoCo v3: https://github.com/facebookresearch/moco-v3\n# --------------------------------------------------------\ndef get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):\n    \"\"\"\n    grid_size: int of the grid height and width\n    return:\n    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)\n    \"\"\"\n    grid_h = np.arange(grid_size, dtype=np.float32)\n    grid_w = np.arange(grid_size, dtype=np.float32)\n    grid = np.meshgrid(grid_w, grid_h)  # here w goes first\n    grid = np.stack(grid, axis=0)\n\n    grid = grid.reshape([2, 1, grid_size, grid_size])\n    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)\n    if cls_token:\n        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)\n    return pos_embed\n\n\ndef get_2d_sincos_pos_embed_from_grid(embed_dim, grid):\n    assert embed_dim % 2 == 0\n\n    # use half of dimensions to encode grid_h\n    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)\n    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)\n\n    emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)\n    return emb\n\n\ndef get_1d_sincos_pos_embed_from_grid(embed_dim, pos):\n    \"\"\"\n    embed_dim: output dimension for each position\n    pos: a list of positions to be encoded: size (M,)\n    out: (M, D)\n    \"\"\"\n    assert embed_dim % 2 == 0\n    omega = np.arange(embed_dim // 2, dtype=np.float)\n    omega = omega / embed_dim / 2.\n    omega = 1. / 10000**omega  # (D/2,)\n\n    pos = pos.reshape(-1)  # (M,)\n    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product\n\n    emb_sin = np.sin(out) # (M, D/2)\n    emb_cos = np.cos(out) # (M, D/2)\n\n    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)\n    return emb\n\n\n# --------------------------------------------------------\n# Interpolate position embeddings for high-resolution\n# References:\n# DeiT: https://github.com/facebookresearch/deit\n# --------------------------------------------------------\ndef interpolate_pos_embed(model, checkpoint_model):\n    if 'pos_embed' in checkpoint_model:\n        pos_embed_checkpoint = checkpoint_model['pos_embed']\n        embedding_size = pos_embed_checkpoint.shape[-1]\n        num_patches = model.patch_embed.num_patches\n        num_extra_tokens = model.pos_embed.shape[-2] - num_patches\n        # height (== width) for the checkpoint position embedding\n        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)\n        # height (== width) for the new position embedding\n        new_size = int(num_patches ** 0.5)\n        # class_token and dist_token are kept unchanged\n        if orig_size != new_size:\n            print(\"Position interpolate from %dx%d to %dx%d\" % (orig_size, orig_size, new_size, new_size))\n            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]\n            # only the position tokens are interpolated\n            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]\n            pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)\n            pos_tokens = torch.nn.functional.interpolate(\n                pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)\n            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)\n            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)\n            checkpoint_model['pos_embed'] = new_pos_embed\n"
  },
  {
    "path": "xraypulse/models/xray_pulse.py",
    "content": "import logging\nimport random\n\nimport torch\nfrom torch.cuda.amp import autocast as autocast\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom xraypulse.common.registry import registry\nfrom xraypulse.models.blip2 import Blip2Base, disabled_train\nfrom transformers import BloomForCausalLM\nfrom transformers import BloomTokenizerFast\nfrom transformers import StoppingCriteria, StoppingCriteriaList #for imprgpt eval\nimport csv #for imprgpt rogue eval\nfrom xraypulse.conversation.conversation import Conversation\nfrom enum import auto, Enum\nfrom typing import List, Tuple, Any\n\nclass StoppingCriteriaSub(StoppingCriteria):\n\n    def __init__(self, stops=[], encounters=1):\n        super().__init__()\n        self.stops = stops\n\n    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor):\n        for stop in self.stops:\n            if torch.all((stop == input_ids[0][-len(stop):])).item():\n                return True\n\n        return False\n    \nclass SeparatorStyle(Enum):\n    \"\"\"Different separator style.\"\"\"\n    SINGLE = auto()\n    TWO = auto()\n    \n\n@registry.register_model(\"xray_pulse\")\nclass XrayPulse(Blip2Base):\n    \"\"\"\n    BLIP2 GPT-Bloom model.\n    \"\"\"\n\n    PRETRAINED_MODEL_CONFIG_DICT = {\n        \"pulse\": \"configs/models/xraypulse.yaml\",\n    }\n\n    def __init__(\n        self,\n        vit_model=\"eva_clip_g\",\n        q_former_model=\"https://storage.googleapis.com/sfr-vision-language-research/LAVIS/models/BLIP2/blip2_pretrained_flant5xxl.pth\",\n        img_size=224,\n        drop_path_rate=0,\n        use_grad_checkpoint=False,\n        vit_precision=\"fp16\",\n        freeze_vit=True,\n        freeze_qformer=True,\n        num_query_token=32,\n        bloom_model=\"\",\n        prompt_path=\"\",\n        prompt_template=\"\",\n        max_txt_len=32,\n        low_resource=False,  # use 8 bit and put vit in cpu\n        end_sym='\\n',\n    ):\n        super().__init__()\n\n        self.tokenizer = self.init_tokenizer()\n        self.low_resource = low_resource\n\n        print('Loading VIT')\n        self.visual_encoder, self.ln_vision = self.init_vision_encoder(\n            vit_model, img_size, drop_path_rate, use_grad_checkpoint, vit_precision\n        )\n        if freeze_vit:\n            for name, param in self.visual_encoder.named_parameters():\n                param.requires_grad = False\n            self.visual_encoder = self.visual_encoder.eval()\n            self.visual_encoder.train = disabled_train\n            for name, param in self.ln_vision.named_parameters():\n                param.requires_grad = False\n            self.ln_vision = self.ln_vision.eval()\n            self.ln_vision.train = disabled_train\n            logging.info(\"freeze vision encoder\")\n        print('Loading VIT Done')\n        \n        print('Loading Q-Former')\n        self.Qformer, self.query_tokens = self.init_Qformer(\n            num_query_token, self.visual_encoder.num_features\n        )\n        self.Qformer.cls = None\n        self.Qformer.bert.embeddings.word_embeddings = None\n        self.Qformer.bert.embeddings.position_embeddings = None\n        for layer in self.Qformer.bert.encoder.layer:\n            layer.output = None\n            layer.intermediate = None\n        self.load_from_pretrained(url_or_filename=q_former_model)\n\n        if freeze_qformer:\n            for name, param in self.Qformer.named_parameters():\n                param.requires_grad = False\n            self.Qformer = self.Qformer.eval()\n            self.Qformer.train = disabled_train\n            self.query_tokens.requires_grad = False\n            logging.info(\"freeze Qformer\")\n        print('Loading Q-Former Done')\n        print('Loading Bloom')\n        self.bloom_tokenizer = BloomTokenizerFast.from_pretrained(bloom_model, use_fast=False)\n        self.bloom_tokenizer.pad_token = self.bloom_tokenizer.eos_token\n\n        if self.low_resource:\n            self.bloom_model = BloomForCausalLM.from_pretrained(\n                bloom_model,\n                torch_dtype=torch.float16,\n                device_map=\"auto\"\n            )\n        else:\n            self.bloom_model = BloomForCausalLM.from_pretrained(\n                bloom_model,\n                torch_dtype=torch.float16,\n            )\n        for name, param in self.bloom_model.named_parameters():\n            param.requires_grad = False\n        print('Loading Bloom Done')\n\n        self.bloom_proj = nn.Linear(\n            self.Qformer.config.hidden_size, self.bloom_model.config.hidden_size\n        )\n        \n        self.max_txt_len = max_txt_len\n        self.end_sym = end_sym\n\n        if prompt_path:\n            with open(prompt_path, 'r') as f:\n                raw_prompts = f.read().splitlines()\n            filted_prompts = [raw_prompt for raw_prompt in raw_prompts if \"<图片>\" in raw_prompt]\n            self.prompt_list = [prompt_template.format(p) for p in filted_prompts]\n            print('Load {} training prompts'.format(len(self.prompt_list)))\n            print('Prompt Example \\n{}'.format(random.choice(self.prompt_list)))\n        else:\n            self.prompt_list = []\n\n        print('#'*100)\n\n    def vit_to_cpu(self):\n        self.ln_vision.to(\"cpu\")\n        self.ln_vision.float()\n        self.visual_encoder.to(\"cpu\")\n        self.visual_encoder.float()\n\n    \n    def encode_img(self, image):\n        device = image.device\n        \n        if self.low_resource:\n            self.vit_to_cpu()\n            image = image.to(\"cpu\")\n\n        with self.maybe_autocast():\n            image_embeds = self.ln_vision(self.visual_encoder(image)).to(device)\n            image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(device)\n\n            query_tokens = self.query_tokens.expand(image_embeds.shape[0], -1, -1)\n            query_output = self.Qformer.bert(\n                query_embeds=query_tokens,\n                encoder_hidden_states=image_embeds,\n                encoder_attention_mask=image_atts,\n                return_dict=True,\n            )\n\n            inputs_bloom = self.bloom_proj(query_output.last_hidden_state)\n            atts_bloom = torch.ones(inputs_bloom.size()[:-1], dtype=torch.long).to(image.device)\n        return inputs_bloom, atts_bloom\n\n    def prompt_wrap(self, img_embeds, atts_img, prompt):\n        if prompt:\n            batch_size = img_embeds.shape[0]\n            p_before, p_after = prompt.split('<图片>')\n            p_before_tokens = self.bloom_tokenizer(\n                p_before, return_tensors=\"pt\", add_special_tokens=False).to(img_embeds.device)\n            p_after_tokens = self.bloom_tokenizer(\n                p_after, return_tensors=\"pt\", add_special_tokens=False).to(img_embeds.device)\n            p_before_embeds = self.bloom_model.transformer.word_embeddings_layernorm(self.bloom_model.transformer.word_embeddings(p_before_tokens.input_ids)).expand(batch_size, -1, -1)\n            p_after_embeds = self.bloom_model.transformer.word_embeddings_layernorm(self.bloom_model.transformer.word_embeddings(p_after_tokens.input_ids)).expand(batch_size, -1, -1)\n            wrapped_img_embeds = torch.cat([p_before_embeds, img_embeds, p_after_embeds], dim=1)\n            wrapped_atts_img = atts_img[:, :1].expand(-1, wrapped_img_embeds.shape[1])\n            return wrapped_img_embeds, wrapped_atts_img\n        else:\n            return img_embeds, atts_img\n\n    def forward(self, samples):\n        image = samples[\"image\"]\n        img_embeds, atts_img = self.encode_img(image)\n        if hasattr(samples, 'question_split'):  # VQA dataset\n            print('VQA Batch')\n            vqa_prompt = 'Instructions: You are PULSE, a large language model trained by SHAIlab. Answer as concisely as possible.\\nKnowledge cutoff: 2021-09-01\\nCurrent date: 2022-02-01</s> User: <Img><图片></Img> </s> Helper: '\n            img_embeds, atts_img = self.prompt_wrap(img_embeds, atts_img, vqa_prompt)\n        elif self.prompt_list:\n            prompt = random.choice(self.prompt_list)\n\n            img_embeds, atts_img = self.prompt_wrap(img_embeds, atts_img, prompt)\n\n        self.bloom_tokenizer.padding_side = \"right\"\n\n        text = [t + self.end_sym for t in samples[\"caption\"]]\n\n        to_regress_tokens = self.bloom_tokenizer(\n            text,\n            return_tensors=\"pt\",\n            padding=\"longest\",\n            truncation=True,\n            max_length=self.max_txt_len,\n            add_special_tokens=False\n        ).to(image.device)\n\n        targets = to_regress_tokens.input_ids.masked_fill(\n            to_regress_tokens.input_ids == self.bloom_tokenizer.pad_token_id, -100\n        )\n\n        empty_targets = (\n            torch.ones([atts_img.shape[0], atts_img.shape[1]+1],\n                       dtype=torch.long).to(image.device).fill_(-100)  # plus one for bos\n        )\n        targets = torch.cat([empty_targets, targets], dim=1)\n\n        batch_size = img_embeds.shape[0]\n        bos = torch.ones([batch_size, 1],\n                         dtype=to_regress_tokens.input_ids.dtype,\n                         device=to_regress_tokens.input_ids.device) * self.bloom_tokenizer.bos_token_id\n        bos_embeds = self.bloom_model.transformer.word_embeddings_layernorm(self.bloom_model.transformer.word_embeddings(bos))\n        atts_bos = atts_img[:, :1]\n\n        to_regress_embeds = self.bloom_model.transformer.word_embeddings_layernorm(self.bloom_model.transformer.word_embeddings(to_regress_tokens.input_ids))\n        inputs_embeds = torch.cat([bos_embeds, img_embeds, to_regress_embeds], dim=1)\n        attention_mask = torch.cat([atts_bos, atts_img, to_regress_tokens.attention_mask], dim=1)\n\n        with self.maybe_autocast():\n            outputs = self.bloom_model(\n                inputs_embeds=inputs_embeds,\n                attention_mask=attention_mask,\n                return_dict=True,\n                labels=targets,\n            )\n        loss = outputs.loss\n        return {\"loss\": loss}\n    \n    def test(self, samples, max_new_tokens=300, num_beams=1, min_length=1, top_p=0.9,\n               repetition_penalty=1.0, length_penalty=1, temperature=1.0, max_length=2000):\n        \n        CONV_ZH = Conversation(\n            system=\"Instructions: You are PULSE, a large language model trained by SHAIlab. Answer as concisely as possible.\\nKnowledge cutoff: 2021-09-01\\nCurrent date: 2022-02-01</s> User: {} </s> Helper: \",\n                # \"Please answer the medical questions based on the patient's description. Give the following medical scan: <Img>图片</Img>.\"\n                # \"You will be able to see the medical scan once I provide it to you. Please answer the patients questions.\",\n            roles=(\"User\", \"Helper\"),\n            messages=[],\n            offset=0,\n            sep_style=SeparatorStyle.SINGLE,\n            sep=\"</s>\",\n            sep2=\"###\",  \n        )\n\n        stop_words_ids = [torch.tensor([835]).to(samples[\"image\"].device),\n                          torch.tensor([2277, 29937]).to(samples[\"image\"].device)]\n        stopping_criteria = StoppingCriteriaList([StoppingCriteriaSub(stops=stop_words_ids)])\n\n        conv.append_message(conv.roles[1], None)\n\n        image = samples[\"image\"]\n        img_embeds, atts_img = self.encode_img(image)\n        conv.append_message(conv.roles[0], \"<Img><图片></Img>\")\n\n\n        embs = self.get_context_emb(conv, img_embeds)\n\n        current_max_len = embs.shape[1] + max_new_tokens\n        if current_max_len - max_length > 0:\n            print('Warning: The number of tokens in current conversation exceeds the max length. '\n                  'The model will not see the contexts outside the range.')\n        begin_idx = max(0, current_max_len - max_length)\n\n        embs = embs[:, begin_idx:]\n\n        outputs = self.bloom_model.generate(\n            inputs_embeds=embs,\n            max_new_tokens=max_new_tokens,\n            stopping_criteria = stopping_criteria,\n            num_beams=num_beams,\n            do_sample=True,\n            min_length=min_length,\n            top_p=top_p,\n            repetition_penalty=repetition_penalty,\n            length_penalty=length_penalty,\n            temperature=temperature,\n        )\n        output_token = outputs[0]\n        if output_token[0] == 0:  # the model might output a unknow token <unk> at the beginning. remove it\n            output_token = output_token[1:]\n        if output_token[0] == 1:  # some users find that there is a start token <s> at the beginning. remove it\n            output_token = output_token[1:]\n        output_text = self.bloom_tokenizer.decode(output_token, add_special_tokens=False)\n        output_text = output_text.split('</s>')[0]\n        conv.messages[-1][1] = output_text\n        return output_text, output_token.cpu().numpy()\n    \n\n    def get_context_emb(self, conv, img):\n        prompt = random.choice(self.prompt_list)\n        text = prompt.split(\"<Img><图片></Img>\")[-1]\n\n        if len(conv.messages) > 0 and conv.messages[-1][0] == conv.roles[0] \\\n                and conv.messages[-1][1][-6:] == '</Img>':  # last message is image.\n            conv.messages[-1][1] = ' '.join([conv.messages[-1][1], text])\n        else:\n            conv.append_message(conv.roles[0], prompt)\n\n        prompt_segs = prompt.split('<图片>')\n        img_list = [img]\n        seg_tokens = [\n            self.bloom_tokenizer(\n                seg, return_tensors=\"pt\", add_special_tokens=i == 0).to(self.device).input_ids\n            # only add bos to the first seg\n            for i, seg in enumerate(prompt_segs)\n        ]\n        seg_embs = [self.bloom_model.transformer.word_embeddings_layernorm(self.bloom_model.transformer.word_embeddings(seg_t)) for seg_t in seg_tokens]\n        mixed_embs = [emb for pair in zip(seg_embs[:-1], img_list) for emb in pair] + [seg_embs[-1]]\n        mixed_embs = torch.cat(mixed_embs, dim=1)\n        return mixed_embs\n\n    @classmethod\n    def from_config(cls, cfg):\n        vit_model = cfg.get(\"vit_model\", \"eva_clip_g\")\n        q_former_model = cfg.get(\"q_former_model\", \"https://storage.googleapis.com/sfr-vision-language-research/LAVIS/models/BLIP2/blip2_pretrained_flant5xxl.pth\")\n        img_size = cfg.get(\"image_size\")\n        num_query_token = cfg.get(\"num_query_token\")\n        bloom_model = cfg.get(\"bloom_model\")\n\n        drop_path_rate = cfg.get(\"drop_path_rate\", 0)\n        use_grad_checkpoint = cfg.get(\"use_grad_checkpoint\", False)\n        vit_precision = cfg.get(\"vit_precision\", \"fp16\")\n        freeze_vit = cfg.get(\"freeze_vit\", True)\n        freeze_qformer = cfg.get(\"freeze_qformer\", True)\n        low_resource = cfg.get(\"low_resource\", False)\n\n        prompt_path = cfg.get(\"prompt_path\", \"\")\n        prompt_template = cfg.get(\"prompt_template\", \"\")\n        max_txt_len = cfg.get(\"max_txt_len\", 32)\n        end_sym = cfg.get(\"end_sym\", '\\n')\n\n        model = cls(\n            vit_model=vit_model,\n            q_former_model=q_former_model,\n            img_size=img_size,\n            drop_path_rate=drop_path_rate,\n            use_grad_checkpoint=use_grad_checkpoint,\n            vit_precision=vit_precision,\n            freeze_vit=freeze_vit,\n            freeze_qformer=freeze_qformer,\n            num_query_token=num_query_token,\n            bloom_model=bloom_model,\n            prompt_path=prompt_path,\n            prompt_template=prompt_template,\n            max_txt_len=max_txt_len,\n            low_resource=low_resource,\n            end_sym=end_sym\n        )\n\n        ckpt_path = cfg.get(\"ckpt\", \"\")\n        if ckpt_path:\n            print(\"Load BLIP2-LLM Checkpoint: {}\".format(ckpt_path))\n            ckpt = torch.load(ckpt_path, map_location=\"cpu\")\n            msg = model.load_state_dict(ckpt['model'], strict=False)\n\n        return model\n"
  },
  {
    "path": "xraypulse/processors/__init__.py",
    "content": "\"\"\"\n Copyright (c) 2022, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\n\nfrom xraypulse.processors.base_processor import BaseProcessor\nfrom xraypulse.processors.blip_processors import (\n    Blip2ImageTrainProcessor,\n    Blip2ImageEvalProcessor,\n    BlipCaptionProcessor,\n)\n\nfrom xraypulse.common.registry import registry\n\n__all__ = [\n    \"BaseProcessor\",\n    \"Blip2ImageTrainProcessor\",\n    \"Blip2ImageEvalProcessor\",\n    \"BlipCaptionProcessor\",\n]\n\n\ndef load_processor(name, cfg=None):\n    \"\"\"\n    Example\n\n    >>> processor = load_processor(\"alpro_video_train\", cfg=None)\n    \"\"\"\n    processor = registry.get_processor_class(name).from_config(cfg)\n\n    return processor\n"
  },
  {
    "path": "xraypulse/processors/base_processor.py",
    "content": "\"\"\"\n Copyright (c) 2022, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\n\nfrom omegaconf import OmegaConf\n\n\nclass BaseProcessor:\n    def __init__(self):\n        self.transform = lambda x: x\n        return\n\n    def __call__(self, item):\n        return self.transform(item)\n\n    @classmethod\n    def from_config(cls, cfg=None):\n        return cls()\n\n    def build(self, **kwargs):\n        cfg = OmegaConf.create(kwargs)\n\n        return self.from_config(cfg)\n"
  },
  {
    "path": "xraypulse/processors/blip_processors.py",
    "content": "\"\"\"\n Copyright (c) 2022, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\n\nimport re\n\nfrom xraypulse.common.registry import registry\nfrom xraypulse.processors.base_processor import BaseProcessor\nfrom xraypulse.processors.randaugment import RandomAugment\nfrom omegaconf import OmegaConf\nfrom torchvision import transforms\nfrom torchvision.transforms.functional import InterpolationMode\n\n\nclass BlipImageBaseProcessor(BaseProcessor):\n    def __init__(self, mean=None, std=None):\n        if mean is None:\n            mean = (0.48145466, 0.4578275, 0.40821073)\n        if std is None:\n            std = (0.26862954, 0.26130258, 0.27577711)\n\n        self.normalize = transforms.Normalize(mean, std)\n\n\n@registry.register_processor(\"blip_caption\")\nclass BlipCaptionProcessor(BaseProcessor):\n    def __init__(self, prompt=\"\", max_words=50):\n        self.prompt = prompt\n        self.max_words = max_words\n\n    def __call__(self, caption):\n        caption = self.prompt + self.pre_caption(caption)\n\n        return caption\n\n    @classmethod\n    def from_config(cls, cfg=None):\n        if cfg is None:\n            cfg = OmegaConf.create()\n\n        prompt = cfg.get(\"prompt\", \"\")\n        max_words = cfg.get(\"max_words\", 50)\n\n        return cls(prompt=prompt, max_words=max_words)\n\n    def pre_caption(self, caption):\n        caption = re.sub(\n            r\"([.!\\\"()*#:;~])\",\n            \" \",\n            caption.lower(),\n        )\n        caption = re.sub(\n            r\"\\s{2,}\",\n            \" \",\n            caption,\n        )\n        caption = caption.rstrip(\"\\n\")\n        caption = caption.strip(\" \")\n\n        # truncate caption\n        caption_words = caption.split(\" \")\n        if len(caption_words) > self.max_words:\n            caption = \" \".join(caption_words[: self.max_words])\n\n        return caption\n\n\n@registry.register_processor(\"blip2_image_train\")\nclass Blip2ImageTrainProcessor(BlipImageBaseProcessor):\n    def __init__(self, image_size=224, mean=None, std=None, min_scale=0.5, max_scale=1.0):\n        super().__init__(mean=mean, std=std)\n\n        self.transform = transforms.Compose(\n            [\n                transforms.RandomResizedCrop(\n                    image_size,\n                    scale=(min_scale, max_scale),\n                    interpolation=InterpolationMode.BICUBIC,\n                ),\n                transforms.ToTensor(),\n                self.normalize,\n            ]\n        )\n\n    def __call__(self, item):\n        return self.transform(item)\n\n    @classmethod\n    def from_config(cls, cfg=None):\n        if cfg is None:\n            cfg = OmegaConf.create()\n\n        image_size = cfg.get(\"image_size\", 224)\n\n        mean = cfg.get(\"mean\", None)\n        std = cfg.get(\"std\", None)\n\n        min_scale = cfg.get(\"min_scale\", 0.5)\n        max_scale = cfg.get(\"max_scale\", 1.0)\n\n        return cls(\n            image_size=image_size,\n            mean=mean,\n            std=std,\n            min_scale=min_scale,\n            max_scale=max_scale,\n        )\n\n\n@registry.register_processor(\"blip2_image_eval\")\nclass Blip2ImageEvalProcessor(BlipImageBaseProcessor):\n    def __init__(self, image_size=224, mean=None, std=None):\n        super().__init__(mean=mean, std=std)\n\n        self.transform = transforms.Compose(\n            [\n                transforms.Resize(\n                    (image_size, image_size), interpolation=InterpolationMode.BICUBIC\n                ),\n                transforms.ToTensor(),\n                self.normalize,\n            ]\n        )\n\n    def __call__(self, item):\n        return self.transform(item)\n\n    @classmethod\n    def from_config(cls, cfg=None):\n        if cfg is None:\n            cfg = OmegaConf.create()\n\n        image_size = cfg.get(\"image_size\", 224)\n\n        mean = cfg.get(\"mean\", None)\n        std = cfg.get(\"std\", None)\n\n        return cls(image_size=image_size, mean=mean, std=std)"
  },
  {
    "path": "xraypulse/processors/randaugment.py",
    "content": "\"\"\"\n Copyright (c) 2022, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\n\nimport cv2\nimport numpy as np\n\nimport torch\n\n\n## aug functions\ndef identity_func(img):\n    return img\n\n\ndef autocontrast_func(img, cutoff=0):\n    \"\"\"\n    same output as PIL.ImageOps.autocontrast\n    \"\"\"\n    n_bins = 256\n\n    def tune_channel(ch):\n        n = ch.size\n        cut = cutoff * n // 100\n        if cut == 0:\n            high, low = ch.max(), ch.min()\n        else:\n            hist = cv2.calcHist([ch], [0], None, [n_bins], [0, n_bins])\n            low = np.argwhere(np.cumsum(hist) > cut)\n            low = 0 if low.shape[0] == 0 else low[0]\n            high = np.argwhere(np.cumsum(hist[::-1]) > cut)\n            high = n_bins - 1 if high.shape[0] == 0 else n_bins - 1 - high[0]\n        if high <= low:\n            table = np.arange(n_bins)\n        else:\n            scale = (n_bins - 1) / (high - low)\n            offset = -low * scale\n            table = np.arange(n_bins) * scale + offset\n            table[table < 0] = 0\n            table[table > n_bins - 1] = n_bins - 1\n        table = table.clip(0, 255).astype(np.uint8)\n        return table[ch]\n\n    channels = [tune_channel(ch) for ch in cv2.split(img)]\n    out = cv2.merge(channels)\n    return out\n\n\ndef equalize_func(img):\n    \"\"\"\n    same output as PIL.ImageOps.equalize\n    PIL's implementation is different from cv2.equalize\n    \"\"\"\n    n_bins = 256\n\n    def tune_channel(ch):\n        hist = cv2.calcHist([ch], [0], None, [n_bins], [0, n_bins])\n        non_zero_hist = hist[hist != 0].reshape(-1)\n        step = np.sum(non_zero_hist[:-1]) // (n_bins - 1)\n        if step == 0:\n            return ch\n        n = np.empty_like(hist)\n        n[0] = step // 2\n        n[1:] = hist[:-1]\n        table = (np.cumsum(n) // step).clip(0, 255).astype(np.uint8)\n        return table[ch]\n\n    channels = [tune_channel(ch) for ch in cv2.split(img)]\n    out = cv2.merge(channels)\n    return out\n\n\ndef rotate_func(img, degree, fill=(0, 0, 0)):\n    \"\"\"\n    like PIL, rotate by degree, not radians\n    \"\"\"\n    H, W = img.shape[0], img.shape[1]\n    center = W / 2, H / 2\n    M = cv2.getRotationMatrix2D(center, degree, 1)\n    out = cv2.warpAffine(img, M, (W, H), borderValue=fill)\n    return out\n\n\ndef solarize_func(img, thresh=128):\n    \"\"\"\n    same output as PIL.ImageOps.posterize\n    \"\"\"\n    table = np.array([el if el < thresh else 255 - el for el in range(256)])\n    table = table.clip(0, 255).astype(np.uint8)\n    out = table[img]\n    return out\n\n\ndef color_func(img, factor):\n    \"\"\"\n    same output as PIL.ImageEnhance.Color\n    \"\"\"\n    ## implementation according to PIL definition, quite slow\n    #  degenerate = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)[:, :, np.newaxis]\n    #  out = blend(degenerate, img, factor)\n    #  M = (\n    #      np.eye(3) * factor\n    #      + np.float32([0.114, 0.587, 0.299]).reshape(3, 1) * (1. - factor)\n    #  )[np.newaxis, np.newaxis, :]\n    M = np.float32(\n        [[0.886, -0.114, -0.114], [-0.587, 0.413, -0.587], [-0.299, -0.299, 0.701]]\n    ) * factor + np.float32([[0.114], [0.587], [0.299]])\n    out = np.matmul(img, M).clip(0, 255).astype(np.uint8)\n    return out\n\n\ndef contrast_func(img, factor):\n    \"\"\"\n    same output as PIL.ImageEnhance.Contrast\n    \"\"\"\n    mean = np.sum(np.mean(img, axis=(0, 1)) * np.array([0.114, 0.587, 0.299]))\n    table = (\n        np.array([(el - mean) * factor + mean for el in range(256)])\n        .clip(0, 255)\n        .astype(np.uint8)\n    )\n    out = table[img]\n    return out\n\n\ndef brightness_func(img, factor):\n    \"\"\"\n    same output as PIL.ImageEnhance.Contrast\n    \"\"\"\n    table = (np.arange(256, dtype=np.float32) * factor).clip(0, 255).astype(np.uint8)\n    out = table[img]\n    return out\n\n\ndef sharpness_func(img, factor):\n    \"\"\"\n    The differences the this result and PIL are all on the 4 boundaries, the center\n    areas are same\n    \"\"\"\n    kernel = np.ones((3, 3), dtype=np.float32)\n    kernel[1][1] = 5\n    kernel /= 13\n    degenerate = cv2.filter2D(img, -1, kernel)\n    if factor == 0.0:\n        out = degenerate\n    elif factor == 1.0:\n        out = img\n    else:\n        out = img.astype(np.float32)\n        degenerate = degenerate.astype(np.float32)[1:-1, 1:-1, :]\n        out[1:-1, 1:-1, :] = degenerate + factor * (out[1:-1, 1:-1, :] - degenerate)\n        out = out.astype(np.uint8)\n    return out\n\n\ndef shear_x_func(img, factor, fill=(0, 0, 0)):\n    H, W = img.shape[0], img.shape[1]\n    M = np.float32([[1, factor, 0], [0, 1, 0]])\n    out = cv2.warpAffine(\n        img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR\n    ).astype(np.uint8)\n    return out\n\n\ndef translate_x_func(img, offset, fill=(0, 0, 0)):\n    \"\"\"\n    same output as PIL.Image.transform\n    \"\"\"\n    H, W = img.shape[0], img.shape[1]\n    M = np.float32([[1, 0, -offset], [0, 1, 0]])\n    out = cv2.warpAffine(\n        img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR\n    ).astype(np.uint8)\n    return out\n\n\ndef translate_y_func(img, offset, fill=(0, 0, 0)):\n    \"\"\"\n    same output as PIL.Image.transform\n    \"\"\"\n    H, W = img.shape[0], img.shape[1]\n    M = np.float32([[1, 0, 0], [0, 1, -offset]])\n    out = cv2.warpAffine(\n        img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR\n    ).astype(np.uint8)\n    return out\n\n\ndef posterize_func(img, bits):\n    \"\"\"\n    same output as PIL.ImageOps.posterize\n    \"\"\"\n    out = np.bitwise_and(img, np.uint8(255 << (8 - bits)))\n    return out\n\n\ndef shear_y_func(img, factor, fill=(0, 0, 0)):\n    H, W = img.shape[0], img.shape[1]\n    M = np.float32([[1, 0, 0], [factor, 1, 0]])\n    out = cv2.warpAffine(\n        img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR\n    ).astype(np.uint8)\n    return out\n\n\ndef cutout_func(img, pad_size, replace=(0, 0, 0)):\n    replace = np.array(replace, dtype=np.uint8)\n    H, W = img.shape[0], img.shape[1]\n    rh, rw = np.random.random(2)\n    pad_size = pad_size // 2\n    ch, cw = int(rh * H), int(rw * W)\n    x1, x2 = max(ch - pad_size, 0), min(ch + pad_size, H)\n    y1, y2 = max(cw - pad_size, 0), min(cw + pad_size, W)\n    out = img.copy()\n    out[x1:x2, y1:y2, :] = replace\n    return out\n\n\n### level to args\ndef enhance_level_to_args(MAX_LEVEL):\n    def level_to_args(level):\n        return ((level / MAX_LEVEL) * 1.8 + 0.1,)\n\n    return level_to_args\n\n\ndef shear_level_to_args(MAX_LEVEL, replace_value):\n    def level_to_args(level):\n        level = (level / MAX_LEVEL) * 0.3\n        if np.random.random() > 0.5:\n            level = -level\n        return (level, replace_value)\n\n    return level_to_args\n\n\ndef translate_level_to_args(translate_const, MAX_LEVEL, replace_value):\n    def level_to_args(level):\n        level = (level / MAX_LEVEL) * float(translate_const)\n        if np.random.random() > 0.5:\n            level = -level\n        return (level, replace_value)\n\n    return level_to_args\n\n\ndef cutout_level_to_args(cutout_const, MAX_LEVEL, replace_value):\n    def level_to_args(level):\n        level = int((level / MAX_LEVEL) * cutout_const)\n        return (level, replace_value)\n\n    return level_to_args\n\n\ndef solarize_level_to_args(MAX_LEVEL):\n    def level_to_args(level):\n        level = int((level / MAX_LEVEL) * 256)\n        return (level,)\n\n    return level_to_args\n\n\ndef none_level_to_args(level):\n    return ()\n\n\ndef posterize_level_to_args(MAX_LEVEL):\n    def level_to_args(level):\n        level = int((level / MAX_LEVEL) * 4)\n        return (level,)\n\n    return level_to_args\n\n\ndef rotate_level_to_args(MAX_LEVEL, replace_value):\n    def level_to_args(level):\n        level = (level / MAX_LEVEL) * 30\n        if np.random.random() < 0.5:\n            level = -level\n        return (level, replace_value)\n\n    return level_to_args\n\n\nfunc_dict = {\n    \"Identity\": identity_func,\n    \"AutoContrast\": autocontrast_func,\n    \"Equalize\": equalize_func,\n    \"Rotate\": rotate_func,\n    \"Solarize\": solarize_func,\n    \"Color\": color_func,\n    \"Contrast\": contrast_func,\n    \"Brightness\": brightness_func,\n    \"Sharpness\": sharpness_func,\n    \"ShearX\": shear_x_func,\n    \"TranslateX\": translate_x_func,\n    \"TranslateY\": translate_y_func,\n    \"Posterize\": posterize_func,\n    \"ShearY\": shear_y_func,\n}\n\ntranslate_const = 10\nMAX_LEVEL = 10\nreplace_value = (128, 128, 128)\narg_dict = {\n    \"Identity\": none_level_to_args,\n    \"AutoContrast\": none_level_to_args,\n    \"Equalize\": none_level_to_args,\n    \"Rotate\": rotate_level_to_args(MAX_LEVEL, replace_value),\n    \"Solarize\": solarize_level_to_args(MAX_LEVEL),\n    \"Color\": enhance_level_to_args(MAX_LEVEL),\n    \"Contrast\": enhance_level_to_args(MAX_LEVEL),\n    \"Brightness\": enhance_level_to_args(MAX_LEVEL),\n    \"Sharpness\": enhance_level_to_args(MAX_LEVEL),\n    \"ShearX\": shear_level_to_args(MAX_LEVEL, replace_value),\n    \"TranslateX\": translate_level_to_args(translate_const, MAX_LEVEL, replace_value),\n    \"TranslateY\": translate_level_to_args(translate_const, MAX_LEVEL, replace_value),\n    \"Posterize\": posterize_level_to_args(MAX_LEVEL),\n    \"ShearY\": shear_level_to_args(MAX_LEVEL, replace_value),\n}\n\n\nclass RandomAugment(object):\n    def __init__(self, N=2, M=10, isPIL=False, augs=[]):\n        self.N = N\n        self.M = M\n        self.isPIL = isPIL\n        if augs:\n            self.augs = augs\n        else:\n            self.augs = list(arg_dict.keys())\n\n    def get_random_ops(self):\n        sampled_ops = np.random.choice(self.augs, self.N)\n        return [(op, 0.5, self.M) for op in sampled_ops]\n\n    def __call__(self, img):\n        if self.isPIL:\n            img = np.array(img)\n        ops = self.get_random_ops()\n        for name, prob, level in ops:\n            if np.random.random() > prob:\n                continue\n            args = arg_dict[name](level)\n            img = func_dict[name](img, *args)\n        return img\n\n\nclass VideoRandomAugment(object):\n    def __init__(self, N=2, M=10, p=0.0, tensor_in_tensor_out=True, augs=[]):\n        self.N = N\n        self.M = M\n        self.p = p\n        self.tensor_in_tensor_out = tensor_in_tensor_out\n        if augs:\n            self.augs = augs\n        else:\n            self.augs = list(arg_dict.keys())\n\n    def get_random_ops(self):\n        sampled_ops = np.random.choice(self.augs, self.N, replace=False)\n        return [(op, self.M) for op in sampled_ops]\n\n    def __call__(self, frames):\n        assert (\n            frames.shape[-1] == 3\n        ), \"Expecting last dimension for 3-channels RGB (b, h, w, c).\"\n\n        if self.tensor_in_tensor_out:\n            frames = frames.numpy().astype(np.uint8)\n\n        num_frames = frames.shape[0]\n\n        ops = num_frames * [self.get_random_ops()]\n        apply_or_not = num_frames * [np.random.random(size=self.N) > self.p]\n\n        frames = torch.stack(\n            list(map(self._aug, frames, ops, apply_or_not)), dim=0\n        ).float()\n\n        return frames\n\n    def _aug(self, img, ops, apply_or_not):\n        for i, (name, level) in enumerate(ops):\n            if not apply_or_not[i]:\n                continue\n            args = arg_dict[name](level)\n            img = func_dict[name](img, *args)\n        return torch.from_numpy(img)\n\n\nif __name__ == \"__main__\":\n    a = RandomAugment()\n    img = np.random.randn(32, 32, 3)\n    a(img)\n"
  },
  {
    "path": "xraypulse/runners/__init__.py",
    "content": "\"\"\"\n Copyright (c) 2022, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\n\nfrom xraypulse.runners.runner_base import RunnerBase\n\n__all__ = [\"RunnerBase\"]\n"
  },
  {
    "path": "xraypulse/runners/runner_base.py",
    "content": "\"\"\"\n Copyright (c) 2022, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\n\nimport datetime\nimport json\nimport logging\nimport os\nimport time\nfrom pathlib import Path\n\nimport torch\nimport torch.distributed as dist\nimport webdataset as wds\nfrom xraypulse.common.dist_utils import (\n    download_cached_file,\n    get_rank,\n    get_world_size,\n    is_main_process,\n    main_process,\n)\nfrom xraypulse.common.registry import registry\nfrom xraypulse.common.utils import is_url\nfrom xraypulse.datasets.data_utils import concat_datasets, reorg_datasets_by_split, ChainDataset\nfrom xraypulse.datasets.datasets.dataloader_utils import (\n    IterLoader,\n    MultiIterLoader,\n    PrefetchLoader,\n)\nfrom torch.nn.parallel import DistributedDataParallel as DDP\nfrom torch.utils.data import DataLoader, DistributedSampler\n\n\n@registry.register_runner(\"runner_base\")\nclass RunnerBase:\n    \"\"\"\n    A runner class to train and evaluate a model given a task and datasets.\n\n    The runner uses pytorch distributed data parallel by default. Future release\n    will support other distributed frameworks.\n    \"\"\"\n\n    def __init__(self, cfg, task, model, datasets, job_id):\n        self.config = cfg\n        self.job_id = job_id\n\n        self.task = task\n        self.datasets = datasets\n\n        self._model = model\n\n        self._wrapped_model = None\n        self._device = None\n        self._optimizer = None\n        self._scaler = None\n        self._dataloaders = None\n        self._lr_sched = None\n\n        self.start_epoch = 0\n\n        # self.setup_seeds()\n        self.setup_output_dir()\n\n    @property\n    def device(self):\n        if self._device is None:\n            self._device = torch.device(self.config.run_cfg.device)\n\n        return self._device\n\n    @property\n    def use_distributed(self):\n        return self.config.run_cfg.distributed\n\n    @property\n    def model(self):\n        \"\"\"\n        A property to get the DDP-wrapped model on the device.\n        \"\"\"\n        # move model to device\n        if self._model.device != self.device:\n            self._model = self._model.to(self.device)\n\n            # distributed training wrapper\n            if self.use_distributed:\n                if self._wrapped_model is None:\n                    self._wrapped_model = DDP(\n                        self._model, device_ids=[self.config.run_cfg.gpu]\n                    )\n            else:\n                self._wrapped_model = self._model\n\n        return self._wrapped_model\n\n    @property\n    def optimizer(self):\n        # TODO make optimizer class and configurations\n        if self._optimizer is None:\n            num_parameters = 0\n            p_wd, p_non_wd = [], []\n            for n, p in self.model.named_parameters():\n                if not p.requires_grad:\n                    continue  # frozen weights\n                print(n)\n                if p.ndim < 2 or \"bias\" in n or \"ln\" in n or \"bn\" in n:\n                    p_non_wd.append(p)\n                else:\n                    p_wd.append(p)\n                num_parameters += p.data.nelement()\n            logging.info(\"number of trainable parameters: %d\" % num_parameters)\n            optim_params = [\n                {\n                    \"params\": p_wd,\n                    \"weight_decay\": float(self.config.run_cfg.weight_decay),\n                },\n                {\"params\": p_non_wd, \"weight_decay\": 0},\n            ]\n            beta2 = self.config.run_cfg.get(\"beta2\", 0.999)\n            self._optimizer = torch.optim.AdamW(\n                optim_params,\n                lr=float(self.config.run_cfg.init_lr),\n                weight_decay=float(self.config.run_cfg.weight_decay),\n                betas=(0.9, beta2),\n            )\n\n        return self._optimizer\n\n    @property\n    def scaler(self):\n        amp = self.config.run_cfg.get(\"amp\", False)\n\n        if amp:\n            if self._scaler is None:\n                self._scaler = torch.cuda.amp.GradScaler()\n\n        return self._scaler\n\n    @property\n    def lr_scheduler(self):\n        \"\"\"\n        A property to get and create learning rate scheduler by split just in need.\n        \"\"\"\n        if self._lr_sched is None:\n            lr_sched_cls = registry.get_lr_scheduler_class(self.config.run_cfg.lr_sched)\n\n            # max_epoch = self.config.run_cfg.max_epoch\n            max_epoch = self.max_epoch\n            # min_lr = self.config.run_cfg.min_lr\n            min_lr = self.min_lr\n            # init_lr = self.config.run_cfg.init_lr\n            init_lr = self.init_lr\n\n            # optional parameters\n            decay_rate = self.config.run_cfg.get(\"lr_decay_rate\", None)\n            warmup_start_lr = self.config.run_cfg.get(\"warmup_lr\", -1)\n            warmup_steps = self.config.run_cfg.get(\"warmup_steps\", 0)\n            iters_per_epoch = self.config.run_cfg.get(\"iters_per_epoch\", None)\n\n            if iters_per_epoch is None:\n                try:\n                    iters_per_epoch = len(self.dataloaders['train'])\n                except (AttributeError, TypeError):\n                    iters_per_epoch = 10000\n\n            self._lr_sched = lr_sched_cls(\n                optimizer=self.optimizer,\n                max_epoch=max_epoch,\n                iters_per_epoch=iters_per_epoch,\n                min_lr=min_lr,\n                init_lr=init_lr,\n                decay_rate=decay_rate,\n                warmup_start_lr=warmup_start_lr,\n                warmup_steps=warmup_steps,\n            )\n\n        return self._lr_sched\n\n    @property\n    def dataloaders(self) -> dict:\n        \"\"\"\n        A property to get and create dataloaders by split just in need.\n\n        If no train_dataset_ratio is provided, concatenate map-style datasets and\n        chain wds.DataPipe datasets separately. Training set becomes a tuple\n        (ConcatDataset, ChainDataset), both are optional but at least one of them is\n        required. The resultant ConcatDataset and ChainDataset will be sampled evenly.\n\n        If train_dataset_ratio is provided, create a MultiIterLoader to sample\n        each dataset by ratios during training.\n\n        Currently do not support multiple datasets for validation and test.\n\n        Returns:\n            dict: {split_name: (tuples of) dataloader}\n        \"\"\"\n        if self._dataloaders is None:\n\n            # concatenate map-style datasets and chain wds.DataPipe datasets separately\n            # training set becomes a tuple (ConcatDataset, ChainDataset), both are\n            # optional but at least one of them is required. The resultant ConcatDataset\n            # and ChainDataset will be sampled evenly.\n            logging.info(\n                \"dataset_ratios not specified, datasets will be concatenated (map-style datasets) or chained (webdataset.DataPipeline).\"\n            )\n\n            datasets = reorg_datasets_by_split(self.datasets)\n            self.datasets = datasets\n            # self.datasets = concat_datasets(datasets)\n\n            # print dataset statistics after concatenation/chaining\n            for split_name in self.datasets:\n                if isinstance(self.datasets[split_name], tuple) or isinstance(\n                    self.datasets[split_name], list\n                ):\n                    # mixed wds.DataPipeline and torch.utils.data.Dataset\n                    num_records = sum(\n                        [\n                            len(d)\n                            if not type(d) in [wds.DataPipeline, ChainDataset]\n                            else 0\n                            for d in self.datasets[split_name]\n                        ]\n                    )\n\n                else:\n                    if hasattr(self.datasets[split_name], \"__len__\"):\n                        # a single map-style dataset\n                        num_records = len(self.datasets[split_name])\n                    else:\n                        # a single wds.DataPipeline\n                        num_records = -1\n                        logging.info(\n                            \"Only a single wds.DataPipeline dataset, no __len__ attribute.\"\n                        )\n\n                if num_records >= 0:\n                    logging.info(\n                        \"Loaded {} records for {} split from the dataset.\".format(\n                            num_records, split_name\n                        )\n                    )\n\n            # create dataloaders\n            split_names = sorted(self.datasets.keys())\n\n            datasets = [self.datasets[split] for split in split_names]\n            is_trains = [split in self.train_splits for split in split_names]\n\n            batch_sizes = [\n                self.config.run_cfg.batch_size_train\n                if split == \"train\"\n                else self.config.run_cfg.batch_size_eval\n                for split in split_names\n            ]\n\n            collate_fns = []\n            for dataset in datasets:\n                if isinstance(dataset, tuple) or isinstance(dataset, list):\n                    collate_fns.append([getattr(d, \"collater\", None) for d in dataset])\n                else:\n                    collate_fns.append(getattr(dataset, \"collater\", None))\n\n            dataloaders = self.create_loaders(\n                datasets=datasets,\n                num_workers=self.config.run_cfg.num_workers,\n                batch_sizes=batch_sizes,\n                is_trains=is_trains,\n                collate_fns=collate_fns,\n            )\n\n            self._dataloaders = {k: v for k, v in zip(split_names, dataloaders)}\n\n        return self._dataloaders\n\n    @property\n    def cuda_enabled(self):\n        return self.device.type == \"cuda\"\n\n    @property\n    def max_epoch(self):\n        return int(self.config.run_cfg.max_epoch)\n\n    @property\n    def log_freq(self):\n        log_freq = self.config.run_cfg.get(\"log_freq\", 50)\n        return int(log_freq)\n\n    @property\n    def init_lr(self):\n        return float(self.config.run_cfg.init_lr)\n\n    @property\n    def min_lr(self):\n        return float(self.config.run_cfg.min_lr)\n\n    @property\n    def accum_grad_iters(self):\n        return int(self.config.run_cfg.get(\"accum_grad_iters\", 1))\n\n    @property\n    def valid_splits(self):\n        valid_splits = self.config.run_cfg.get(\"valid_splits\", [])\n\n        if len(valid_splits) == 0:\n            logging.info(\"No validation splits found.\")\n\n        return valid_splits\n\n    @property\n    def test_splits(self):\n        test_splits = self.config.run_cfg.get(\"test_splits\", [])\n\n        return test_splits\n\n    @property\n    def train_splits(self):\n        train_splits = self.config.run_cfg.get(\"train_splits\", [])\n\n        if len(train_splits) == 0:\n            logging.info(\"Empty train splits.\")\n\n        return train_splits\n\n    @property\n    def evaluate_only(self):\n        \"\"\"\n        Set to True to skip training.\n        \"\"\"\n        return self.config.run_cfg.evaluate\n\n    @property\n    def use_dist_eval_sampler(self):\n        return self.config.run_cfg.get(\"use_dist_eval_sampler\", True)\n\n    @property\n    def resume_ckpt_path(self):\n        return self.config.run_cfg.get(\"resume_ckpt_path\", None)\n\n    @property\n    def train_loader(self):\n        train_dataloader = self.dataloaders[\"train\"]\n\n        return train_dataloader\n\n    def setup_output_dir(self):\n        lib_root = Path(registry.get_path(\"library_root\"))\n\n        output_dir = lib_root / self.config.run_cfg.output_dir / self.job_id\n        result_dir = output_dir / \"result\"\n\n        output_dir.mkdir(parents=True, exist_ok=True)\n        result_dir.mkdir(parents=True, exist_ok=True)\n\n        registry.register_path(\"result_dir\", str(result_dir))\n        registry.register_path(\"output_dir\", str(output_dir))\n\n        self.result_dir = result_dir\n        self.output_dir = output_dir\n\n    def train(self):\n        start_time = time.time()\n        best_agg_metric = 0\n        best_epoch = 0\n\n        self.log_config()\n        # resume from checkpoint if specified\n        if not self.evaluate_only and self.resume_ckpt_path is not None:\n            self._load_checkpoint(self.resume_ckpt_path)\n\n        for cur_epoch in range(self.start_epoch, self.max_epoch):\n            # training phase\n            if not self.evaluate_only:\n                logging.info(\"Start training\")\n                train_stats = self.train_epoch(cur_epoch)\n                self.log_stats(split_name=\"train\", stats=train_stats)\n\n            # evaluation phase\n            if len(self.valid_splits) > 0:\n                for split_name in self.valid_splits:\n                    logging.info(\"Evaluating on {}.\".format(split_name))\n\n                    val_log = self.eval_epoch(\n                        split_name=split_name, cur_epoch=cur_epoch\n                    )\n                    if val_log is not None:\n                        if is_main_process():\n                            assert (\n                                \"agg_metrics\" in val_log\n                            ), \"No agg_metrics found in validation log.\"\n\n                            agg_metrics = val_log[\"agg_metrics\"]\n                            if agg_metrics > best_agg_metric and split_name == \"val\":\n                                best_epoch, best_agg_metric = cur_epoch, agg_metrics\n\n                                self._save_checkpoint(cur_epoch, is_best=True)\n\n                            val_log.update({\"best_epoch\": best_epoch})\n                            self.log_stats(val_log, split_name)\n\n            else:\n                # if no validation split is provided, we just save the checkpoint at the end of each epoch.\n                if not self.evaluate_only:\n                    self._save_checkpoint(cur_epoch, is_best=False)\n\n            if self.evaluate_only:\n                break\n\n            if self.config.run_cfg.distributed:\n                dist.barrier()\n\n        # # testing phase\n        # test_epoch = \"best\" if len(self.valid_splits) > 0 else cur_epoch\n        # self.evaluate(cur_epoch=test_epoch, skip_reload=self.evaluate_only)\n\n        total_time = time.time() - start_time\n        total_time_str = str(datetime.timedelta(seconds=int(total_time)))\n        logging.info(\"Training time {}\".format(total_time_str))\n\n    def test(self):\n        start_time = time.time()\n\n        # resume from checkpoint if specified\n        if not self.evaluate_only and self.resume_ckpt_path is not None:\n            self._load_checkpoint(self.resume_ckpt_path)\n\n        test_stats = self.test_epoch(1)\n                \n        total_time = time.time() - start_time\n        total_time_str = str(datetime.timedelta(seconds=int(total_time)))\n        print(\"Testing time {}\".format(total_time_str))\n\n    def evaluate(self, ckpt, cur_epoch=\"best\", skip_reload=False):\n        test_logs = dict()\n\n        if len(self.test_splits) > 0:\n            for split_name in self.test_splits:\n                test_logs[split_name] = self.eval_epoch(\n                    ckpt, split_name=split_name, cur_epoch=cur_epoch, skip_reload=skip_reload\n                )\n\n            return test_logs\n\n    def train_epoch(self, epoch):\n        # train\n        self.model.train()\n\n        return self.task.train_epoch(\n            epoch=epoch,\n            model=self.model,\n            data_loader=self.train_loader,\n            optimizer=self.optimizer,\n            scaler=self.scaler,\n            lr_scheduler=self.lr_scheduler,\n            cuda_enabled=self.cuda_enabled,\n            log_freq=self.log_freq,\n            accum_grad_iters=self.accum_grad_iters,\n        )\n    \n    def test_epoch(self, epoch):\n        # train\n        self.model.eval()\n\n        return self.task.test_epoch(\n            epoch=epoch,\n            model=self.model,\n            data_loader=self.train_loader,\n            optimizer=self.optimizer,\n            scaler=self.scaler,\n            lr_scheduler=self.lr_scheduler,\n            cuda_enabled=self.cuda_enabled,\n            log_freq=self.log_freq,\n            accum_grad_iters=self.accum_grad_iters,\n        )\n\n    @torch.no_grad()\n    def eval_epoch(self, ckpt, split_name, cur_epoch, skip_reload=False):\n        \"\"\"\n        Evaluate the model on a given split.\n\n        Args:\n            split_name (str): name of the split to evaluate on.\n            cur_epoch (int): current epoch.\n            skip_reload_best (bool): whether to skip reloading the best checkpoint.\n                During training, we will reload the best checkpoint for validation.\n                During testing, we will use provided weights and skip reloading the best checkpoint .\n        \"\"\"\n        data_loader = self.dataloaders.get(split_name, None)\n        assert data_loader, \"data_loader for split {} is None.\".format(split_name)\n\n        # TODO In validation, you need to compute loss as well as metrics\n        # TODO consider moving to model.before_evaluation()\n        model = self.unwrap_dist_model(self.model)\n        if not skip_reload and cur_epoch == \"best\":\n            model = self._reload_model(model, ckpt)\n        model.eval()\n\n        self.task.before_evaluation(\n            model=model,\n            dataset=self.datasets[split_name],\n        )\n        results = self.task.evaluation(model, data_loader)\n\n        if results is not None:\n            return self.task.after_evaluation(\n                val_result=results,\n                split_name=split_name,\n                epoch=cur_epoch,\n            )\n\n    def unwrap_dist_model(self, model):\n        if self.use_distributed:\n            return model.module\n        else:\n            return model\n\n    def create_loaders(\n        self,\n        datasets,\n        num_workers,\n        batch_sizes,\n        is_trains,\n        collate_fns,\n        dataset_ratios=None,\n    ):\n        \"\"\"\n        Create dataloaders for training and validation.\n        \"\"\"\n\n        def _create_loader(dataset, num_workers, bsz, is_train, collate_fn):\n            # create a single dataloader for each split\n            if isinstance(dataset, ChainDataset) or isinstance(\n                dataset, wds.DataPipeline\n            ):\n                # wds.WebdDataset instance are chained together\n                # webdataset.DataPipeline has its own sampler and collate_fn\n                loader = iter(\n                    DataLoader(\n                        dataset,\n                        batch_size=bsz,\n                        num_workers=num_workers,\n                        pin_memory=True,\n                    )\n                )\n            else:\n                # map-style dataset are concatenated together\n                # setup distributed sampler\n                if self.use_distributed:\n                    sampler = DistributedSampler(\n                        dataset,\n                        shuffle=is_train,\n                        num_replicas=get_world_size(),\n                        rank=get_rank(),\n                    )\n                    if not self.use_dist_eval_sampler:\n                        # e.g. retrieval evaluation\n                        sampler = sampler if is_train else None\n                else:\n                    sampler = None\n\n                loader = DataLoader(\n                    dataset,\n                    batch_size=bsz,\n                    num_workers=num_workers,\n                    pin_memory=True,\n                    sampler=sampler,\n                    shuffle=sampler is None and is_train,\n                    collate_fn=collate_fn,\n                    drop_last=True if is_train else False,\n                )\n                loader = PrefetchLoader(loader)\n\n                if is_train:\n                    loader = IterLoader(loader, use_distributed=self.use_distributed)\n\n            return loader\n\n        loaders = []\n\n        for dataset, bsz, is_train, collate_fn in zip(\n            datasets, batch_sizes, is_trains, collate_fns\n        ):\n            if isinstance(dataset, list) or isinstance(dataset, tuple):\n                if hasattr(dataset[0], 'sample_ratio') and dataset_ratios is None:\n                    dataset_ratios = [d.sample_ratio for d in dataset]\n                loader = MultiIterLoader(\n                    loaders=[\n                        _create_loader(d, num_workers, bsz, is_train, collate_fn[i])\n                        for i, d in enumerate(dataset)\n                    ],\n                    ratios=dataset_ratios,\n                )\n            else:\n                loader = _create_loader(dataset, num_workers, bsz, is_train, collate_fn)\n\n            loaders.append(loader)\n\n        return loaders\n\n    @main_process\n    def _save_checkpoint(self, cur_epoch, is_best=False):\n        \"\"\"\n        Save the checkpoint at the current epoch.\n        \"\"\"\n        model_no_ddp = self.unwrap_dist_model(self.model)\n        param_grad_dic = {\n            k: v.requires_grad for (k, v) in model_no_ddp.named_parameters()\n        }\n        state_dict = model_no_ddp.state_dict()\n        for k in list(state_dict.keys()):\n            if k in param_grad_dic.keys() and not param_grad_dic[k]:\n                # delete parameters that do not require gradient\n                del state_dict[k]\n        save_obj = {\n            \"model\": state_dict,\n            \"optimizer\": self.optimizer.state_dict(),\n            \"config\": self.config.to_dict(),\n            \"scaler\": self.scaler.state_dict() if self.scaler else None,\n            \"epoch\": cur_epoch,\n        }\n        save_to = os.path.join(\n            self.output_dir,\n            \"checkpoint_{}.pth\".format(\"best\" if is_best else cur_epoch),\n        )\n        logging.info(\"Saving checkpoint at epoch {} to {}.\".format(cur_epoch, save_to))\n        torch.save(save_obj, save_to)\n\n    def _reload_best_model(self, model):\n        \"\"\"\n        Load the best checkpoint for evaluation.\n        \"\"\"\n        checkpoint_path = os.path.join(self.output_dir, \"checkpoint_best.pth\")\n\n        logging.info(\"Loading checkpoint from {}.\".format(checkpoint_path))\n        checkpoint = torch.load(checkpoint_path, map_location=\"cpu\")\n        try:\n            model.load_state_dict(checkpoint[\"model\"])\n        except RuntimeError as e:\n            logging.warning(\n                \"\"\"\n                Key mismatch when loading checkpoint. This is expected if only part of the model is saved.\n                Trying to load the model with strict=False.\n                \"\"\"\n            )\n            model.load_state_dict(checkpoint[\"model\"], strict=False)\n        return model\n    \n    def _reload_model(self, model,ckpt):\n        \"\"\"\n        Load the best checkpoint for evaluation.\n        \"\"\"\n\n        logging.info(\"Loading checkpoint from {}.\".format(ckpt))\n        checkpoint = torch.load(ckpt, map_location=\"cpu\")\n        try:\n            model.load_state_dict(checkpoint[\"model\"])\n        except RuntimeError as e:\n            logging.warning(\n                \"\"\"\n                Key mismatch when loading checkpoint. This is expected if only part of the model is saved.\n                Trying to load the model with strict=False.\n                \"\"\"\n            )\n            model.load_state_dict(checkpoint[\"model\"], strict=False)\n        return model\n\n    def _load_checkpoint(self, url_or_filename):\n        \"\"\"\n        Resume from a checkpoint.\n        \"\"\"\n        if is_url(url_or_filename):\n            cached_file = download_cached_file(\n                url_or_filename, check_hash=False, progress=True\n            )\n            checkpoint = torch.load(cached_file, map_location=self.device)\n        elif os.path.isfile(url_or_filename):\n            checkpoint = torch.load(url_or_filename, map_location=self.device)\n        else:\n            raise RuntimeError(\"checkpoint url or path is invalid\")\n\n        state_dict = checkpoint[\"model\"]\n        self.unwrap_dist_model(self.model).load_state_dict(state_dict,strict=False)\n\n        self.optimizer.load_state_dict(checkpoint[\"optimizer\"])\n        if self.scaler and \"scaler\" in checkpoint:\n            self.scaler.load_state_dict(checkpoint[\"scaler\"])\n\n        self.start_epoch = checkpoint[\"epoch\"] + 1\n        logging.info(\"Resume checkpoint from {}\".format(url_or_filename))\n\n    @main_process\n    def log_stats(self, stats, split_name):\n        if isinstance(stats, dict):\n            log_stats = {**{f\"{split_name}_{k}\": v for k, v in stats.items()}}\n            with open(os.path.join(self.output_dir, \"log.txt\"), \"a\") as f:\n                f.write(json.dumps(log_stats) + \"\\n\")\n        elif isinstance(stats, list):\n            pass\n\n    @main_process\n    def log_config(self):\n        with open(os.path.join(self.output_dir, \"log.txt\"), \"a\") as f:\n            f.write(json.dumps(self.config.to_dict(), indent=4) + \"\\n\")\n"
  },
  {
    "path": "xraypulse/tasks/__init__.py",
    "content": "\"\"\"\n Copyright (c) 2022, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\n\nfrom xraypulse.common.registry import registry\nfrom xraypulse.tasks.base_task import BaseTask\nfrom xraypulse.tasks.image_text_pretrain import ImageTextPretrainTask\n\n\ndef setup_task(cfg):\n    assert \"task\" in cfg.run_cfg, \"Task name must be provided.\"\n\n    task_name = cfg.run_cfg.task\n    task = registry.get_task_class(task_name).setup_task(cfg=cfg)\n    assert task is not None, \"Task {} not properly registered.\".format(task_name)\n\n    return task\n\n\n__all__ = [\n    \"BaseTask\",\n    \"ImageTextPretrainTask\",\n]\n"
  },
  {
    "path": "xraypulse/tasks/base_task.py",
    "content": "\"\"\"\n Copyright (c) 2022, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\n\nimport logging\nimport os\n\nimport torch\nimport torch.distributed as dist\nfrom xraypulse.common.dist_utils import get_rank, get_world_size, is_main_process, is_dist_avail_and_initialized\nfrom xraypulse.common.logger import MetricLogger, SmoothedValue\nfrom xraypulse.common.registry import registry\nfrom xraypulse.datasets.data_utils import prepare_sample\n\nimport csv #for imprgpt rogue eval\n\n\nclass BaseTask:\n    def __init__(self, **kwargs):\n        super().__init__()\n\n        self.inst_id_key = \"instance_id\"\n\n    @classmethod\n    def setup_task(cls, **kwargs):\n        return cls()\n\n    def build_model(self, cfg):\n        model_config = cfg.model_cfg\n\n        model_cls = registry.get_model_class(model_config.arch)\n        return model_cls.from_config(model_config)\n\n    def build_datasets(self, cfg):\n        \"\"\"\n        Build a dictionary of datasets, keyed by split 'train', 'valid', 'test'.\n        Download dataset and annotations automatically if not exist.\n\n        Args:\n            cfg (common.config.Config): _description_\n\n        Returns:\n            dict: Dictionary of torch.utils.data.Dataset objects by split.\n        \"\"\"\n\n        datasets = dict()\n\n        datasets_config = cfg.datasets_cfg\n\n        assert len(datasets_config) > 0, \"At least one dataset has to be specified.\"\n\n        for name in datasets_config:\n            dataset_config = datasets_config[name]\n\n            builder = registry.get_builder_class(name)(dataset_config)\n            dataset = builder.build_datasets()\n\n            dataset['train'].name = name\n            if 'sample_ratio' in dataset_config:\n                dataset['train'].sample_ratio = dataset_config.sample_ratio\n\n            datasets[name] = dataset\n\n        return datasets\n\n    def train_step(self, model, samples):\n        loss = model(samples)[\"loss\"]\n        return loss\n    \n    def test_step(self, model, samples):\n        output_text, output_token = model.test(samples)\n        return output_text\n\n    def valid_step(self, model, samples):\n        raise NotImplementedError\n\n    def before_evaluation(self, model, dataset, **kwargs):\n        model.before_evaluation(dataset=dataset, task_type=type(self))\n\n    def after_evaluation(self, **kwargs):\n        pass\n\n    def inference_step(self):\n        raise NotImplementedError\n\n    def evaluation(self, model, data_loader, cuda_enabled=True):\n        metric_logger = MetricLogger(delimiter=\"  \")\n        header = \"Evaluation\"\n        # TODO make it configurable\n        print_freq = 10\n\n        results = []\n\n        for samples in metric_logger.log_every(data_loader, print_freq, header):\n            samples = prepare_sample(samples, cuda_enabled=cuda_enabled)\n\n            eval_output = self.valid_step(model=model, samples=samples)\n            results.extend(eval_output)\n\n        if is_dist_avail_and_initialized():\n            dist.barrier()\n\n        return results\n\n    def train_epoch(\n        self,\n        epoch,\n        model,\n        data_loader,\n        optimizer,\n        lr_scheduler,\n        scaler=None,\n        cuda_enabled=False,\n        log_freq=50,\n        accum_grad_iters=1,\n    ):\n        return self._train_inner_loop(\n            epoch=epoch,\n            iters_per_epoch=lr_scheduler.iters_per_epoch,\n            model=model,\n            data_loader=data_loader,\n            optimizer=optimizer,\n            scaler=scaler,\n            lr_scheduler=lr_scheduler,\n            log_freq=log_freq,\n            cuda_enabled=cuda_enabled,\n            accum_grad_iters=accum_grad_iters,\n        )\n    \n    def test_epoch(\n        self,\n        epoch,\n        model,\n        data_loader,\n        optimizer,\n        lr_scheduler,\n        scaler=None,\n        cuda_enabled=False,\n        log_freq=50,\n        accum_grad_iters=1,\n    ):\n        return self._test_inner_loop(\n            epoch=epoch,\n            iters_per_epoch=lr_scheduler.iters_per_epoch,\n            model=model,\n            data_loader=data_loader,\n            optimizer=optimizer,\n            scaler=scaler,\n            lr_scheduler=lr_scheduler,\n            log_freq=log_freq,\n            cuda_enabled=cuda_enabled,\n            accum_grad_iters=accum_grad_iters,\n        )\n\n    def train_iters(\n        self,\n        epoch,\n        start_iters,\n        iters_per_inner_epoch,\n        model,\n        data_loader,\n        optimizer,\n        lr_scheduler,\n        scaler=None,\n        cuda_enabled=False,\n        log_freq=50,\n        accum_grad_iters=1,\n    ):\n        return self._train_inner_loop(\n            epoch=epoch,\n            start_iters=start_iters,\n            iters_per_epoch=iters_per_inner_epoch,\n            model=model,\n            data_loader=data_loader,\n            optimizer=optimizer,\n            scaler=scaler,\n            lr_scheduler=lr_scheduler,\n            log_freq=log_freq,\n            cuda_enabled=cuda_enabled,\n            accum_grad_iters=accum_grad_iters,\n        )\n    \n    def test_iters(\n        self,\n        epoch,\n        start_iters,\n        iters_per_inner_epoch,\n        model,\n        data_loader,\n        optimizer,\n        lr_scheduler,\n        scaler=None,\n        cuda_enabled=False,\n        log_freq=50,\n        accum_grad_iters=1,\n    ):\n        return self._test_inner_loop(\n            epoch=epoch,\n            start_iters=start_iters,\n            iters_per_epoch=iters_per_inner_epoch,\n            model=model,\n            data_loader=data_loader,\n            optimizer=optimizer,\n            scaler=scaler,\n            lr_scheduler=lr_scheduler,\n            log_freq=log_freq,\n            cuda_enabled=cuda_enabled,\n            accum_grad_iters=accum_grad_iters,\n        )\n\n    def _train_inner_loop(\n        self,\n        epoch,\n        iters_per_epoch,\n        model,\n        data_loader,\n        optimizer,\n        lr_scheduler,\n        scaler=None,\n        start_iters=None,\n        log_freq=50,\n        cuda_enabled=False,\n        accum_grad_iters=1,\n    ):\n        \"\"\"\n        An inner training loop compatible with both epoch-based and iter-based training.\n\n        When using epoch-based, training stops after one epoch; when using iter-based,\n        training stops after #iters_per_epoch iterations.\n        \"\"\"\n        use_amp = scaler is not None\n\n        if not hasattr(data_loader, \"__next__\"):\n            # convert to iterator if not already\n            data_loader = iter(data_loader)\n        metric_logger = MetricLogger(delimiter=\"  \")\n        metric_logger.add_meter(\"lr\", SmoothedValue(window_size=1, fmt=\"{value:.6f}\"))\n        metric_logger.add_meter(\"loss\", SmoothedValue(window_size=1, fmt=\"{value:.4f}\"))\n\n        # if iter-based runner, schedule lr based on inner epoch.\n        logging.info(\n            \"Start training epoch {}, {} iters per inner epoch.\".format(\n                epoch, iters_per_epoch\n            )\n        )\n        header = \"Train: data epoch: [{}]\".format(epoch)\n        if start_iters is None:\n            # epoch-based runner\n            inner_epoch = epoch\n        else:\n            # In iter-based runner, we schedule the learning rate based on iterations.\n            inner_epoch = start_iters // iters_per_epoch\n            header = header + \"; inner epoch [{}]\".format(inner_epoch)\n\n        for i in metric_logger.log_every(range(iters_per_epoch), log_freq, header):\n            # if using iter-based runner, we stop after iters_per_epoch iterations.\n            if i >= iters_per_epoch:\n                break\n\n            samples = next(data_loader)\n\n            samples = prepare_sample(samples, cuda_enabled=cuda_enabled)\n            samples.update(\n                {\n                    \"epoch\": inner_epoch,\n                    \"num_iters_per_epoch\": iters_per_epoch,\n                    \"iters\": i,\n                }\n            )\n\n            lr_scheduler.step(cur_epoch=inner_epoch, cur_step=i)\n\n            with torch.cuda.amp.autocast(enabled=use_amp):\n                loss = self.train_step(model=model, samples=samples)\n\n            # after_train_step()\n            if use_amp:\n                scaler.scale(loss).backward()\n            else:\n                loss.backward()\n\n            # update gradients every accum_grad_iters iterations\n            if (i + 1) % accum_grad_iters == 0:\n                if use_amp:\n                    scaler.step(optimizer)\n                    scaler.update()                     \n                else:    \n                    optimizer.step()\n                optimizer.zero_grad()\n\n            metric_logger.update(loss=loss.item())\n            metric_logger.update(lr=optimizer.param_groups[0][\"lr\"])\n                \n\n        # after train_epoch()\n        # gather the stats from all processes\n        metric_logger.synchronize_between_processes()\n        logging.info(\"Averaged stats: \" + str(metric_logger.global_avg()))\n        return {\n            k: \"{:.3f}\".format(meter.global_avg)\n            for k, meter in metric_logger.meters.items()\n        }\n    \n    def _test_inner_loop(\n        self,\n        epoch,\n        iters_per_epoch,\n        model,\n        data_loader,\n        optimizer,\n        lr_scheduler,\n        scaler=None,\n        start_iters=None,\n        log_freq=50,\n        cuda_enabled=False,\n        accum_grad_iters=1,\n    ):\n        \"\"\"\n        An inner testing loop .\n        \"\"\"\n        use_amp = scaler is not None\n\n        if not hasattr(data_loader, \"__next__\"):\n            # convert to iterator if not already\n            data_loader = iter(data_loader)\n\n        for I in range(iters_per_epoch) :\n            samples = next(data_loader)\n\n            samples = prepare_sample(samples, cuda_enabled=cuda_enabled)\n            \n            with torch.cuda.amp.autocast(enabled=use_amp):\n                output_text = self.test_step(model=model, samples=samples)\n\n            output_text = output_text.replace(\"<Img>\", \"\")\n            output_text = output_text.replace(\"Impression:\", \"\")\n            output_text = output_text.replace(\"\\n\", \"\")\n            output_text = output_text.replace(\"#\", \"\")\n            output_text = output_text.replace(\"___\", \"\")\n\n            fields=[str(samples['image_id'][0].tolist()), output_text, samples['caption'][0]]\n            with open(r'/mnt/lustre/huangzhongzhen/pretrain_our/ImprMiniGPT/ours_medclip/minigpt4_stage3_all_v2.2(mimic_chatgpt)_radiology_finetune/20230509154/result/vanilla_minigpt.csv', 'a') as f:\n                writer = csv.writer(f)\n                writer.writerow(fields)\n\n            print('{}/{}'.format(I,iters_per_epoch))\n\n\n\n    @staticmethod\n    def save_result(result, result_dir, filename, remove_duplicate=\"\"):\n        import json\n\n        result_file = os.path.join(\n            result_dir, \"%s_rank%d.json\" % (filename, get_rank())\n        )\n        final_result_file = os.path.join(result_dir, \"%s.json\" % filename)\n\n        json.dump(result, open(result_file, \"w\"))\n\n        if is_dist_avail_and_initialized():\n            dist.barrier()\n\n        if is_main_process():\n            logging.warning(\"rank %d starts merging results.\" % get_rank())\n            # combine results from all processes\n            result = []\n\n            for rank in range(get_world_size()):\n                result_file = os.path.join(\n                    result_dir, \"%s_rank%d.json\" % (filename, rank)\n                )\n                res = json.load(open(result_file, \"r\"))\n                result += res\n\n            if remove_duplicate:\n                result_new = []\n                id_list = []\n                for res in result:\n                    if res[remove_duplicate] not in id_list:\n                        id_list.append(res[remove_duplicate])\n                        result_new.append(res)\n                result = result_new\n\n            json.dump(result, open(final_result_file, \"w\"))\n            print(\"result file saved to %s\" % final_result_file)\n\n        return final_result_file\n"
  },
  {
    "path": "xraypulse/tasks/image_text_pretrain.py",
    "content": "\"\"\"\n Copyright (c) 2022, salesforce.com, inc.\n All rights reserved.\n SPDX-License-Identifier: BSD-3-Clause\n For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause\n\"\"\"\n\nfrom xraypulse.common.registry import registry\nfrom xraypulse.tasks.base_task import BaseTask\n\n\n@registry.register_task(\"image_text_pretrain\")\nclass ImageTextPretrainTask(BaseTask):\n    def __init__(self):\n        super().__init__()\n\n    def evaluation(self, model, data_loader, cuda_enabled=True):\n        pass\n"
  }
]