<jupyter_start><jupyter_text>IHM Example Using LSTMThis notebook shows an example of using `LSTM` to model In-hospital mortality from MIMIC-III dataset.Data is presumed to have been already extracted from cohort and defined via a yaml configuration as below:  Pre-ambleThe following code cell imports the required libraries and sets up the notebook<jupyter_code># Jupyter notebook specific imports
%matplotlib inline

import warnings
warnings.filterwarnings('ignore')

# Imports injecting into namespace
from tqdm.auto import tqdm
tqdm.pandas()

# General imports
import os
import json
import pickle
from pathlib import Path

import pandas as pd
import numpy as np
from getpass import getpass
import argparse

from sklearn.preprocessing import StandardScaler
from sklearn.exceptions import NotFittedError

import torch as T
from torch import nn
from pytorch_lightning import Trainer

from lightsaber import constants as C
import lightsaber.data_utils.utils as du
from lightsaber.data_utils import pt_dataset as ptd
from lightsaber.trainers import pt_trainer as ptr

from lightsaber.model_lib.pt_sota_models import rnn
import logging
log = logging.getLogger()
data_dir = Path(os.environ.get('LS_DATA_PATH', './data'))
assert data_dir.is_dir()

conf_path = os.environ.get('LS_CONF_PATH', os.path.abspath('./ihm_expt_config.yml')) 
expt_conf = du.yaml.load(open(conf_path).read().format(DATA_DIR=data_dir),
                         Loader=du._Loader)<jupyter_output><empty_output><jupyter_text>IHM Model TrainingIn general, user need to follow the following steps to train a `HistGBT` for IHM model.* _Data Ingestion_: The first step involves setting up the pre-processors to train an IHM model. In this example, we will use a `StandardScaler` from `scikit-learn` using filters defined within lightsaber.  - We would next read the train, test, and validation dataset. In some cases, users may also want to define a calibration dataset    * _Model Definition_: We would next need to define a base model for classification. In this example, we will use a pre-packaged `LSTM` model from  `lightsaber`* _Model Training_: Once the models are defined, we can use `lightsaber` to train the model via the pre-packaged `PyModel` and the corresponding trainer code. This step will also generate the relevant `metrics` for this problem.  Data IngestionWe firs start by reading extracted cohort data and use a `StandardScaler` demonstrating the proper usage of a pre-processor<jupyter_code>preprocessor = StandardScaler()
train_filter = [ptd.filter_preprocessor(cols=expt_conf['numerical'], 
                                        preprocessor=preprocessor,
                                        refit=True),
                ptd.filter_fillna(fill_value=expt_conf['normal_values'],
                                  time_order_col=expt_conf['time_order_col'])
                ]
transform = ptd.transform_drop_cols(cols_to_drop=expt_conf['time_order_col'])
train_dataset = ptd.BaseDataset(tgt_file=expt_conf['train']['tgt_file'],
                                feat_file=expt_conf['train']['feat_file'],
                                idx_col=expt_conf['idx_cols'],
                                tgt_col=expt_conf['tgt_col'],
                                feat_columns=expt_conf['feat_cols'],
                                time_order_col=expt_conf['time_order_col'],
                                category_map=expt_conf['category_map'],
                                transform=transform,
                                filter=train_filter,
                               )
# print(train_dataset.data.head())
print(train_dataset.shape, len(train_dataset))
# For other datasets use fitted preprocessors
fitted_filter = [ptd.filter_preprocessor(cols=expt_conf['numerical'], 
                                         preprocessor=preprocessor, refit=False),
                 ptd.filter_fillna(fill_value=expt_conf['normal_values'],
                                   time_order_col=expt_conf['time_order_col'])
                 ]

val_dataset = ptd.BaseDataset(tgt_file=expt_conf['val']['tgt_file'],
                              feat_file=expt_conf['val']['feat_file'],
                              idx_col=expt_conf['idx_cols'],
                              tgt_col=expt_conf['tgt_col'],
                              feat_columns=expt_conf['feat_cols'],
                              time_order_col=expt_conf['time_order_col'],
                              category_map=expt_conf['category_map'],
                              transform=transform,
                              filter=fitted_filter,
                              )

