steps/forward_pass/Forward Pass.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "673c70c6",
   "metadata": {},
   "source": [
    "## STEPS"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f4a0a77",
   "metadata": {},
   "source": [
    "* We go through the 4 steps that are required to de-identify a dataset (i.e run the forward pass on this dataset using a trained model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0704cda3",
   "metadata": {},
   "source": [
    "## STEP 0: LIBRARIES"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea299f28",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2113a64e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import HfArgumentParser, TrainingArguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff0c5b5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from robust_deid.ner_datasets import DatasetCreator\n",
    "from robust_deid.sequence_tagging import SequenceTagger\n",
    "from robust_deid.sequence_tagging.arguments import (\n",
    "    ModelArguments,\n",
    "    DataTrainingArguments,\n",
    "    EvaluationArguments,\n",
    ")\n",
    "from robust_deid.deid import TextDeid"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15b183ce",
   "metadata": {},
   "source": [
    "## STEP 1: INITIALIZE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0ec20f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the path where the dataset is located (input_file).\n",
    "# Input dataset\n",
    "input_file = '../../data/notes/notes.jsonl'\n",
    "# Initialize the location where we will store the sentencized and tokenized dataset (ner_dataset_file)\n",
    "ner_dataset_file = '../../data/ner_datasets/test.jsonl'\n",
    "# Initialize the location where we will store the model predictions (predictions_file)\n",
    "# Verify this file location - Ensure it's the same location that you will pass in the json file\n",
    "# to the sequence tagger model. i.e. output_predictions_file in the json file should have the same\n",
    "# value as below\n",
    "predictions_file = '../../data/predictions/predictions.jsonl'\n",
    "# Initialize the file that will contain the original note text and the de-identified note text\n",
    "deid_file = '../../data/notes/deid.jsonl'\n",
    "# Initialize the model config. This config file contains the various parameters of the model.\n",
    "model_config = './run/i2b2/predict_i2b2.json'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4efb34e0",
   "metadata": {},
   "source": [
    "## STEP 2: NER DATASET\n",
    "* Sentencize and tokenize the raw text. We used sentences of length 128, which includes an additional 32 context tokens on either side of the sentence. These 32 tokens serve (from the previous & next sentence) serve as additional context to the current sentence.\n",
    "* We used the en_core_sci_lg sentencizer and a custom tokenizer (can be found in the preprocessing module)\n",
    "* The dataset stored in the ner_dataset_file will be used as input to the sequence tagger model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3da763f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the dataset creator object\n",
    "dataset_creator = DatasetCreator(\n",
    "    sentencizer='en_core_sci_sm',\n",
    "    tokenizer='clinical',\n",
    "    max_tokens=128,\n",
    "    max_prev_sentence_token=32,\n",
    "    max_next_sentence_token=32,\n",
    "    default_chunk_size=32,\n",
    "    ignore_label='NA'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df7783a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This function call sentencizes and tokenizes the dataset\n",
    "# It returns a generator that iterates through the sequences.\n",
    "# We write the output to the ner_dataset_file (in json format)\n",
    "ner_notes = dataset_creator.create(\n",
    "    input_file=input_file,\n",
    "    mode='predict',\n",
    "    notation='BILOU',\n",
    "    token_text_key='text',\n",
    "    metadata_key='meta',\n",
    "    note_id_key='note_id',\n",
    "    label_key='label',\n",
    "    span_text_key='spans'\n",
    ")\n",
    "# Write to file\n",
    "with open(ner_dataset_file, 'w') as file:\n",
    "    for ner_sentence in ner_notes:\n",
    "        file.write(json.dumps(ner_sentence) + '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c11a06a",
   "metadata": {},
   "source": [
    "## STEP 3: SEQUENCE TAGGING\n",
    "* Run the sequence model - specify parameters to the sequence model in the config file (model_config). The model will be run with the specified parameters. For more information of these parameters, please refer to huggingface (or use the docs provided).\n",
    "* This file uses the argmax output. To use the recall threshold models (running the forward pass with a recall biased threshold for aggressively removing PHI) use the other config files.\n",
    "* The config files in the i2b2 direct`ory specify the model trained on only the i2b2 dataset. The config files in the mgb_i2b2 directory is for the model trained on both MGB and I2B2 datasets.\n",
    "* You can manually pass in the parameters instead of using the config file. The config file option is recommended. In our example we are passing the parameters through a config file. If you do not want to use the config file, skip the next code block and manually enter the values in the following code blocks. You will still need to read in the training args using huggingface and change values in the training args according to your needs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4cd2d28",
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = HfArgumentParser((\n",
    "    ModelArguments,\n",
    "    DataTrainingArguments,\n",
    "    EvaluationArguments,\n",
    "    TrainingArguments\n",
    "))\n",
    "# If we pass only one argument to the script and it's the path to a json file,\n",
    "# let's parse it to get our arguments.\n",
    "model_args, data_args, evaluation_args, training_args = parser.parse_json_file(json_file=model_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4fbaf42",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the sequence tagger\n",
    "sequence_tagger = SequenceTagger(\n",
    "    task_name=data_args.task_name,\n",
    "    notation=data_args.notation,\n",
    "    ner_types=data_args.ner_types,\n",
    "    model_name_or_path=model_args.model_name_or_path,\n",
    "    config_name=model_args.config_name,\n",
    "    tokenizer_name=model_args.tokenizer_name,\n",
    "    post_process=model_args.post_process,\n",
    "    cache_dir=model_args.cache_dir,\n",
    "    model_revision=model_args.model_revision,\n",
    "    use_auth_token=model_args.use_auth_token,\n",
    "    threshold=model_args.threshold,\n",
    "    do_lower_case=data_args.do_lower_case,\n",
    "    fp16=training_args.fp16,\n",
    "    seed=training_args.seed,\n",
    "    local_rank=training_args.local_rank\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "865fd6f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the required functions of the sequence tagger\n",
    "sequence_tagger.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e44d2375",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the required data and predictions of the sequence tagger\n",
    "# Can also use data_args.test_file instead of ner_dataset_file (make sure it matches ner_dataset_file)\n",
    "sequence_tagger.set_predict(\n",
    "    test_file=ner_dataset_file,\n",
    "    max_test_samples=data_args.max_predict_samples,\n",
    "    preprocessing_num_workers=data_args.preprocessing_num_workers,\n",
    "    overwrite_cache=data_args.overwrite_cache\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "840f6b23",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the huggingface trainer\n",
    "sequence_tagger.setup_trainer(training_args=training_args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb5207c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Store predictions in the specified file\n",
    "predictions = sequence_tagger.predict()\n",
    "# Write predictions to a file\n",
    "with open(predictions_file, 'w') as file:\n",
    "    for prediction in predictions:\n",
    "        file.write(json.dumps(prediction) + '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce2fa5b3",
   "metadata": {},
   "source": [
    "## STEP 4: DE-IDENTIFY TEXT\n",
    "\n",
    "* This step uses the predictions from the previous step to de-id the text. We pass the original input file where the original text is present. We look at this text and the predictions and use both of these to de-id the text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "230b504b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the text deid object\n",
    "text_deid = TextDeid(notation='BILOU', span_constraint='super_strict')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3baa3bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# De-identify the text - using deid_strategy=replace_informative doesn't drop the PHI from the text, but instead\n",
    "# labels the PHI - which you can use to drop the PHI or do any other processing.\n",
    "# If you want to drop the PHI automatically, you can use deid_strategy=remove\n",
    "deid_notes = text_deid.run_deid(\n",
    "    input_file=input_file,\n",
    "    predictions_file=predictions_file,\n",
    "    deid_strategy='replace_informative',\n",
    "    keep_age=False,\n",
    "    metadata_key='meta',\n",
    "    note_id_key='note_id',\n",
    "    tokens_key='tokens',\n",
    "    predictions_key='predictions',\n",
    "    text_key='text',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08127afe",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Write the deidentified output to a file\n",
    "with open(deid_file, 'w') as file:\n",
    "    for deid_note in deid_notes:\n",
    "        file.write(json.dumps(deid_note) + '\\n')"
   ]
  }
 ],
 "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.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}