upload_to_hub_cxrmate_ed.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/apps/python/3.12.0/lib/python3.12/site-packages/numpy/core/getlimits.py:549: UserWarning: The value of the smallest subnormal for <class 'numpy.float64'> type is zero.\n",
      "  setattr(self, word, getattr(machar, word).flat[0])\n",
      "/apps/python/3.12.0/lib/python3.12/site-packages/numpy/core/getlimits.py:89: UserWarning: The value of the smallest subnormal for <class 'numpy.float64'> type is zero.\n",
      "  return self._float_to_str(self.smallest_subnormal)\n",
      "/apps/python/3.12.0/lib/python3.12/site-packages/numpy/core/getlimits.py:549: UserWarning: The value of the smallest subnormal for <class 'numpy.float32'> type is zero.\n",
      "  setattr(self, word, getattr(machar, word).flat[0])\n",
      "/apps/python/3.12.0/lib/python3.12/site-packages/numpy/core/getlimits.py:89: UserWarning: The value of the smallest subnormal for <class 'numpy.float32'> type is zero.\n",
      "  return self._float_to_str(self.smallest_subnormal)\n"
     ]
    }
   ],
   "source": [
    "from cxrmate_ed.configuration_cxrmate_ed import CXRMateEDConfig\n",
    "from cxrmate_ed.modelling_cxrmate_ed import CXRMateEDModel\n",
    "from cxrmate_ed.lightning_module import SCSTSectionsWeightedCXRBERTBERTScoreARNReward\n",
    "import torch\n",
    "import transformers\n",
    "import warnings\n",
    "import yaml\n",
    "from pathlib import Path\n",
    "import os\n",
    "from huggingface_hub import HfApi\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('4.46.3', '2.1.1')"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "transformers.__version__, torch.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hub checkpoint name:\n",
    "hub_ckpt_name = 'aehrc/cxrmate-ed'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Paths:\n",
    "ckpt_zoo_dir = '/datasets/work/hb-mlaifsp-mm/work/checkpoints'\n",
    "database_dir = '/scratch3/nic261/database/cxrmate_ed'\n",
    "ckpt_path = '/datasets/work/hb-mlaifsp-mm/work/experiments/cxrmate_ed/rl/003_scst_sections_weighted_cxrbert_bertscore_ngram/trial_0/epoch=31-step=76416-val_findings_bertscore_f1=0.459060.ckpt'\n",
    "args_path = '/datasets/work/hb-mlaifsp-mm/work/experiments/cxrmate_ed/rl/003_scst_sections_weighted_cxrbert_bertscore_ngram/trial_0/arguments/session_6.yaml'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "transformers.AutoConfig.register('cxrmate-ed', CXRMateEDConfig)\n",
    "transformers.AutoModelForCausalLM.register(CXRMateEDConfig, CXRMateEDModel)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "CXRMateEDConfig.register_for_auto_class()\n",
    "CXRMateEDModel.register_for_auto_class('AutoModelForCausalLM')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Description, Special token, Index\n",
      "bos_token, [BOS], 1\n",
      "eos_token, [EOS], 2\n",
      "unk_token, [UNK], 0\n",
      "sep_token, [SEP], 3\n",
      "pad_token, [PAD], 4\n",
      "cls_token, [BOS], 1\n",
      "mask_token, [MASK], 5\n"
     ]
    }
   ],
   "source": [
    "# HF Model and tokenizer:\n",
    "args = yaml.safe_load(Path(args_path).read_text())\n",
    "lightning_module = SCSTSectionsWeightedCXRBERTBERTScoreARNReward(warm_start_modules=False, **args)\n",
    "model = lightning_module.model\n",
    "tokenizer = lightning_module.tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/apps/pytorch/2.1.1-py312-cu122-mpi/lib/python3.12/site-packages/torch/_utils.py:831: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  return self.fget.__get__(instance, owner)()\n"
     ]
    }
   ],
   "source": [
    "# Rename state dict keys:\n",
    "state_dict = torch.load(ckpt_path, map_location='cpu', weights_only=True)['state_dict']\n",
    "\n",
    "for key in list(state_dict.keys()):\n",
    "    if 'encoder_decoder.' in key:\n",
    "        state_dict[key.replace('encoder_decoder.', '')] = state_dict.pop(key)\n",
    "\n",
    "for key in list(state_dict.keys()):\n",
    "    if 'encoder.uniformer.' in key:\n",
    "        state_dict[key.replace('encoder.uniformer.', 'image_encoder.encoder.uniformer.')] = state_dict.pop(key)\n",
    "\n",
    "for key in list(state_dict.keys()):\n",
    "    if 'decoder.model.' in key:\n",
    "        state_dict[key.replace('decoder.model.', 'language_model.model.')] = state_dict.pop(key)\n",
    "\n",
    "for key in list(state_dict.keys()):\n",
    "    if 'decoder.lm_head.weight' in key:\n",
    "        state_dict[key.replace('decoder.lm_head.weight', 'language_model.lm_head.weight')] = state_dict.pop(key)\n",
    "\n",
    "for key in list(state_dict.keys()):\n",
    "    if 'encoder.projection_head.' in key:\n",
    "        state_dict[key.replace('encoder.projection_head.', 'image_encoder.adapter.')] = state_dict.pop(key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.load_state_dict(state_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-12-13 11:47:45,678] [INFO] [real_accelerator.py:191:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/nic261/.local/lib/python3.12/site-packages/pydantic/_internal/_fields.py:160: UserWarning: Field \"model_server_url\" has conflict with protected namespace \"model_\".\n",
      "\n",
      "You may be able to resolve this warning by setting `model_config['protected_namespaces'] = ()`.\n",
      "  warnings.warn(\n",
      "/home/nic261/.local/lib/python3.12/site-packages/pydantic/_internal/_config.py:334: UserWarning: Valid config keys have changed in V2:\n",
      "* 'schema_extra' has been renamed to 'json_schema_extra'\n",
      "  warnings.warn(message, UserWarning)\n"
     ]
    }
   ],
   "source": [
    "# Save model:\n",
    "save_path = '/scratch3/nic261/checkpoints/cxrmate_ed'\n",
    "model.save_pretrained(save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1510586cf9c44be4ada0f0dfacc3391d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "README.md:   0%|          | 0.00/5.26k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/aehrc/cxrmate-ed/commit/d92f6a0fbdd4a81c9b8727d2da625ba585d4056f', commit_message='Upload tokenizer', commit_description='', oid='d92f6a0fbdd4a81c9b8727d2da625ba585d4056f', pr_url=None, pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Push to hub:\n",
    "model.push_to_hub(hub_ckpt_name)\n",
    "tokenizer.push_to_hub(hub_ckpt_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "api = HfApi()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/aehrc/cxrmate-ed/commit/88235a43d3f0f45599c60ce99770c5b70f475ebe', commit_message='Upload mimic_iv_ed_mimic_cxr_jpg_test_study_ids.json with huggingface_hub', commit_description='', oid='88235a43d3f0f45599c60ce99770c5b70f475ebe', pr_url=None, pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "api.upload_file(\n",
    "    path_or_fileobj='cxrmate_ed/lookup_tables.json',\n",
    "    path_in_repo='lookup_tables.json',\n",
    "    repo_id=hub_ckpt_name,\n",
    ")\n",
    "api.upload_file(\n",
    "    path_or_fileobj='cxrmate_ed/tables.json',\n",
    "    path_in_repo='tables.json',\n",
    "    repo_id=hub_ckpt_name,\n",
    ")\n",
    "api.upload_file(\n",
    "    path_or_fileobj='cxrmate_ed/token_type_ids.json',\n",
    "    path_in_repo='token_type_ids.json',\n",
    "    repo_id=hub_ckpt_name,\n",
    ")\n",
    "api.upload_file(\n",
    "    path_or_fileobj='cxrmate_ed/mimic_cxr_jpg_train_study_ids.json',\n",
    "    path_in_repo='mimic_cxr_jpg_train_study_ids.json',\n",
    "    repo_id=hub_ckpt_name,\n",
    ")\n",
    "api.upload_file(\n",
    "    path_or_fileobj='cxrmate_ed/mimic_cxr_jpg_validate_study_ids.json',\n",
    "    path_in_repo='mimic_cxr_jpg_validate_study_ids.json',\n",
    "    repo_id=hub_ckpt_name,\n",
    ")\n",
    "api.upload_file(\n",
    "    path_or_fileobj='cxrmate_ed/mimic_cxr_jpg_test_study_ids.json',\n",
    "    path_in_repo='mimic_cxr_jpg_test_study_ids.json',\n",
    "    repo_id=hub_ckpt_name,\n",
    ")\n",
    "api.upload_file(\n",
    "    path_or_fileobj='cxrmate_ed/mimic_iv_ed_mimic_cxr_jpg_train_study_ids.json',\n",
    "    path_in_repo='mimic_iv_ed_mimic_cxr_jpg_train_study_ids.json',\n",
    "    repo_id=hub_ckpt_name,\n",
    ")\n",
    "api.upload_file(\n",
    "    path_or_fileobj='cxrmate_ed/mimic_iv_ed_mimic_cxr_jpg_validate_study_ids.json',\n",
    "    path_in_repo='mimic_iv_ed_mimic_cxr_jpg_validate_study_ids.json',\n",
    "    repo_id=hub_ckpt_name,\n",
    ")\n",
    "api.upload_file(\n",
    "    path_or_fileobj='cxrmate_ed/mimic_iv_ed_mimic_cxr_jpg_test_study_ids.json',\n",
    "    path_in_repo='mimic_iv_ed_mimic_cxr_jpg_test_study_ids.json',\n",
    "    repo_id=hub_ckpt_name,\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cxrmate_ed",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.0"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}