test_dataset = ptd.BaseDataset(tgt_file=expt_conf['test']['tgt_file'],
                               feat_file=expt_conf['test']['feat_file'],
                               idx_col=expt_conf['idx_cols'],
                               tgt_col=expt_conf['tgt_col'],
                               feat_columns=expt_conf['feat_cols'],
                               time_order_col=expt_conf['time_order_col'],
                               category_map=expt_conf['category_map'],
                               transform=transform,
                               filter=fitted_filter,
                               )

print(val_dataset.shape, len(val_dataset))
print(test_dataset.shape, len(test_dataset))
# Handling imbala
input_dim, target_dim = train_dataset.shape
output_dim = 2

weight_labels = train_dataset.target.iloc[:, 0].value_counts()
weight_labels = (weight_labels.max() / ((weight_labels + 0.0000001) ** (1)))
weight_labels.sort_index(inplace=True)
weights = T.FloatTensor(weight_labels.values).to(train_dataset.device)
print(weights)<jupyter_output><empty_output><jupyter_text>Single Run<jupyter_code># For most models you need to change only this part
hparams = argparse.Namespace(gpus=[0],
                             lr=0.01,
                             max_epochs=50,
                             batch_size=32,
                             hidden_dim=32,
                             rnn_class='LSTM',
                             n_layers=2,
                             dropout=0.1,
                             recurrent_dropout=0.1,
                             bidirectional=False,
                             )

hparams.rnn_class = C.PYTORCH_CLASS_DICT[hparams.rnn_class]

base_model = rnn.RNNClassifier(input_dim, output_dim, 
                               hidden_dim=hparams.hidden_dim,
                               rnn_class=hparams.rnn_class,
                               n_layers=hparams.n_layers,
                               dropout=hparams.dropout,
                               recurrent_dropout=hparams.recurrent_dropout,
                               bidirectional=hparams.bidirectional
                              )

criterion = nn.CrossEntropyLoss(weight=weights)
# optimizer = T.optim.Adam(base_model.parameters(),
#                          lr=hparams.lr,
#                          weight_decay=1e-5  # standard value)
#                          )

# scheduler = T.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min')

# Creating the wrapped model
wrapped_model = ptr.PyModel(hparams, base_model,
                            train_dataset=train_dataset,
                            val_dataset=val_dataset, # None
                            test_dataset=None, #test_dataset, # test_dataset
                            #optimizer=optimizer,
                            loss_func=criterion,
                            #scheduler=scheduler,
                            collate_fn=ptd.collate_fn
                            )
# Training 
overfit_pct, fast_dev_run, terminate_on_nan, auto_lr_find = 0, False, False, False

trainer = Trainer(max_epochs=hparams.max_epochs, 
                  gpus=hparams.gpus,
                  default_root_dir=os.path.join('./out/', 'classifier_ihm'),
                  terminate_on_nan=terminate_on_nan,
                  auto_lr_find=auto_lr_find,
                  overfit_pct=overfit_pct,
                  fast_dev_run=fast_dev_run #True if devugging
                 )
mlflow_conf = dict(experiment_name=f'classifier_ihm')
artifacts = dict(preprocessor=preprocessor, 
                 weight_labels=weight_labels,
                )
experiment_tags = dict(model='RNNClassifier',
                       input_dim=input_dim,
                       output_dim=output_dim
                      )

(run_id, metrics, 
 val_y, val_yhat, val_pred_proba, 
 test_y, test_yhat, test_pred_proba) = ptr.run_training_with_mlflow(mlflow_conf, 
                                                                    trainer, 
                                                                    wrapped_model, 
                                                                    overfit_pct=overfit_pct,
                                                                    artifacts=artifacts,
                                                                    **experiment_tags)

