nearest-neighbors.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a5b478bb-e54f-41de-bb4e-0b75dee8f087",
   "metadata": {},
   "source": [
    "# Goal\n",
    "\n",
    "Use https://blog.paperspace.com/sentence-embeddings-pytorch-lightning/ with as few lines of code/text as possible :) \n",
    "\n",
    "Budget: $50 in Anthropic API credits!\n",
    "\n",
    "Constraints: 2 hours of active \"coding\" or \"prompting\" time on the day of the talk (https://lu.ma/LAIxNYC).\n",
    "\n",
    "Data from:\n",
    "\n",
    "https://colab.research.google.com/github/onefact/electronic-health-records-analysis/blob/main/notebooks/loading_physionet_mimic_iii_data.ipynb\n",
    "\n",
    "https://colab.research.google.com/github/onefact/electronic-health-records-analysis/blob/main/notebooks/loading_physionet_mimic_iv_data.ipynb \n",
    "\n",
    "https://colab.research.google.com/github/onefact/electronic-health-records-analysis/blob/main/notebooks/loading_physionet_mimic_iv_clinical_notes.ipynb"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40fbc828-7a02-4e34-9357-bba448184e60",
   "metadata": {},
   "source": [
    "# Setup and prerequisites\n",
    "\n",
    "Create a new text file, populated from `requirements.in` file in https://colab.research.google.com/github/jaanli/language-model-notebooks/blob/main/notebooks/getting-started.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a1a7ee4e-924f-4b79-8cba-d63ea5de0392",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "\n",
    "!pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2f15c84-3e5e-417d-ac17-950dd72cda27",
   "metadata": {},
   "source": [
    "# Load packages and extensions needed to get help from large language models like Claude and ChatGPT :)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a2a01d73-46df-480e-ad81-814c82aa3a8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext jupyter_ai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "20290aaf-52f6-4538-8223-b635e1ef7e07",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext dotenv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ba9ebb85-3ecd-457d-8a90-d48e9299d298",
   "metadata": {},
   "outputs": [],
   "source": [
    "%dotenv "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5937900d-c73c-460c-8040-b1f472deffe1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/me/projects/loving-the-baseline/.venv/lib/python3.10/site-packages/sql/traits.py:20: FutureWarning: named_parameters: boolean values are now deprecated. Value True will be treated as \"enabled\". \n",
      "Please use a valid option: \"warn\", \"enabled\", or \"disabled\". \n",
      "For more information, see the docs: https://jupysql.ploomber.io/en/latest/api/configuration.html#named-parameters\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "# Load duckdb, which lets us efficiently load large files\n",
    "import duckdb\n",
    "\n",
    "# Load pandas, which lets us manipulate dataframes\n",
    "import pandas as pd\n",
    "\n",
    "# Import jupysql Jupyter extension to create SQL cells\n",
    "%load_ext sql\n",
    "\n",
    "# Set configrations on jupysql to directly output data to Pandas and to simplify the output that is printed to the notebook.\n",
    "%config SqlMagic.autopandas = True\n",
    "\n",
    "%config SqlMagic.feedback = False\n",
    "%config SqlMagic.displaycon = False\n",
    "\n",
    "# Allow named parameters (python variables) in SQL cells\n",
    "%config SqlMagic.named_parameters=True\n",
    "\n",
    "# Connect jupysql to DuckDB using a SQLAlchemy-style connection string. Either connect to an in memory DuckDB, or a file backed db.\n",
    "%sql duckdb:///:memory:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "daa9e6e0-656f-48e3-a82d-5d1ad4a833aa",
   "metadata": {},
   "source": [
    "# Look at input data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7f1a7326-2e2d-4d31-9ceb-a7c75d42fc99",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/Users/me/data/mimic/mimic-iv-note-deidentified-free-text-clinical-notes-2.2/discharge.parquet\n"
     ]
    }
   ],
   "source": [
    "!ls ~/data/mimic/mimic-iv-note-deidentified-free-text-clinical-notes-2.2/discharge.parquet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "31140ce0-5012-40e5-8b14-530c9f11b8bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture discharge\n",
    "%%bash\n",
    "for file in /Users/me/data/mimic/mimic-iv-note-deidentified-free-text-clinical-notes-2.2/discharge.parquet; do \n",
    "    echo \"$file\"\n",
    "    echo \"----------\"\n",
    "    duckdb -c \"DESCRIBE SELECT * FROM parquet_scan('$file')\"\n",
    "    echo \"==========\"\n",
    "done"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "078b8a2d-7c03-47ee-b1ce-761600fa4f56",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/Users/me/data/mimic/mimic-iv-note-deidentified-free-text-clinical-notes-2.2/discharge.parquet\n",
      "----------\n",
      "┌─────────────┬─────────────┬─────────┬─────────┬─────────┬─────────┐\n",
      "│ column_name │ column_type │  null   │   key   │ default │  extra  │\n",
      "│   varchar   │   varchar   │ varchar │ varchar │ varchar │ varchar │\n",
      "├─────────────┼─────────────┼─────────┼─────────┼─────────┼─────────┤\n",
      "│ note_id     │ VARCHAR     │ YES     │         │         │         │\n",
      "│ subject_id  │ INTEGER     │ YES     │         │         │         │\n",
      "│ hadm_id     │ INTEGER     │ YES     │         │         │         │\n",
      "│ note_type   │ VARCHAR     │ YES     │         │         │         │\n",
      "│ note_seq    │ INTEGER     │ YES     │         │         │         │\n",
      "│ charttime   │ TIMESTAMP   │ YES     │         │         │         │\n",
      "│ storetime   │ TIMESTAMP   │ YES     │         │         │         │\n",
      "│ text        │ VARCHAR     │ YES     │         │         │         │\n",
      "└─────────────┴─────────────┴─────────┴─────────┴─────────┴─────────┘\n",
      "==========\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(discharge)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "050c3f76-36bf-449b-be9d-d7e977c8ac05",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/Users/me/data/mimic/mimic-iv-2.2/d_icd_diagnoses.parquet\n"
     ]
    }
   ],
   "source": [
    "!ls ~/data/mimic/mimic-iv-2.2/d_icd_diagnoses.parquet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "4a392fc6-2927-48d0-a9cd-a7dcd3a546fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture d_icd_diagnoses\n",
    "%%bash\n",
    "for file in /Users/me/data/mimic/mimic-iv-2.2/d_icd_diagnoses.parquet; do \n",
    "    echo \"$file\"\n",
    "    echo \"----------\"\n",
    "    duckdb -c \"DESCRIBE SELECT * FROM parquet_scan('$file')\"\n",
    "    echo \"==========\"\n",
    "done"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "5285f6f8-8108-4c96-b9d4-68e218b90402",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/Users/me/data/mimic/mimic-iv-2.2/d_icd_diagnoses.parquet\n",
      "----------\n",
      "┌─────────────┬─────────────┬─────────┬─────────┬─────────┬─────────┐\n",
      "│ column_name │ column_type │  null   │   key   │ default │  extra  │\n",
      "│   varchar   │   varchar   │ varchar │ varchar │ varchar │ varchar │\n",
      "├─────────────┼─────────────┼─────────┼─────────┼─────────┼─────────┤\n",
      "│ icd_code    │ VARCHAR     │ YES     │         │         │         │\n",
      "│ icd_version │ VARCHAR     │ YES     │         │         │         │\n",
      "│ long_title  │ VARCHAR     │ YES     │         │         │         │\n",
      "└─────────────┴─────────────┴─────────┴─────────┴─────────┴─────────┘\n",
      "==========\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(d_icd_diagnoses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "e64621e9-3dc2-478e-b38a-b60b18e08d0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture d_icd_procedures\n",
    "%%bash\n",
    "for file in /Users/me/data/mimic/mimic-iv-2.2/d_icd_procedures.parquet; do \n",
    "    echo \"$file\"\n",
    "    echo \"----------\"\n",
    "    duckdb -c \"DESCRIBE SELECT * FROM parquet_scan('$file')\"\n",
    "    echo \"==========\"\n",
    "done"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "5577c903-cc65-4dc3-8178-07d6a5dd3845",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/Users/me/data/mimic/mimic-iv-2.2/d_icd_procedures.parquet\n",
      "----------\n",
      "┌─────────────┬─────────────┬─────────┬─────────┬─────────┬─────────┐\n",
      "│ column_name │ column_type │  null   │   key   │ default │  extra  │\n",
      "│   varchar   │   varchar   │ varchar │ varchar │ varchar │ varchar │\n",
      "├─────────────┼─────────────┼─────────┼─────────┼─────────┼─────────┤\n",
      "│ icd_code    │ VARCHAR     │ YES     │         │         │         │\n",
      "│ icd_version │ BIGINT      │ YES     │         │         │         │\n",
      "│ long_title  │ VARCHAR     │ YES     │         │         │         │\n",
      "└─────────────┴─────────────┴─────────┴─────────┴─────────┴─────────┘\n",
      "==========\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(d_icd_procedures)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "5cef2063-f72e-427b-8e5a-5a5f57db2d93",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture diagnoses_icd\n",
    "%%bash\n",
    "for file in /Users/me/data/mimic/mimic-iv-2.2/diagnoses_icd.parquet; do \n",
    "    echo \"$file\"\n",
    "    echo \"----------\"\n",
    "    duckdb -c \"DESCRIBE SELECT * FROM parquet_scan('$file')\"\n",
    "    echo \"==========\"\n",
    "done"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "a6fb3b03-8afd-492a-aa92-96d6c1c491ba",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/Users/me/data/mimic/mimic-iv-2.2/diagnoses_icd.parquet\n",
      "----------\n",
      "┌─────────────┬─────────────┬─────────┬─────────┬─────────┬─────────┐\n",
      "│ column_name │ column_type │  null   │   key   │ default │  extra  │\n",
      "│   varchar   │   varchar   │ varchar │ varchar │ varchar │ varchar │\n",
      "├─────────────┼─────────────┼─────────┼─────────┼─────────┼─────────┤\n",
      "│ subject_id  │ INTEGER     │ YES     │         │         │         │\n",
      "│ hadm_id     │ INTEGER     │ YES     │         │         │         │\n",
      "│ seq_num     │ INTEGER     │ YES     │         │         │         │\n",
      "│ icd_code    │ VARCHAR     │ YES     │         │         │         │\n",
      "│ icd_version │ INTEGER     │ YES     │         │         │         │\n",
      "└─────────────┴─────────────┴─────────┴─────────┴─────────┴─────────┘\n",
      "==========\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(diagnoses_icd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "c5b888e7-7222-4e23-b581-3ffddd160ccb",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture procedures_icd\n",
    "%%bash\n",
    "for file in /Users/me/data/mimic/mimic-iv-2.2/procedures_icd.parquet; do \n",
    "    echo \"$file\"\n",
    "    echo \"----------\"\n",
    "    duckdb -c \"DESCRIBE SELECT * FROM parquet_scan('$file')\"\n",
    "    echo \"==========\"\n",
    "done"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "960a0542-9e5e-4e93-aa1e-5d7322c1b731",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/Users/me/data/mimic/mimic-iv-2.2/procedures_icd.parquet\n",
      "----------\n",
      "┌─────────────┬─────────────┬─────────┬─────────┬─────────┬─────────┐\n",
      "│ column_name │ column_type │  null   │   key   │ default │  extra  │\n",
      "│   varchar   │   varchar   │ varchar │ varchar │ varchar │ varchar │\n",
      "├─────────────┼─────────────┼─────────┼─────────┼─────────┼─────────┤\n",
      "│ subject_id  │ BIGINT      │ YES     │         │         │         │\n",
      "│ hadm_id     │ BIGINT      │ YES     │         │         │         │\n",
      "│ seq_num     │ BIGINT      │ YES     │         │         │         │\n",
      "│ chartdate   │ DATE        │ YES     │         │         │         │\n",
      "│ icd_code    │ VARCHAR     │ YES     │         │         │         │\n",
      "│ icd_version │ BIGINT      │ YES     │         │         │         │\n",
      "└─────────────┴─────────────┴─────────┴─────────┴─────────┴─────────┘\n",
      "==========\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(procedures_icd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "b9b5969a-4eea-4733-9fdd-d0cb9fc20d1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture icu_icustays\n",
    "%%bash\n",
    "for file in /Users/me/data/mimic/mimic-iv-2.2/icu_icustays.parquet; do \n",
    "    echo \"$file\"\n",
    "    echo \"----------\"\n",
    "    duckdb -c \"DESCRIBE SELECT * FROM parquet_scan('$file')\"\n",
    "    echo \"==========\"\n",
    "done"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "e74d2e11-febf-4808-8564-9a9609ad81d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/Users/me/data/mimic/mimic-iv-2.2/icu_icustays.parquet\n",
      "----------\n",
      "┌────────────────┬─────────────┬─────────┬─────────┬─────────┬─────────┐\n",
      "│  column_name   │ column_type │  null   │   key   │ default │  extra  │\n",
      "│    varchar     │   varchar   │ varchar │ varchar │ varchar │ varchar │\n",
      "├────────────────┼─────────────┼─────────┼─────────┼─────────┼─────────┤\n",
      "│ subject_id     │ BIGINT      │ YES     │         │         │         │\n",
      "│ hadm_id        │ BIGINT      │ YES     │         │         │         │\n",
      "│ stay_id        │ BIGINT      │ YES     │         │         │         │\n",
      "│ first_careunit │ VARCHAR     │ YES     │         │         │         │\n",
      "│ last_careunit  │ VARCHAR     │ YES     │         │         │         │\n",
      "│ intime         │ TIMESTAMP   │ YES     │         │         │         │\n",
      "│ outtime        │ TIMESTAMP   │ YES     │         │         │         │\n",
      "│ los            │ DOUBLE      │ YES     │         │         │         │\n",
      "└────────────────┴─────────────┴─────────┴─────────┴─────────┴─────────┘\n",
      "==========\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(icu_icustays)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1557a21-6a1a-4fd9-80d2-c6834969d130",
   "metadata": {},
   "source": [
    "# Include documentation as context for newer libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "54a5a931-dec7-4f69-8ffb-949d381add2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture lightning_docs\n",
    "\n",
    "!curl -s \"https://blog.paperspace.com/sentence-embeddings-pytorch-lightning/\" | sed -e 's/<[^>]*>//g; /^$/d' | tr -s '\\n'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "abca56a4-19fd-4cb0-90b8-824784b6de76",
   "metadata": {},
   "outputs": [],
   "source": [
    "example = \"\"\"\n",
    "import torch\n",
    "\n",
    "import flash\n",
    "from flash.text import TextClassificationData, TextEmbedder\n",
    "from sentence_transformers import util\n",
    "predict_data=[\"I like to watch tv\",\"Watching Television is my favorite time pass\",\"The cat was running after the butterfly\",\"It is so windy today\"]\n",
    "# Wrapping the prediction data inside a datamodule\n",
    "datamodule = TextClassificationData.from_lists(\n",
    "    predict_data=predict_data,\n",
    "    batch_size=8,\n",
    ")\n",
    "\n",
    "# We are loading a pre-trained SentenceEmbedder\n",
    "model = TextEmbedder(backbone=\"sentence-transformers/all-MiniLM-L6-v2\")\n",
    "\n",
    "trainer = flash.Trainer(gpus=1)\n",
    "\n",
    "#Since this task is tackled unsupervised, the predict method generates sentence embeddings using the prediction input\n",
    "embeddings = trainer.predict(model, datamodule=datamodule)\n",
    "\n",
    "for i in range(0,len(predict_data),2):\n",
    "  embed1,embed2=embeddings[0][i],embeddings[0][i+1]\n",
    "  # we are using cosine similarity to compute the similarity score\n",
    "  cosine_scores = util.cos_sim(embed1, embed2)\n",
    "  if cosine_scores>=0.5:\n",
    "      label=\"Similar\"\n",
    "  else:\n",
    "      label=\"Not Similar\"\n",
    "  print(\"sentence 1:{} | sentence 2:{}| prediction: {}\".format(predict_data[i],predict_data[i+1],label))\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "642396fc-b16b-4ff8-8b39-dce7a7f3aa40",
   "metadata": {},
   "source": [
    "# Write prompts to get help from large language models for lightning stack"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "4c8bb762-dfb5-4e2d-a1c9-f7bc775dac5c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "AI generated code inserted below &#11015;&#65039;"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "execution_count": 37,
     "metadata": {
      "text/html": {
       "jupyter_ai": {
        "model_id": "claude-3-opus-20240229",
        "provider_id": "anthropic-chat"
       }
      }
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%ai anthropic-chat:claude-3-opus-20240229 --format code\n",
    "\n",
    "Given the schema in these files: \n",
    "\n",
    "{icu_icustays}\n",
    "{procedures_icd}\n",
    "{diagnoses_icd}\n",
    "{d_icd_procedures}\n",
    "{d_icd_diagnoses}\n",
    "{discharge}\n",
    "\n",
    "And this example code for using Lightning with sentence transformers (ignore the fine-tuning parts, but remember the key abstractions, methods, etc!), from https://blog.paperspace.com/sentence-embeddings-pytorch-lightning/:\n",
    "\n",
    "```\n",
    "```\n",
    "\n",
    "\n",
    "And this example duckdb dialect SQL query that creates a histogram of length of stays in addition to their procedures and diagnoses (from https://colab.research.google.com/github/onefact/electronic-health-records-analysis/blob/main/notebooks/loading_physionet_mimic_iv_clinical_notes.ipynb): \n",
    "\n",
    "```\n",
    "SELECT \n",
    "    i.hadm_id,\n",
    "    ROUND(i.los) AS los_days,\n",
    "    dp.icd_code AS procedure_icd_code,\n",
    "    dp.long_title AS procedure_description,\n",
    "    dd.icd_code AS diagnosis_icd_code,  \n",
    "    dd.long_title AS diagnosis_description\n",
    "FROM icu_icustays i\n",
    "LEFT JOIN procedures_icd p ON i.hadm_id = p.hadm_id\n",
    "LEFT JOIN d_icd_procedures dp ON p.icd_code = dp.icd_code\n",
    "LEFT JOIN diagnoses_icd d ON i.hadm_id = d.hadm_id\n",
    "LEFT JOIN d_icd_diagnoses dd ON d.icd_code = dd.icd_code\n",
    "```\n",
    "\n",
    "Proceed to create simple PyTorch Lightning code using the `jaxtyping` and `beartype` libraries for runtime static type checking. \n",
    "\n",
    "Ensure that critical methods are decorated using the `@beartype` decorator, and type hints are given using the jaxtyping syntax (https://kidger.site/thoughts/jaxtyping/).\n",
    "\n",
    "Proceed step-by-step as an elite site reliability/devops/L20 principal warez engineer at google, returning as few tokens as possible, to debug this python & lightning code. Give the complete corrected code ready to be run in a Jupyter Notebook cell!\n",
    "\n",
    "The output should be the following: \n",
    "\n",
    "* for every `procedure_description` and `diagnosis_description` in the data (`icd_code` in most tables/parquet files)\n",
    "* segment the `text` field in the `discharge` table in the `hadm_id` (hospital admission identifer) that uniquely identifies an intensive care unit admission\n",
    "* do this using the above example methods like sentence transformers, faiss, etc\n",
    "* remember to use runtime static type checking with beartyping and jaxtyping throughout\n",
    "* create code that takes a random sample of sentences from clinical notes, and computes the cosine similarity from every sentence to every `procedure_description` or `diagnosis_description` in the data\n",
    "* then use the `torchmetrics` library to compute the accuracy / recall / standard metrics for this task and report the performance \n",
    "* do this for a small number of hospital admission identifiers (selected at random) and a small number of sentences (also selected at random) for demonstration purposes\n",
    "\n",
    "You can do it! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8837dd56-d5b7-4053-8958-4484e4f50da8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "from typing import List, Tuple, Dict\n",
    "\n",
    "import duckdb\n",
    "import pandas as pd\n",
    "import torch\n",
    "from beartype import beartype\n",
    "from jaxtyping import Float, Int, Array, PyTree\n",
    "from lightning.pytorch import LightningModule\n",
    "from sentence_transformers import SentenceTransformer\n",
    "from sklearn.metrics.pairwise import cosine_similarity\n",
    "from torchmetrics import MetricCollection, Accuracy, Precision, Recall, F1Score\n",
    "\n",
    "\n",
    "@beartype\n",
    "def load_data(hadm_ids: List[Int[32]]) -> Tuple[pd.DataFrame, pd.DataFrame]:\n",
    "    query = \"\"\"\n",
    "        SELECT \n",
    "            i.hadm_id,\n",
    "            ROUND(i.los) AS los_days,\n",
    "            dp.icd_code AS procedure_icd_code,\n",
    "            dp.long_title AS procedure_description,\n",
    "            dd.icd_code AS diagnosis_icd_code,  \n",
    "            dd.long_title AS diagnosis_description\n",
    "        FROM icu_icustays i\n",
    "        LEFT JOIN procedures_icd p ON i.hadm_id = p.hadm_id\n",
    "        LEFT JOIN d_icd_procedures dp ON p.icd_code = dp.icd_code\n",
    "        LEFT JOIN diagnoses_icd d ON i.hadm_id = d.hadm_id\n",
    "        LEFT JOIN d_icd_diagnoses dd ON d.icd_code = dd.icd_code\n",
    "        WHERE i.hadm_id IN ({})\n",
    "    \"\"\".format(', '.join(map(str, hadm_ids)))\n",
    "\n",
    "    con = duckdb.connect()\n",
    "    df = con.execute(query).df()\n",
    "\n",
    "    query = \"\"\"\n",
    "        SELECT \n",
    "            hadm_id,\n",
    "            text\n",
    "        FROM discharge\n",
    "        WHERE hadm_id IN ({})\n",
    "    \"\"\".format(', '.join(map(str, hadm_ids)))\n",
    "\n",
    "    notes_df = con.execute(query).df()\n",
    "\n",
    "    return df, notes_df\n",
    "\n",
    "\n",
    "@beartype\n",
    "def get_embeddings(model: SentenceTransformer, sentences: List[str]) -> Array[Float[64]]:\n",
    "    embeddings = model.encode(sentences)\n",
    "    return embeddings\n",
    "\n",
    "\n",
    "class SimilarityModel(LightningModule):\n",
    "    def __init__(self, encoder_model: str):\n",
    "        super().__init__()\n",
    "        self.encoder = SentenceTransformer(encoder_model)\n",
    "        self.metrics = MetricCollection([Accuracy(), Precision(), Recall(), F1Score()])\n",
    "\n",
    "    @beartype\n",
    "    def forward(self, sentences: List[str], labels: List[str]) -> PyTree:\n",
    "        sentence_embeddings = get_embeddings(self.encoder, sentences)\n",
    "        label_embeddings = get_embeddings(self.encoder, labels)\n",
    "\n",
    "        similarities = cosine_similarity(sentence_embeddings, label_embeddings)\n",
    "        predicted_labels = [labels[i] for i in similarities.argmax(axis=1)]\n",
    "\n",
    "        self.metrics.update(predicted_labels, labels)\n",
    "\n",
    "        return {\"loss\": torch.tensor(0), \"predicted_labels\": predicted_labels}\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        sentences, labels = batch\n",
    "        return self(sentences, labels)\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        sentences, labels = batch\n",
    "        return self(sentences, labels)\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        sentences, labels = batch\n",
    "        return self(sentences, labels)\n",
    "\n",
    "    def training_epoch_end(self, outputs):\n",
    "        self.log_dict(self.metrics.compute())\n",
    "        self.metrics.reset()\n",
    "\n",
    "    def validation_epoch_end(self, outputs):\n",
    "        self.log_dict(self.metrics.compute())\n",
    "        self.metrics.reset()\n",
    "\n",
    "    def test_epoch_end(self, outputs):\n",
    "        self.log_dict(self.metrics.compute())\n",
    "        self.metrics.reset()\n",
    "\n",
    "\n",
    "@beartype\n",
    "def main():\n",
    "    random.seed(42)\n",
    "    hadm_ids = random.sample(range(100000, 200000), 10)  # Random hospital admission IDs\n",
    "    df, notes_df = load_data(hadm_ids)\n",
    "\n",
    "    labels = list(df[\"procedure_description\"].unique()) + list(df[\"diagnosis_description\"].unique())\n",
    "\n",
    "    sentences ="
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}