text: import gc
  import copy
  import os
  import math
  
  import config
  import utils
  import argparse
  from dataset_group import DKTDataset
  from lana_arch import LANA
  
  import numpy as np
  import pandas as pd
  from tqdm import tqdm
  from tqdm import trange
  
  import torch
  from torch import nn
  from torch.utils.data import Dataset, DataLoader
  from torch.nn import TransformerEncoder, TransformerEncoderLayer, TransformerDecoder, TransformerDecoderLayer
  from torch.optim.lr_scheduler import LambdaLR
  
  import pytorch_lightning as pl
  from pytorch_lightning.callbacks import ModelCheckpoint
  
  from sklearn.model_selection import train_test_split
  from sklearn.metrics import roc_auc_score, accuracy_score
  
  # Config
  tqdm.pandas()
  utils.set_seed(config.SEED)
  
  if __name__ == "__main__":
      parser = argparse.ArgumentParser(description="Leveled Learning")
      parser.add_argument('-d', '--data', type=str, required=True,
                          help="Filepath of the preprocessed data")
      parser.add_argument('-m', '--model', type=str, required=True,
                          help="Saved model parameters by Pytorch-Lightning")
      parser.add_argument('-n', '--num', type=int, required=True,
                          help="Number of levels")
      parser.add_argument('-t', '--top', type=int, required=True,
                          help="Top-k")
      parser.add_argument('-i', '--mu_itv', type=float, required=True,
                          help="Mean interval")
  
      args = parser.parse_args()
  
      n_models = args.num
      train_df = pd.read_pickle(f"{args.data}.train")
      val_df = pd.read_pickle(f"{args.data}.valid")
      user_performance = utils.read_pickle(f"{args.data}.user")
  
      print("train size: ", train_df.shape, "validation size: ", val_df.shape)
  
      train_loaders = []
      print("Generating dataset...")
      train_dataset = DKTDataset(train_df, max_seq=config.MAX_SEQ,
                                 min_seq=config.MIN_SEQ, overlap_seq=config.OVERLAP_SEQ,
                                 user_performance=user_performance, n_levels=n_models, mu_itv=args.mu_itv)
      train_loader = DataLoader(train_dataset,
                                batch_size=config.BATCH_SIZE,
                                num_workers=8,
                                shuffle=False,
                                pin_memory=True)
      val_dataset = DKTDataset(val_df, max_seq=config.MAX_SEQ,
                               min_seq=config.MIN_SEQ, overlap_seq=config.OVERLAP_SEQ,
                               user_performance=user_performance, n_levels=n_models, mu_itv=args.mu_itv)
      val_loader = DataLoader(val_dataset,
                              batch_size=config.BATCH_SIZE,
                              num_workers=8,
                              shuffle=False,
                              pin_memory=True)
      print(f"All dataloaders are generated")
      del train_dataset
      del val_dataset
      gc.collect()
  
      ARGS = {"d_model": config.MODEL_DIMS,
              'n_head': config.N_HEADS,
              'n_encoder': config.NUM_ENCODER,
              'n_decoder': config.NUM_DECODER,
              'dim_feedforward': config.FEEDFORWARD_DIMS,
              'dropout': config.DROPOUT,
              'max_seq': config.MAX_SEQ,
              'n_exercises': config.TOTAL_EID,
              'n_parts': config.TOTAL_PART,
              'n_resp': config.TOTAL_RESP,
              'n_etime': config.TOTAL_ETIME,
              'n_ltime_s': config.TOTAL_LTIME_S,
              'n_ltime_m': config.TOTAL_LTIME_M,
              'n_ltime_d': config.TOTAL_LTIME_D}
  
      DEVICE = f"cuda:{config.DEVICE[0]}" if torch.cuda.is_available() else "cpu"
      models = [LANA(**ARGS).to(DEVICE) for _ in range(n_models)]
      optimizers = [torch.optim.AdamW(models[i].parameters(), lr=config.LEARNING_RATE / 1000) for i in range(n_models)] # Finetune
  
      baseline, optimizer_state, model_state = utils.load_model(args.model, device=DEVICE)
      model_state = utils.remove_prefix_from_dict(model_state, "model.")  # remove prefix added by pytorch-lightning
  
      print("Baseline AUC: ", baseline)
      print("Loading models...")
      for cluster in range(n_models):
          models[cluster].load_state_dict(model_state)
  
      print("All models loaded")
      for epoch in range(config.EPOCH):
          with tqdm(total=len(train_loader), dynamic_ncols=True) as t:
              t.set_description(f"Epoch {epoch}")
              for batch_idx, batch in enumerate(train_loader):
                  for cluster in range(n_models):
                      models[cluster].train()
                      optimizers[cluster].zero_grad()
                  inputs, target, probs = batch
                  inputs = utils.dict_to_device(inputs, DEVICE)
                  target = target.to(DEVICE)
                  target_mask = (inputs["content_id"] != config.PAD)
                  probs = probs.to(DEVICE)
                  total_loss = 0.
                  target = torch.masked_select(target, target_mask)
                  for cluster in range(n_models):
                      output = models[cluster](inputs)
                      output = torch.masked_select(output, target_mask)
                      weight = torch.masked_select(probs[:, cluster].unsqueeze(-1), target_mask)
                      loss = nn.BCEWithLogitsLoss(reduction="none")(output.float(), target.float())
                      loss = loss * weight
                      loss = loss.mean()
                      total_loss += loss.item()
                      loss.backward()
                      optimizers[cluster].step()
                  t.set_postfix({"t_loss": ("%.4f" % total_loss)})
                  t.update()
          for cluster in range(n_models):
              models[cluster].eval()
  
          with torch.no_grad():
              val_labels = []
              val_outs = []
              with tqdm(total=len(val_loader), dynamic_ncols=True) as t:
                  t.set_description("Validating")
                  for batch_idx, batch in enumerate(val_loader):
                      inputs, target, probs = batch
                      inputs = utils.dict_to_device(inputs, DEVICE)
                      target = target.to(DEVICE)
                      target_mask = (inputs["content_id"] != config.PAD)
                      probs = probs.to(DEVICE)
                      probs = utils.prob_topk(probs, args.top)
                      output = torch.zeros(target.shape, device=DEVICE)
                      for cluster in range(n_models):
                          m_output = models[cluster](inputs)
                          output += (m_output * probs[:, cluster].unsqueeze(-1))
  
                      output = torch.masked_select(output, target_mask)
                      target = torch.masked_select(target, target_mask)
                      loss = nn.BCEWithLogitsLoss()(output.float(), target.float())
                      auc = roc_auc_score(target.cpu(), output.detach().float().cpu())
                      val_labels.extend(target.view(-1).data.cpu().numpy())
                      val_outs.extend(output.view(-1).data.cpu().numpy())
                      t.set_postfix({"v_loss": ("%.4f" % loss.item()), "v_auc": ("%.4f" % auc)})
                      t.update()
              real_auc = roc_auc_score(val_labels, val_outs)
              print(f"AUC: ", real_auc)
  