print(f"MLFlow Experiment: {mlflow_conf['experiment_name']} \t | Run ID: {run_id}")
print(metrics)<jupyter_output><empty_output><jupyter_text>IHM Model RegistrationThis block shows how to register a model for subsequent steps. Given a `run_id` this block can be run independtly of other aspectsInternally, the following steps happen:- a saved model (along with hyper-params and weights) is retrieved using `run_id`- model is initialized using the weights- model is logged to mlflow under registered model name<jupyter_code>print(f"Registering model for run: {run_id}")
# Reading things from mlflow
# Model coders can create functions to repeat this - part of model init
import torch
from lightsaber.trainers import helper
data_dir = Path(os.environ.get('LS_DATA_PATH', './data'))
assert data_dir.is_dir()

conf_path = os.environ.get('LS_CONF_PATH', os.path.abspath('./ihm_expt_config.yml')) 
expt_conf = du.yaml.load(open(conf_path).read().format(DATA_DIR=data_dir),
                         Loader=du._Loader)

mlflow_conf = dict(experiment_name=f'classifier_ihm')
registered_model_name = 'classifier_ihm_rnn_v0'

## Loading model attributes from mlflow
mlflow_setup = helper.setup_mlflow(**mlflow_conf)
run_data = helper.fetch_mlflow_run(run_id, 
                                   mlflow_uri=mlflow_setup['mlflow_uri'],
                                   artifacts_prefix=['artifact/weight_labels'],
                                   parse_params=True
                                  )
hparams = run_data['params']

hparams = argparse.Namespace(**hparams)
hparams.rnn_class = helper.import_model_class(hparams.rnn_class.split("'")[1::2][0])

weight_labels = pickle.load(open(helper.get_artifact_path(run_data['artifact_paths'][0],
                                   artifact_uri=run_data['info'].artifact_uri), 'rb'))

weights = T.FloatTensor(weight_labels.values)

## Setting model weights
base_model = rnn.RNNClassifier(input_dim=input_dim, 
                               output_dim=output_dim, 
                               hidden_dim=hparams.hidden_dim,
                               rnn_class=hparams.rnn_class,
                               n_layers=hparams.n_layers,
                               dropout=hparams.dropout,
                               recurrent_dropout=hparams.recurrent_dropout,
                               bidirectional=hparams.bidirectional
                              )

criterion = nn.CrossEntropyLoss(weight=weights)

wrapped_model = ptr.PyModel(hparams, base_model,
                            train_dataset=None,
                            val_dataset=None, # None
                            test_dataset=None, # test_dataset
                            #optimizer=optimizer,
                            loss_func=criterion,
                            #scheduler=scheduler,
                            collate_fn=ptd.collate_fn
                            )
# Recreate models
base_model = rnn.RNNClassifier(input_dim=int(run_data['tags']['input_dim']),
                               output_dim=int(run_data['tags']['output_dim']), 
                               hidden_dim=hparams.hidden_dim,
                               rnn_class=hparams.rnn_class,
                               n_layers=hparams.n_layers,
                               dropout=hparams.dropout,
                               recurrent_dropout=hparams.recurrent_dropout,
                               bidirectional=hparams.bidirectional
                               )
criterion = nn.CrossEntropyLoss(weight=weights)


# Creating the wrapped model
wrapped_model = ptr.PyModel(hparams, base_model,
                            train_dataset=None,
                            val_dataset=None, # None
                            test_dataset=None, # test_dataset
                            cal_dataset=None,
                            loss_func=criterion,
                            collate_fn=ptd.collate_fn
                            )
print('model ready for logging')
# Register model
ptr.register_model_with_mlflow(run_id, mlflow_conf, wrapped_model, 
                               registered_model_name=registered_model_name,
                               test_feat_file=expt_conf['test']['feat_file'],
                               test_tgt_file=expt_conf['test']['tgt_file'],
                               config=os.path.abspath('./ihm_expt_config.yml'),
                               model_path='model_checkpoint'
                              )<jupyter_output><empty_output><jupyter_text>IHM Model Inference `Lightsaber` also natively supports conducting inferences on new patients using the registered model. The key steps involve:* loading the registerd model from mlflow* Ingest the new test data using `BaseDataset` in inference mode (setting `tgt_file` to `None`)* Use the `PyModel.predict_patient` method to generate inference for the patient of interest It is to be noted, for the first step, users may need to perform additional setup as show below<jupyter_code>print(f"Inference using model for run: {run_id}")
