| from lightning.pytorch import seed_everything |
| from lightning.pytorch.callbacks import ModelCheckpoint |
| from lightning.pytorch.callbacks.early_stopping import EarlyStopping |
| import lightning.pytorch as pl |
| import pandas as pd |
| from sklearn.model_selection import train_test_split |
| from transformers import AutoTokenizer |
| from ast import literal_eval |
| from pytorch_lightning.loggers import TensorBoardLogger |
|
|
| |
| import config |
| from model import CL_model |
| from dataset import CLDataModule |
|
|
| if __name__ == "__main__": |
| seed_everything(0, workers=True) |
| logger = TensorBoardLogger( |
| "/data/aiiih/projects/sunx/ccf_fuzzy_diag/train/prompt/logs", name="CL" |
| ) |
|
|
| query_df = pd.read_csv( |
| "/home/sunx/data/aiiih/projects/sunx/ccf_fuzzy_diag/data_proc/query_df.csv" |
| ) |
| query_df["concepts"] = query_df["concepts"].apply(literal_eval) |
| query_df["codes"] = query_df["codes"].apply(literal_eval) |
| query_df["codes"] = query_df["codes"].apply( |
| lambda x: [val for val in x if val is not None] |
| ) |
| train_df, val_df = train_test_split(query_df, test_size=config.split_ratio) |
|
|
| all_d = pd.read_csv( |
| "/home/sunx/data/aiiih/projects/sunx/ccf_fuzzy_diag/data_proc/query_all_d.csv" |
| ) |
| all_d.drop(columns=["finding_sites", "morphology"], inplace=True) |
| all_d["synonyms"] = all_d["synonyms"].apply(literal_eval) |
| all_d["ancestors"] = all_d["ancestors"].apply(literal_eval) |
| dictionary = dict(zip(all_d["concept"], all_d["synonyms"])) |
|
|
| pairs = pd.read_csv("/data/aiiih/projects/sunx/ccf_fuzzy_diag/data_proc/pairs.csv") |
|
|
| tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT") |
|
|
| data_module = CLDataModule(train_df, val_df, tokenizer, dictionary, all_d, pairs) |
| data_module.setup() |
|
|
| model = CL_model( |
| n_batches=len(data_module.train_dataset) / config.batch_size, |
| n_epochs=config.max_epochs, |
| lr=config.learning_rate, |
| mlm_weight=config.mlm_weight, |
| unfreeze=config.unfreeze_ratio, |
| ) |
|
|
| checkpoint = ModelCheckpoint( |
| dirpath="/home/sunx/data/aiiih/projects/sunx/ccf_fuzzy_diag/train/ckpt/v2", |
| filename="{epoch}-{step}", |
| save_weights_only=True, |
| save_last=True, |
| every_n_train_steps=config.ckcpt_every_n_steps, |
| monitor=None, |
| save_top_k=-1, |
| ) |
|
|
| trainer = pl.Trainer( |
| accelerator=config.accelerator, |
| devices=config.devices, |
| strategy="ddp", |
| logger=logger, |
| max_epochs=config.max_epochs, |
| min_epochs=config.min_epochs, |
| precision=config.precision, |
| callbacks=[ |
| EarlyStopping(monitor="val_loss", min_delta=1e-3, patience=2, mode="min"), |
| checkpoint, |
| ], |
| profiler="simple", |
| log_every_n_steps=config.log_every_n_steps, |
| ) |
|
|
| trainer.fit(model, data_module) |
|
|
