The dataset also contains ICD-9-DM codes to classify diagnoses associated with each visit. Notice that there are three diagnostic code columns. This is consistent with what we said about coded variables in the Reason-for-Visit codes section. Because ICD-9 codes are usually assigned to visits after the workup has been performed and the cause of the symptoms determined, we will have to omit them from this modeling task:

diag_cols= [
    'DIAG1','DIAG2','DIAG3',
    'PRDIAG1','PRDIAG2','PRDIAG3',
    'DIAG1R','DIAG2R','DIAG3R'
]

X_train.drop(diag_cols, axis=1, inplace=True)
X_test.drop(diag_cols, axis=1, inplace=True)