# Reading things from mlflow
# Model coders can create functions to repeat this - part of model init
import torch
from lightsaber.trainers import helper
data_dir = Path(os.environ.get('LS_DATA_PATH', './data'))
assert data_dir.is_dir()

conf_path = os.environ.get('LS_CONF_PATH', os.path.abspath('./ihm_expt_config.yml')) 
expt_conf = du.yaml.load(open(conf_path).read().format(DATA_DIR=data_dir),
                         Loader=du._Loader)

mlflow_conf = dict(experiment_name=f'classifier_ihm')
registered_model_name = 'classifier_ihm_rnn_v0'

## Loading model attributes from mlflow
mlflow_setup = helper.setup_mlflow(**mlflow_conf)
run_data = helper.fetch_mlflow_run(run_id, 
                                   mlflow_uri=mlflow_setup['mlflow_uri'],
                                   artifacts_prefix=['artifact/weight_labels'],
                                   parse_params=True
                                  )
hparams = run_data['params']

hparams = argparse.Namespace(**hparams)
hparams.rnn_class = helper.import_model_class(hparams.rnn_class.split("'")[1::2][0])

weight_labels = pickle.load(open(helper.get_artifact_path(run_data['artifact_paths'][0],
                                   artifact_uri=run_data['info'].artifact_uri), 'rb'))

weights = T.FloatTensor(weight_labels.values)

## Setting model weights
base_model = rnn.RNNClassifier(input_dim=input_dim, 
                               output_dim=output_dim, 
                               hidden_dim=hparams.hidden_dim,
                               rnn_class=hparams.rnn_class,
                               n_layers=hparams.n_layers,
                               dropout=hparams.dropout,
                               recurrent_dropout=hparams.recurrent_dropout,
                               bidirectional=hparams.bidirectional
                              )

criterion = nn.CrossEntropyLoss(weight=weights)

wrapped_model = ptr.PyModel(hparams, base_model,
                            train_dataset=None,
                            val_dataset=None, # None
                            test_dataset=None, # test_dataset
                            #optimizer=optimizer,
                            loss_func=criterion,
                            #scheduler=scheduler,
                            collate_fn=ptd.collate_fn
                            )
# Loading saved model from mlflow
wrapped_model = ptr.load_model_from_mlflow(run_id, mlflow_conf, wrapped_model)
inference_dataloader = ptd.BaseDataset(tgt_file=None,
                                       feat_file=expt_conf['test']['feat_file'],
                                       idx_col=expt_conf['idx_cols'],
                                       tgt_col=expt_conf['tgt_col'],
                                       feat_columns=expt_conf['feat_cols'],
                                       time_order_col=expt_conf['time_order_col'],
                                       category_map=expt_conf['category_map'],
                                       transform=transform,
                                       filter=fitted_filter,
                                       )

print(inference_dataloader.shape, len(inference_dataloader))
patient_id = inference_dataloader.sample_idx.index[0]
print(f"Inference for patient: {patient_id}")

# patient_id = '10011_episode1_timeseries.csv'
wrapped_model.predict_patient(patient_id, inference_dataloader)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.07862084738978786, 'fraction_numerical': 0.0026302247126765393, 'mean_word_length': 1.864864864864865, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 33, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '487525', 'license': 'Apache-2.0', 'n_tokens_mistral': 4811, 'n_tokens_neox': 4658, 'n_words': 874, 'path': 'lightsaber/examples/in_hosptial_mortality/Exp_LSTM.ipynb', 'repo_name': 'IBM/DPM360', 'size': '17489'}