def create_video_note(
            self,
            title,
            video_path: str,
            desc: str,
            cover_path: str = None,
            ats: list = None,
            post_time: str = None,
            topics: list = None,
            is_private: bool = False,
            wait_time: int = 3,
    ):
        """发布视频笔记

        :param title: 笔记标题
        :param video_path: 视频文件路径，目前只支持本地路径
        :param desc: 笔记详情
        :param cover_path: 可选，封面文件路径
        :param ats: 可选，@用户信息
        :param post_time: 可选，发布时间
        :param topics: 可选，话题信息
        :param is_private: 可选，是否私密发布
        :param wait_time: 可选，默认 3 s，循环等待获取视频第一帧为笔记封面
        :return:
        :rtype: object
        """
        if ats is None:
            ats = []
        if topics is None:
            topics = []

        file_id, token = self.get_upload_files_permit("video")
        res = self.upload_file(
            file_id,
            token,
            video_path,
            content_type="video/mp4",
        )
        video_id, is_upload = res.headers["X-Ros-Video-Id"], False

        image_id = None
        if cover_path is None:
            for _ in range(10):
                time.sleep(wait_time)
                image_id = self.get_video_first_frame_image_id(video_id)
                if image_id:
                    break

        if cover_path:
            is_upload = True
            image_id, token = self.get_upload_files_permit("image")
            self.upload_file(image_id, token, cover_path)

        cover_info = {
            "file_id": image_id,
            "frame": {"ts": 0, "is_user_select": False, "is_upload": is_upload},
        }

        video_info = {
            "file_id": file_id,
            "timelines": [],
            "cover": cover_info,
            "chapters": [],
            "chapter_sync_text": False,
            "entrance": "web",
        }
        return self.create_note(title, desc, NoteType.VIDEO.value, ats=ats, topics=topics, video_info=video_info,
                                post_time=post_time, is_private=is_private)

def create_video_note(
            self,
            title,
            video_path: str,
            desc: str,
            cover_path: str = None,
            ats: list = None,
            post_time: str = None,
            topics: list = None,
            is_private: bool = False,
            wait_time: int = 3,
    ):

        if ats is None:
            ats = []
        if topics is None:
            topics = []

        file_id, token = self.get_upload_files_permit("video")
        res = self.upload_file(
            file_id,
            token,
            video_path,
            content_type="video/mp4",
        )
        video_id, is_upload = res.headers["X-Ros-Video-Id"], False

        image_id = None
        if cover_path is None:
            for _ in range(10):
                time.sleep(wait_time)
                image_id = self.get_video_first_frame_image_id(video_id)
                if image_id:
                    break

        if cover_path:
            is_upload = True
            image_id, token = self.get_upload_files_permit("image")
            self.upload_file(image_id, token, cover_path)

        cover_info = {
            "file_id": image_id,
            "frame": {"ts": 0, "is_user_select": False, "is_upload": is_upload},
        }

        video_info = {
            "file_id": file_id,
            "timelines": [],
            "cover": cover_info,
            "chapters": [],
            "chapter_sync_text": False,
            "entrance": "web",
        }
        return self.create_note(title, desc, NoteType.VIDEO.value, ats=ats, topics=topics, video_info=video_info,
                                post_time=post_time, is_private=is_private)


@abstractproperty
    def name(self) -> str:
        """
        returns the name of the encoder
        """

"""
        returns the name of the encoder
        """


def __init__(self, equ_h: int, equ_w: int):
        """Args:
        equ_h: (int) the height of the generated equirect
        equ_w: (int) the width of the generated equirect
        """

        # Cubemap input
        input_projections = get_cubemap_projections(equ_h,equ_h)

        # Equirectangular output
        output_projection = EquirectProjection(equ_h, equ_w)
        super(Cube2Equirect, self).__init__(
            input_projections, output_projection
        )

def __init__(self, equ_h: int, equ_w: int):

        # Cubemap input
        input_projections = get_cubemap_projections(equ_h,equ_h)

        # Equirectangular output
        output_projection = EquirectProjection(equ_h, equ_w)
        super(Cube2Equirect, self).__init__(
            input_projections, output_projection
        )


def get_coor_colors(obj_labels):
    """
    Args:
        obj_labels: 1 is ground, labels > 1 indicates different instance cluster

    Returns:
        rgb: [N, 3]. color for each point.
    """
    colors = matplotlib.colors.XKCD_COLORS.values()
    max_color_num = obj_labels.max()

    color_list = list(colors)[:max_color_num+1]
    colors_rgba = [matplotlib.colors.to_rgba_array(color) for color in color_list]
    label_rgba = np.array(colors_rgba)[obj_labels]
    label_rgba = label_rgba.squeeze()[:, :3]

    return label_rgba

def get_coor_colors(obj_labels):

    colors = matplotlib.colors.XKCD_COLORS.values()
    max_color_num = obj_labels.max()

    color_list = list(colors)[:max_color_num+1]
    colors_rgba = [matplotlib.colors.to_rgba_array(color) for color in color_list]
    label_rgba = np.array(colors_rgba)[obj_labels]
    label_rgba = label_rgba.squeeze()[:, :3]

    return label_rgba


def apply_coords(self, coords: np.ndarray, original_size: Tuple[int, ...]) -> np.ndarray:
        """
        Expects a numpy array of length 2 in the final dimension. Requires the
        original image size in (H, W) format.
        """
        old_h, old_w = original_size
        new_h, new_w = self.get_preprocess_shape(
            original_size[0], original_size[1], self.target_length
        )
        coords = deepcopy(coords).astype(float)
        coords[..., 0] = coords[..., 0] * (new_w / old_w)
        coords[..., 1] = coords[..., 1] * (new_h / old_h)
        return coords

def apply_coords(self, coords: np.ndarray, original_size: Tuple[int, ...]) -> np.ndarray:

        old_h, old_w = original_size
        new_h, new_w = self.get_preprocess_shape(
            original_size[0], original_size[1], self.target_length
        )
        coords = deepcopy(coords).astype(float)
        coords[..., 0] = coords[..., 0] * (new_w / old_w)
        coords[..., 1] = coords[..., 1] * (new_h / old_h)
        return coords


def f_train_wrapper(args, pt_model, loss_wrapper, device, \
                    optimizer_wrapper, \
                    train_dataset_wrapper, \
                    val_dataset_wrapper = None, \
                    checkpoint = None):
    """
    f_train_wrapper(args, pt_model, loss_wrapper, device,
                    optimizer_wrapper
                    train_dataset_wrapper, val_dataset_wrapper = None,
                    check_point = None):
      A wrapper to run the training process

    Args:
       args:         argument information given by argpase
       pt_model:     pytorch model (torch.nn.Module)
       loss_wrapper: a wrapper over loss function
                     loss_wrapper.compute(generated, target)
       device:       torch.device("cuda") or torch.device("cpu")

       optimizer_wrapper:
           a wrapper over optimizer (defined in op_manager.py)
           optimizer_wrapper.optimizer is torch.optimizer

       train_dataset_wrapper:
           a wrapper over training data set (data_io/default_data_io.py)
           train_dataset_wrapper.get_loader() returns torch.DataSetLoader

       val_dataset_wrapper:
           a wrapper over validation data set (data_io/default_data_io.py)
           it can None.

       check_point:
           a check_point that stores every thing to resume training
    """

    nii_display.f_print_w_date("Start model training")

    ##############
    ## Preparation
    ##############

    # get the optimizer
    optimizer_wrapper.print_info()
    optimizer = optimizer_wrapper.optimizer
    lr_scheduler = optimizer_wrapper.lr_scheduler
    epoch_num = optimizer_wrapper.get_epoch_num()
    no_best_epoch_num = optimizer_wrapper.get_no_best_epoch_num()

    # get data loader for training set
    train_dataset_wrapper.print_info()
    train_data_loader = train_dataset_wrapper.get_loader()
    train_seq_num = train_dataset_wrapper.get_seq_num()

    # get the training process monitor
    monitor_trn = nii_monitor.Monitor(epoch_num, train_seq_num)

    # if validation data is provided, get data loader for val set
    if val_dataset_wrapper is not None:
        val_dataset_wrapper.print_info()
        val_data_loader = val_dataset_wrapper.get_loader()
        val_seq_num = val_dataset_wrapper.get_seq_num()
        monitor_val = nii_monitor.Monitor(epoch_num, val_seq_num)
    else:
        monitor_val = None

    # training log information
    train_log = ''

    # prepare for DataParallism if available
    # pytorch.org/tutorials/beginner/blitz/data_parallel_tutorial.html
    if torch.cuda.device_count() > 1 and args.multi_gpu_data_parallel:
        flag_multi_device = True
        nii_display.f_print("\nUse %d GPUs\n" % (torch.cuda.device_count()))
        # no way to call normtarget_f after pt_model is in DataParallel
        normtarget_f = pt_model.normalize_target
        pt_model = nn.DataParallel(pt_model)
    else:
        nii_display.f_print("\nUse single GPU: %s\n" % \
                            (torch.cuda.get_device_name(device)))
        flag_multi_device = False
        normtarget_f = None
    pt_model.to(device, dtype=nii_dconf.d_dtype)

    # print the network
    nii_nn_tools.f_model_show(pt_model)
    nii_nn_tools.f_loss_show(loss_wrapper)

    ###############################
    ## Resume training if necessary
    ###############################
    # resume training or initialize the model if necessary
    cp_names = nii_nn_manage_conf.CheckPointKey()
    if checkpoint is not None:
        if type(checkpoint) is dict:
            # checkpoint

            # load model parameter and optimizer state
            if cp_names.state_dict in checkpoint:
                # wrap the state_dic in f_state_dict_wrapper
                # in case the model is saved when DataParallel is on
                pt_model.load_state_dict(
                    nii_nn_tools.f_state_dict_wrapper(
                        checkpoint[cp_names.state_dict],
                        flag_multi_device))

            # load optimizer state
            if cp_names.optimizer in checkpoint and \
               not args.ignore_optimizer_statistics_in_trained_model:
                optimizer.load_state_dict(checkpoint[cp_names.optimizer])

            # optionally, load training history
            if not args.ignore_training_history_in_trained_model:
                #nii_display.f_print("Load ")
                if cp_names.trnlog in checkpoint:
                    monitor_trn.load_state_dic(
                        checkpoint[cp_names.trnlog])
                if cp_names.vallog in checkpoint and monitor_val:
                    monitor_val.load_state_dic(
                        checkpoint[cp_names.vallog])
                if cp_names.info in checkpoint:
                    train_log = checkpoint[cp_names.info]
                if cp_names.lr_scheduler in checkpoint and \
                   checkpoint[cp_names.lr_scheduler] and lr_scheduler.f_valid():
                    lr_scheduler.f_load_state_dict(
                        checkpoint[cp_names.lr_scheduler])

                nii_display.f_print("Load check point, resume training")
            else:
                nii_display.f_print("Load pretrained model and optimizer")
        else:
            # only model status
            pt_model.load_state_dict(
                nii_nn_tools.f_state_dict_wrapper(
                    checkpoint, flag_multi_device))
            nii_display.f_print("Load pretrained model")

    ######################
    ### User defined setup
    ######################
    if hasattr(pt_model, "other_setups"):
        nii_display.f_print("Conduct User-defined setup")
        pt_model.other_setups()

    # This should be merged with other_setups
    if hasattr(pt_model, "g_pretrained_model_path") and \
       hasattr(pt_model, "g_pretrained_model_prefix"):
        nii_display.f_print("Load pret-rained models as part of this mode")
        nii_nn_tools.f_load_pretrained_model_partially(
            pt_model, pt_model.g_pretrained_model_path,
            pt_model.g_pretrained_model_prefix)

    ######################
    ### Start training
    ######################
    # other variables
    flag_early_stopped = False
    start_epoch = monitor_trn.get_epoch()
    epoch_num = monitor_trn.get_max_epoch()

    # print
    _ = nii_op_display_tk.print_log_head()
    nii_display.f_print_message(train_log, flush=True, end='')

    # loop over multiple epochs
    for epoch_idx in range(start_epoch, epoch_num):

        # training one epoch
        pt_model.train()
        # set validation flag if necessary
        if hasattr(pt_model, 'validation'):
            pt_model.validation = False
            mes = "Warning: model.validation is deprecated, "
            mes += "please use model.flag_validation"
            nii_display.f_print(mes, 'warning')
        if hasattr(pt_model, 'flag_validation'):
            pt_model.flag_validation = False

        f_run_one_epoch(args, pt_model, loss_wrapper, device, \
                        monitor_trn, train_data_loader, \
                        epoch_idx, optimizer, normtarget_f)
        time_trn = monitor_trn.get_time(epoch_idx)
        loss_trn = monitor_trn.get_loss(epoch_idx)

        # if necessary, do validataion
        if val_dataset_wrapper is not None:
            # set eval() if necessary
            if args.eval_mode_for_validation:
                pt_model.eval()

            # set validation flag if necessary
            if hasattr(pt_model, 'validation'):
                pt_model.validation = True
                mes = "Warning: model.validation is deprecated, "
                mes += "please use model.flag_validation"
                nii_display.f_print(mes, 'warning')
            if hasattr(pt_model, 'flag_validation'):
                pt_model.flag_validation = True

            with torch.no_grad():
                f_run_one_epoch(args, pt_model, loss_wrapper, \
                                device, \
                                monitor_val, val_data_loader, \
                                epoch_idx, None, normtarget_f)
            time_val = monitor_val.get_time(epoch_idx)
            loss_val = monitor_val.get_loss(epoch_idx)

            # update lr rate scheduler if necessary
            if lr_scheduler.f_valid():
                lr_scheduler.f_step(loss_val)

        else:
            time_val, loss_val = 0, 0

        if val_dataset_wrapper is not None:
            flag_new_best = monitor_val.is_new_best()
        else:
            flag_new_best = True

        # print information
        train_log += nii_op_display_tk.print_train_info(
            epoch_idx, time_trn, loss_trn, time_val, loss_val,
            flag_new_best, optimizer_wrapper.get_lr_info())

        # save the best model
        if flag_new_best:
            tmp_best_name = nii_nn_tools.f_save_trained_name(args)
            torch.save(pt_model.state_dict(), tmp_best_name)

        # save intermediate model if necessary
        if not args.not_save_each_epoch:
            tmp_model_name = nii_nn_tools.f_save_epoch_name(args, epoch_idx)

            if monitor_val is not None:
                tmp_val_log = monitor_val.get_state_dic()
            else:
                tmp_val_log = None

            if lr_scheduler.f_valid():
                lr_scheduler_state = lr_scheduler.f_state_dict()
            else:
                lr_scheduler_state = None

            # save
            tmp_dic = {
                cp_names.state_dict : pt_model.state_dict(),
                cp_names.info : train_log,
                cp_names.optimizer : optimizer.state_dict(),
                cp_names.trnlog : monitor_trn.get_state_dic(),
                cp_names.vallog : tmp_val_log,
                cp_names.lr_scheduler : lr_scheduler_state
            }
            torch.save(tmp_dic, tmp_model_name)
            if args.verbose == 1:
                nii_display.f_eprint(str(datetime.datetime.now()))
                nii_display.f_eprint("Save {:s}".format(tmp_model_name),
                                     flush=True)

        # Early stopping
        #  note: if LR scheduler is used, early stopping will be
        #  disabled
        if lr_scheduler.f_allow_early_stopping() and \
           monitor_val is not None and \
           monitor_val.should_early_stop(no_best_epoch_num):
            flag_early_stopped = True
            break

    # loop done
    nii_op_display_tk.print_log_tail()
    if flag_early_stopped:
        nii_display.f_print("Training finished by early stopping")
    else:
        nii_display.f_print("Training finished")
    nii_display.f_print("Model is saved to", end = '')
    nii_display.f_print("{}".format(nii_nn_tools.f_save_trained_name(args)))
    return

def f_train_wrapper(args, pt_model, loss_wrapper, device, \
                    optimizer_wrapper, \
                    train_dataset_wrapper, \
                    val_dataset_wrapper = None, \
                    checkpoint = None):

    nii_display.f_print_w_date("Start model training")

    ##############
    ## Preparation
    ##############

    # get the optimizer
    optimizer_wrapper.print_info()
    optimizer = optimizer_wrapper.optimizer
    lr_scheduler = optimizer_wrapper.lr_scheduler
    epoch_num = optimizer_wrapper.get_epoch_num()
    no_best_epoch_num = optimizer_wrapper.get_no_best_epoch_num()

    # get data loader for training set
    train_dataset_wrapper.print_info()
    train_data_loader = train_dataset_wrapper.get_loader()
    train_seq_num = train_dataset_wrapper.get_seq_num()

    # get the training process monitor
    monitor_trn = nii_monitor.Monitor(epoch_num, train_seq_num)

    # if validation data is provided, get data loader for val set
    if val_dataset_wrapper is not None:
        val_dataset_wrapper.print_info()
        val_data_loader = val_dataset_wrapper.get_loader()
        val_seq_num = val_dataset_wrapper.get_seq_num()
        monitor_val = nii_monitor.Monitor(epoch_num, val_seq_num)
    else:
        monitor_val = None

    # training log information
    train_log = ''

    # prepare for DataParallism if available
    # pytorch.org/tutorials/beginner/blitz/data_parallel_tutorial.html
    if torch.cuda.device_count() > 1 and args.multi_gpu_data_parallel:
        flag_multi_device = True
        nii_display.f_print("\nUse %d GPUs\n" % (torch.cuda.device_count()))
        # no way to call normtarget_f after pt_model is in DataParallel
        normtarget_f = pt_model.normalize_target
        pt_model = nn.DataParallel(pt_model)
    else:
        nii_display.f_print("\nUse single GPU: %s\n" % \
                            (torch.cuda.get_device_name(device)))
        flag_multi_device = False
        normtarget_f = None
    pt_model.to(device, dtype=nii_dconf.d_dtype)

    # print the network
    nii_nn_tools.f_model_show(pt_model)
    nii_nn_tools.f_loss_show(loss_wrapper)

    ###############################
    ## Resume training if necessary
    ###############################
    # resume training or initialize the model if necessary
    cp_names = nii_nn_manage_conf.CheckPointKey()
    if checkpoint is not None:
        if type(checkpoint) is dict:
            # checkpoint

            # load model parameter and optimizer state
            if cp_names.state_dict in checkpoint:
                # wrap the state_dic in f_state_dict_wrapper
                # in case the model is saved when DataParallel is on
                pt_model.load_state_dict(
                    nii_nn_tools.f_state_dict_wrapper(
                        checkpoint[cp_names.state_dict],
                        flag_multi_device))

            # load optimizer state
            if cp_names.optimizer in checkpoint and \
               not args.ignore_optimizer_statistics_in_trained_model:
                optimizer.load_state_dict(checkpoint[cp_names.optimizer])

            # optionally, load training history
            if not args.ignore_training_history_in_trained_model:
                #nii_display.f_print("Load ")
                if cp_names.trnlog in checkpoint:
                    monitor_trn.load_state_dic(
                        checkpoint[cp_names.trnlog])
                if cp_names.vallog in checkpoint and monitor_val:
                    monitor_val.load_state_dic(
                        checkpoint[cp_names.vallog])
                if cp_names.info in checkpoint:
                    train_log = checkpoint[cp_names.info]
                if cp_names.lr_scheduler in checkpoint and \
                   checkpoint[cp_names.lr_scheduler] and lr_scheduler.f_valid():
                    lr_scheduler.f_load_state_dict(
                        checkpoint[cp_names.lr_scheduler])

                nii_display.f_print("Load check point, resume training")
            else:
                nii_display.f_print("Load pretrained model and optimizer")
        else:
            # only model status
            pt_model.load_state_dict(
                nii_nn_tools.f_state_dict_wrapper(
                    checkpoint, flag_multi_device))
            nii_display.f_print("Load pretrained model")

    ######################
    ### User defined setup
    ######################
    if hasattr(pt_model, "other_setups"):
        nii_display.f_print("Conduct User-defined setup")
        pt_model.other_setups()

    # This should be merged with other_setups
    if hasattr(pt_model, "g_pretrained_model_path") and \
       hasattr(pt_model, "g_pretrained_model_prefix"):
        nii_display.f_print("Load pret-rained models as part of this mode")
        nii_nn_tools.f_load_pretrained_model_partially(
            pt_model, pt_model.g_pretrained_model_path,
            pt_model.g_pretrained_model_prefix)

    ######################
    ### Start training
    ######################
    # other variables
    flag_early_stopped = False
    start_epoch = monitor_trn.get_epoch()
    epoch_num = monitor_trn.get_max_epoch()

    # print
    _ = nii_op_display_tk.print_log_head()
    nii_display.f_print_message(train_log, flush=True, end='')

    # loop over multiple epochs
    for epoch_idx in range(start_epoch, epoch_num):

        # training one epoch
        pt_model.train()
        # set validation flag if necessary
        if hasattr(pt_model, 'validation'):
            pt_model.validation = False
            mes = "Warning: model.validation is deprecated, "
            mes += "please use model.flag_validation"
            nii_display.f_print(mes, 'warning')
        if hasattr(pt_model, 'flag_validation'):
            pt_model.flag_validation = False

        f_run_one_epoch(args, pt_model, loss_wrapper, device, \
                        monitor_trn, train_data_loader, \
                        epoch_idx, optimizer, normtarget_f)
        time_trn = monitor_trn.get_time(epoch_idx)
        loss_trn = monitor_trn.get_loss(epoch_idx)

        # if necessary, do validataion
        if val_dataset_wrapper is not None:
            # set eval() if necessary
            if args.eval_mode_for_validation:
                pt_model.eval()

            # set validation flag if necessary
            if hasattr(pt_model, 'validation'):
                pt_model.validation = True
                mes = "Warning: model.validation is deprecated, "
                mes += "please use model.flag_validation"
                nii_display.f_print(mes, 'warning')
            if hasattr(pt_model, 'flag_validation'):
                pt_model.flag_validation = True

            with torch.no_grad():
                f_run_one_epoch(args, pt_model, loss_wrapper, \
                                device, \
                                monitor_val, val_data_loader, \
                                epoch_idx, None, normtarget_f)
            time_val = monitor_val.get_time(epoch_idx)
            loss_val = monitor_val.get_loss(epoch_idx)

            # update lr rate scheduler if necessary
            if lr_scheduler.f_valid():
                lr_scheduler.f_step(loss_val)

        else:
            time_val, loss_val = 0, 0

        if val_dataset_wrapper is not None:
            flag_new_best = monitor_val.is_new_best()
        else:
            flag_new_best = True

        # print information
        train_log += nii_op_display_tk.print_train_info(
            epoch_idx, time_trn, loss_trn, time_val, loss_val,
            flag_new_best, optimizer_wrapper.get_lr_info())

        # save the best model
        if flag_new_best:
            tmp_best_name = nii_nn_tools.f_save_trained_name(args)
            torch.save(pt_model.state_dict(), tmp_best_name)

        # save intermediate model if necessary
        if not args.not_save_each_epoch:
            tmp_model_name = nii_nn_tools.f_save_epoch_name(args, epoch_idx)

            if monitor_val is not None:
                tmp_val_log = monitor_val.get_state_dic()
            else:
                tmp_val_log = None

            if lr_scheduler.f_valid():
                lr_scheduler_state = lr_scheduler.f_state_dict()
            else:
                lr_scheduler_state = None

            # save
            tmp_dic = {
                cp_names.state_dict : pt_model.state_dict(),
                cp_names.info : train_log,
                cp_names.optimizer : optimizer.state_dict(),
                cp_names.trnlog : monitor_trn.get_state_dic(),
                cp_names.vallog : tmp_val_log,
                cp_names.lr_scheduler : lr_scheduler_state
            }
            torch.save(tmp_dic, tmp_model_name)
            if args.verbose == 1:
                nii_display.f_eprint(str(datetime.datetime.now()))
                nii_display.f_eprint("Save {:s}".format(tmp_model_name),
                                     flush=True)

        # Early stopping
        #  note: if LR scheduler is used, early stopping will be
        #  disabled
        if lr_scheduler.f_allow_early_stopping() and \
           monitor_val is not None and \
           monitor_val.should_early_stop(no_best_epoch_num):
            flag_early_stopped = True
            break

    # loop done
    nii_op_display_tk.print_log_tail()
    if flag_early_stopped:
        nii_display.f_print("Training finished by early stopping")
    else:
        nii_display.f_print("Training finished")
    nii_display.f_print("Model is saved to", end = '')
    nii_display.f_print("{}".format(nii_nn_tools.f_save_trained_name(args)))
    return


@staticmethod
    def decode_based_on_file_extension(content: bytes, file_path: str, default_bg: str = "#ffffff", default_fg: str = "#000000") -> 'AnsiArtDocument':
        """Creates a document from the given bytes, detecting the file format.

        Raises FormatReadNotSupported if the file format is not supported for reading. Some are write-only.
        Raises UnicodeDecodeError, which can be a very long message, so make sure to handle it!
        Raises UnidentifiedImageError if the format is not detected.
        """
        format_id = AnsiArtDocument.format_from_extension(file_path)
        # print("Supported image formats for reading:", Image.OPEN.keys())
        # TODO: try loading as image first, then as text if that fails with UnidentifiedImageError
        # That way it can handle images without file extensions.
        if format_id in Image.OPEN:
            return AnsiArtDocument.from_image_format(content)
        elif format_id == "ANSI":
            return AnsiArtDocument.from_ansi(content.decode('utf-8'), default_bg, default_fg)
        elif format_id == "IRC":
            return AnsiArtDocument.from_irc(content.decode('utf-8'), default_bg, default_fg)
        elif format_id == "PLAINTEXT":
            return AnsiArtDocument.from_plain(content.decode('utf-8'), default_bg, default_fg)
        elif format_id == "SVG":
            return AnsiArtDocument.from_svg(content.decode('utf-8'), default_bg, default_fg)
        elif format_id in Image.SAVE or format_id in ["HTML", "RICH_CONSOLE_MARKUP"]:
            # This is a write-only format.
            raise FormatReadNotSupported(localized_message=_("Cannot read files saved as %1 format.", format_id))
        else:
            # This is an unknown format.
            # For now at least, I'm preserving the behavior of loading as ANSI/PLAINTEXT.
            return AnsiArtDocument.from_text(content.decode('utf-8'), default_bg, default_fg)

@staticmethod
    def decode_based_on_file_extension(content: bytes, file_path: str, default_bg: str = "#ffffff", default_fg: str = "#000000") -> 'AnsiArtDocument':

        format_id = AnsiArtDocument.format_from_extension(file_path)
        # print("Supported image formats for reading:", Image.OPEN.keys())
        # TODO: try loading as image first, then as text if that fails with UnidentifiedImageError
        # That way it can handle images without file extensions.
        if format_id in Image.OPEN:
            return AnsiArtDocument.from_image_format(content)
        elif format_id == "ANSI":
            return AnsiArtDocument.from_ansi(content.decode('utf-8'), default_bg, default_fg)
        elif format_id == "IRC":
            return AnsiArtDocument.from_irc(content.decode('utf-8'), default_bg, default_fg)
        elif format_id == "PLAINTEXT":
            return AnsiArtDocument.from_plain(content.decode('utf-8'), default_bg, default_fg)
        elif format_id == "SVG":
            return AnsiArtDocument.from_svg(content.decode('utf-8'), default_bg, default_fg)
        elif format_id in Image.SAVE or format_id in ["HTML", "RICH_CONSOLE_MARKUP"]:
            # This is a write-only format.
            raise FormatReadNotSupported(localized_message=_("Cannot read files saved as %1 format.", format_id))
        else:
            # This is an unknown format.
            # For now at least, I'm preserving the behavior of loading as ANSI/PLAINTEXT.
            return AnsiArtDocument.from_text(content.decode('utf-8'), default_bg, default_fg)


async def update(self, db: AsyncSession, dept_id: int, obj_in: UpdateDeptParam) -> int:
        """
        更新部门

        :param db:
        :param dept_id:
        :param obj_in:
        :return:
        """
        return await self.update_model(db, dept_id, obj_in)

async def update(self, db: AsyncSession, dept_id: int, obj_in: UpdateDeptParam) -> int:

        return await self.update_model(db, dept_id, obj_in)


def enable(self, scene: Scene = None):
        """ Actually loads the learnable params into the scene nodes attr

        Args:
            scene (Scene, optional): An optional target scene to load the learnable params.
                If not provided, `self.scene` will be used. Defaults to None.
        """
        self.is_enabled = True
        scene = self.scene or scene
        if self.refine_ego_motion:
            for node_id, transform in self['ego_motion'].items():
                node = scene.all_nodes[node_id]
                # NOTE: Here, 'transform' of node.frame_data.subattr will be set to \
                #           the handle of the learnable 'transform' nn.Module here.
                #       This allows the learnable transform to be used with gradients during the rendering process,
                #       making the learnable param here part of the computation graph.
                node.frame_data.subattr['transform'] = transform

        if self.refine_other_motion:
            for node_id, transform in self['other_motion'].items():
                node = scene.all_nodes[node_id]
                node.frame_data.subattr['transform'] = transform

        if self.refine_camera_intr:
            pass

        if self.refine_camera_extr:
            pass

        if self.refine_sensor_ts:
            for node_id, ts in self['sensor_ts'].items():
                node = scene.all_nodes[node_id]
                node.frame_data.subattr['global_ts'] = ts

def enable(self, scene: Scene = None):

        self.is_enabled = True
        scene = self.scene or scene
        if self.refine_ego_motion:
            for node_id, transform in self['ego_motion'].items():
                node = scene.all_nodes[node_id]
                # NOTE: Here, 'transform' of node.frame_data.subattr will be set to \
                #           the handle of the learnable 'transform' nn.Module here.
                #       This allows the learnable transform to be used with gradients during the rendering process,
                #       making the learnable param here part of the computation graph.
                node.frame_data.subattr['transform'] = transform

        if self.refine_other_motion:
            for node_id, transform in self['other_motion'].items():
                node = scene.all_nodes[node_id]
                node.frame_data.subattr['transform'] = transform

        if self.refine_camera_intr:
            pass

        if self.refine_camera_extr:
            pass

        if self.refine_sensor_ts:
            for node_id, ts in self['sensor_ts'].items():
                node = scene.all_nodes[node_id]
                node.frame_data.subattr['global_ts'] = ts


@staticmethod
    def find_subject_predicate_object_with_chunk(sentence, chunk):
        """
        Extract subject, predicate, and object from a sentence, using a chunk for context.

        Args:
            sentence (str): The input sentence.
            chunk (str): A chunk of text providing context.

        Returns:
            tuple: A tuple containing (subject, predicate, object), each as a string or None if not found.

        Raises:
            ValueError: If the input sentence or chunk is not a string or is empty.
            Exception: For any other unexpected errors during processing.
        """
        if not isinstance(sentence, str) or not sentence.strip():
            raise ValueError("Input sentence must be a non-empty string")
        if not isinstance(chunk, str) or not chunk.strip():
            raise ValueError("Input chunk must be a non-empty string")

        try:
            doc = nlp(chunk)  # Process the chunk for context
            sentence_doc = nlp(sentence)  # Process the sentence separately

            # Identify named entities using SpaCy NER
            entities = [ent for ent in doc.ents if ent.label_ in ("PERSON", "ORG")]

            subject, predicate, _object = None, None, None

            # Use syntactic dependency labels to identify subject and verb
            for token in sentence_doc:
                if token.dep_ == "nsubj" or token.dep_ == "nsubjpass":
                    subject = token

                    # Filter irrelevant words based on POS tags and additional stop words
                    if subject and isinstance(subject, spacy.tokens.Doc):
                        filtered_subject_words = [
                            word.text
                            for word in subject.words
                            if word.pos_ not in ["STOP", "ADP", "DET", "AUX"]  # Add AUX for auxiliary verbs
                        ]
                    else:
                        filtered_subject_words = [subject.text] if subject else None

                    # Join the filtered words with a space
                    subject_text = " ".join(filtered_subject_words) if filtered_subject_words else None
                    print(f"\nDEBUG CHUNK: \nFiltered subject words: {filtered_subject_words}\nSubject text: {subject_text}\n")

                elif token.pos_ == "VERB":
                    # Check if it's part of a verb phrase indicating the predicate
                    if token.dep_ == "aux" and nlp(token.head.text).pos_ == "VERB":
                        continue  # Skip auxiliary verbs
                    else:
                        predicate = token.head  # Consider the head of the verb phrase as the predicate
                        break

            # If subject not found directly, explore other possibilities
            if not subject:
                if predicate:
                    # Check for subject within relative clauses or previous entities
                    for child in predicate.children:
                        if child.dep_ == "relcl":
                            subject = TripleExtract.find_subject_in_clause_with_chunk(child,
                                                                                entities.copy())  # Pass a copy of entities
                            if subject:
                                break
                        elif child.dep_ == "pobj" and len(entities) > 0:
                            # Check if object from previous sentence is the subject
                            for entity in entities:
                                if entity.text == child.text:
                                    subject = entity
                                    break
                        else:
                            # Try finding a verb phrase as the subject
                            for chunk in doc.noun_chunks:
                                if any(token.pos_ == "VERB" for token in chunk):
                                    subject = chunk
                                    break

            # Look for candidate objects after finding subject and predicate
            if subject and predicate:
                for child in predicate.children:
                    if child.dep_ in ["dobj", "attr", "iobj"]:
                        for grandchild in child.children:  # Iterate over child.children
                            if grandchild.dep_ == "pobj":
                                _object = grandchild
                                break  # Stop iterating after finding the object
                    elif child.dep_ == "prep":
                        for grandchild in child.children:
                            if grandchild.dep_ == "pobj":
                                _object = grandchild
                                break  # Stop iterating after finding the object

            # Convert identified tokens to text if they exist
            subject_text = subject.text if subject else None
            predicate_text = predicate.lemma_ if predicate else None  # Using lemma for base form of verb
            object_text = _object.text if _object else None

            print(
                f"usingContext:\nSubject: {subject_text}\nPredicate: {predicate_text}\nObject: {object_text}")
            return subject_text, predicate_text, object_text

        except Exception as e:
            raise Exception(f"An error occurred while processing the sentence with chunk: {str(e)}")

@staticmethod
    def find_subject_predicate_object_with_chunk(sentence, chunk):

        if not isinstance(sentence, str) or not sentence.strip():
            raise ValueError("Input sentence must be a non-empty string")
        if not isinstance(chunk, str) or not chunk.strip():
            raise ValueError("Input chunk must be a non-empty string")

        try:
            doc = nlp(chunk)  # Process the chunk for context
            sentence_doc = nlp(sentence)  # Process the sentence separately

            # Identify named entities using SpaCy NER
            entities = [ent for ent in doc.ents if ent.label_ in ("PERSON", "ORG")]

            subject, predicate, _object = None, None, None

            # Use syntactic dependency labels to identify subject and verb
            for token in sentence_doc:
                if token.dep_ == "nsubj" or token.dep_ == "nsubjpass":
                    subject = token

                    # Filter irrelevant words based on POS tags and additional stop words
                    if subject and isinstance(subject, spacy.tokens.Doc):
                        filtered_subject_words = [
                            word.text
                            for word in subject.words
                            if word.pos_ not in ["STOP", "ADP", "DET", "AUX"]  # Add AUX for auxiliary verbs
                        ]
                    else:
                        filtered_subject_words = [subject.text] if subject else None

                    # Join the filtered words with a space
                    subject_text = " ".join(filtered_subject_words) if filtered_subject_words else None
                    print(f"\nDEBUG CHUNK: \nFiltered subject words: {filtered_subject_words}\nSubject text: {subject_text}\n")

                elif token.pos_ == "VERB":
                    # Check if it's part of a verb phrase indicating the predicate
                    if token.dep_ == "aux" and nlp(token.head.text).pos_ == "VERB":
                        continue  # Skip auxiliary verbs
                    else:
                        predicate = token.head  # Consider the head of the verb phrase as the predicate
                        break

            # If subject not found directly, explore other possibilities
            if not subject:
                if predicate:
                    # Check for subject within relative clauses or previous entities
                    for child in predicate.children:
                        if child.dep_ == "relcl":
                            subject = TripleExtract.find_subject_in_clause_with_chunk(child,
                                                                                entities.copy())  # Pass a copy of entities
                            if subject:
                                break
                        elif child.dep_ == "pobj" and len(entities) > 0:
                            # Check if object from previous sentence is the subject
                            for entity in entities:
                                if entity.text == child.text:
                                    subject = entity
                                    break
                        else:
                            # Try finding a verb phrase as the subject
                            for chunk in doc.noun_chunks:
                                if any(token.pos_ == "VERB" for token in chunk):
                                    subject = chunk
                                    break

            # Look for candidate objects after finding subject and predicate
            if subject and predicate:
                for child in predicate.children:
                    if child.dep_ in ["dobj", "attr", "iobj"]:
                        for grandchild in child.children:  # Iterate over child.children
                            if grandchild.dep_ == "pobj":
                                _object = grandchild
                                break  # Stop iterating after finding the object
                    elif child.dep_ == "prep":
                        for grandchild in child.children:
                            if grandchild.dep_ == "pobj":
                                _object = grandchild
                                break  # Stop iterating after finding the object

            # Convert identified tokens to text if they exist
            subject_text = subject.text if subject else None
            predicate_text = predicate.lemma_ if predicate else None  # Using lemma for base form of verb
            object_text = _object.text if _object else None

            print(
                f"usingContext:\nSubject: {subject_text}\nPredicate: {predicate_text}\nObject: {object_text}")
            return subject_text, predicate_text, object_text

        except Exception as e:
            raise Exception(f"An error occurred while processing the sentence with chunk: {str(e)}")


def test_complete_chat_sse_kwargs(self):
        """
        Tests that the StreamingClaudeWrapper can be used to perform streaming chat completion with kwargs.
        """
        fixture = StreamingClaudeWrapper(anthropic_api_key, model="claude-v1", format_sse=True, append_stop_token=False,
                                         temperature=0.9, top_k=2)

        test_streaming_complete_chat_sse(self, fixture, check_stop=False, verbose=False)

def test_complete_chat_sse_kwargs(self):

        fixture = StreamingClaudeWrapper(anthropic_api_key, model="claude-v1", format_sse=True, append_stop_token=False,
                                         temperature=0.9, top_k=2)

        test_streaming_complete_chat_sse(self, fixture, check_stop=False, verbose=False)


def get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor:
    """
    Get relative positional embeddings according to the relative positions of
        query and key sizes.
    Args:
        q_size (int): size of query q.
        k_size (int): size of key k.
        rel_pos (Tensor): relative position embeddings (L, C).

    Returns:
        Extracted positional embeddings according to relative positions.
    """
    max_rel_dist = int(2 * max(q_size, k_size) - 1)
    # Interpolate rel pos if needed.
    if rel_pos.shape[0] != max_rel_dist:
        # Interpolate rel pos.
        rel_pos_resized = F.interpolate(
            rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1),
            size=max_rel_dist,
            mode="linear",
        )
        rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0)
    else:
        rel_pos_resized = rel_pos

    # Scale the coords with short length if shapes for q and k are different.
    q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0)
    k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0)
    relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0)

    return rel_pos_resized[relative_coords.long()]

def get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor:

    max_rel_dist = int(2 * max(q_size, k_size) - 1)
    # Interpolate rel pos if needed.
    if rel_pos.shape[0] != max_rel_dist:
        # Interpolate rel pos.
        rel_pos_resized = F.interpolate(
            rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1),
            size=max_rel_dist,
            mode="linear",
        )
        rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0)
    else:
        rel_pos_resized = rel_pos

    # Scale the coords with short length if shapes for q and k are different.
    q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0)
    k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0)
    relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0)

    return rel_pos_resized[relative_coords.long()]


def join_path(self, filepath: Union[str, Path],
                  *filepaths: Union[str, Path]) -> str:
        """Concatenate all file paths.

        Join one or more filepath components intelligently. The return value
        is the concatenation of filepath and any members of *filepaths.

        Args:
            filepath (str or Path): Path to be concatenated.

        Returns:
            str: The result of concatenation.
        """
        return self.client.join_path(filepath, *filepaths)

def join_path(self, filepath: Union[str, Path],
                  *filepaths: Union[str, Path]) -> str:

        return self.client.join_path(filepath, *filepaths)


def normalize(x, mean=IMAGENET_MEAN, std=IMAGENET_STD, inplace=False):
    """Denormalize RGB images x per ImageNet stats in BCHW format, i.e. = (x - mean) / std."""
    return TF.normalize(x, mean, std, inplace=inplace)

def normalize(x, mean=IMAGENET_MEAN, std=IMAGENET_STD, inplace=False):

    return TF.normalize(x, mean, std, inplace=inplace)


@partial(
	jax.custom_vjp,
	nondiff_argnums=[6, 7, 8, 9, 10, 11],
)
def ring_attention_tpu(
	query: chex.Array,
	key: chex.Array,
	value: chex.Array,
	bias: tp.Optional[chex.Array] = None,
	segment_ids: tp.Optional[SegmentIds] = None,
	cache_idx: tp.Optional[int] = None,
	axis_name: tp.Optional[str] = None,
	float32_logits: bool = True,
	softmax_scale: tp.Optional[float] = None,
	blocksize_q: int = 256,
	blocksize_k: int = 256,
	blocksize_c: tp.Optional[int] = None,
) -> chex.Array:
	"""Computes ring attention using FlashAttention on TPU.

	Args:
	    query: Query array of shape (batch, query_len, num_heads, dim_per_head).
	    key: Key array of shape (batch, kv_len, num_heads, dim_per_head).
	    value: Value array of shape (batch, kv_len, num_heads, dim_per_head).
	    bias: tp.Optional bias array.  Its shape depends on the attention mechanism.
	    segment_ids: tp.Optional segment ids for Q and KV sequences.
	    cache_idx: tp.Optional cache index for use with caching.
	    axis_name: tp.Optional name of the axis to ppermute over (for multi-host support).
	    float32_logits: Whether to compute logits in float32.
	    softmax_scale: tp.Optional scaling factor for the softmax function.
	    blocksize_q: Block size for the query sequence.
	    blocksize_k: Block size for the key/value sequence.
	    blocksize_c: tp.Optional block size for causal masking.

	Returns:
	    Output array of shape (batch, query_len, num_heads, dim_per_head).
	"""
	y, _ = _ring_flash_attention_fwd_tpu(
		query,
		key,
		value,
		bias,
		segment_ids,
		cache_idx,
		axis_name,
		float32_logits,
		softmax_scale,
		blocksize_q,
		blocksize_k,
		blocksize_c,
	)
	return y

@partial(
	jax.custom_vjp,
	nondiff_argnums=[6, 7, 8, 9, 10, 11],
)
def ring_attention_tpu(
	query: chex.Array,
	key: chex.Array,
	value: chex.Array,
	bias: tp.Optional[chex.Array] = None,
	segment_ids: tp.Optional[SegmentIds] = None,
	cache_idx: tp.Optional[int] = None,
	axis_name: tp.Optional[str] = None,
	float32_logits: bool = True,
	softmax_scale: tp.Optional[float] = None,
	blocksize_q: int = 256,
	blocksize_k: int = 256,
	blocksize_c: tp.Optional[int] = None,
) -> chex.Array:

	y, _ = _ring_flash_attention_fwd_tpu(
		query,
		key,
		value,
		bias,
		segment_ids,
		cache_idx,
		axis_name,
		float32_logits,
		softmax_scale,
		blocksize_q,
		blocksize_k,
		blocksize_c,
	)
	return y


def pack_sequences(
	dataset: Dataset,
	max_length: int = 512,
	pad_token_id: int = 0,
	reset_position_ids: bool = False,
	num_proc: tp.Optional[int] = None,
):
	"""
	Pack sequences together with their attention masks and position IDs

	# With continuous position IDs
	packed_dataset = pack_sequences(
			dataset,
			max_length=512,
			pad_token_id=0,
			reset_position_ids=False
	)

	# With reset position IDs for each sequence
	packed_dataset = pack_sequences(
			dataset,
			max_length=512,
			pad_token_id=0,
			reset_position_ids=True
	)

	# Example output format for a packed sequence with two sequences:
	# reset_position_ids=False:
	{
			'input_ids': [seq1_tokens + [PAD] + seq2_tokens + [PAD] + padding],
			'attention_mask': [1,1,1,0,1,1,1,0,0,0],
			'position_ids': [0,1,2,3,4,5,6,7,0,0]
	}

	# reset_position_ids=True:
	{
			'input_ids': [seq1_tokens + [PAD] + seq2_tokens + [PAD] + padding],
			'attention_mask': [1,1,1,0,1,1,1,0,0,0],
			'position_ids': [0,1,2,0,0,1,2,0,0,0]
	}

	Args:
	    dataset: Dataset containing 'input_ids' and 'attention_mask'
	    max_length: Maximum length of packed sequence
	    pad_token_id: Token ID used for padding
	    reset_position_ids: If True, reset position IDs for each sequence in the pack

	Returns:
	    packed_dataset: Dataset with packed sequences, attention masks, and position IDs
	"""

	def pack_examples(examples):
		current_packed_input_ids = []
		current_packed_attention_mask = []
		current_packed_position_ids = []
		current_length = 0

		packed_input_ids = []
		packed_attention_mask = []
		packed_position_ids = []

		def get_position_ids(length, start_position=0):
			if reset_position_ids:
				return list(range(length))
			else:
				return list(range(start_position, start_position + length))

		# Iterate through all examples
		for input_ids, attention_mask in zip(
			examples["input_ids"], examples["attention_mask"]
		):
			seq_length = len(input_ids)

			# If adding this sequence would exceed max_length, start a new packed sequence
			if current_length + seq_length + 1 > max_length:
				# Pad the current packed sequence if needed
				if current_length < max_length:
					padding_length = max_length - current_length
					current_packed_input_ids.extend([pad_token_id] * padding_length)
					current_packed_attention_mask.extend([0] * padding_length)
					current_packed_position_ids.extend([0] * padding_length)

				# Add the completed packed sequence to results
				packed_input_ids.append(current_packed_input_ids)
				packed_attention_mask.append(current_packed_attention_mask)
				packed_position_ids.append(current_packed_position_ids)

				# Start new packed sequence
				current_packed_input_ids = []
				current_packed_attention_mask = []
				current_packed_position_ids = []
				current_length = 0

			# Generate position IDs for current sequence
			position_ids = get_position_ids(seq_length, start_position=current_length)

			# Add current sequence
			current_packed_input_ids.extend(input_ids)
			current_packed_attention_mask.extend(attention_mask)
			current_packed_position_ids.extend(position_ids)

			# Add separator token
			current_packed_input_ids.append(pad_token_id)
			current_packed_attention_mask.append(0)
			current_packed_position_ids.append(
				position_ids[-1] + 1 if not reset_position_ids else 0
			)

			current_length += seq_length + 1

		# Handle the last packed sequence
		if current_packed_input_ids:
			# Pad if needed
			if current_length < max_length:
				padding_length = max_length - current_length
				current_packed_input_ids.extend([pad_token_id] * padding_length)
				current_packed_attention_mask.extend([0] * padding_length)
				current_packed_position_ids.extend([0] * padding_length)

			packed_input_ids.append(current_packed_input_ids)
			packed_attention_mask.append(current_packed_attention_mask)
			packed_position_ids.append(current_packed_position_ids)

		return {
			"input_ids": packed_input_ids,
			"attention_mask": packed_attention_mask,
			"position_ids": packed_position_ids,
		}

	# Process the dataset in batches
	packed_dataset = dataset.map(
		pack_examples,
		batched=True,
		remove_columns=dataset.column_names,
		desc="Packing sequences",
		num_proc=num_proc,
	)

	return packed_dataset

def pack_sequences(
	dataset: Dataset,
	max_length: int = 512,
	pad_token_id: int = 0,
	reset_position_ids: bool = False,
	num_proc: tp.Optional[int] = None,
):

	def pack_examples(examples):
		current_packed_input_ids = []
		current_packed_attention_mask = []
		current_packed_position_ids = []
		current_length = 0

		packed_input_ids = []
		packed_attention_mask = []
		packed_position_ids = []

		def get_position_ids(length, start_position=0):
			if reset_position_ids:
				return list(range(length))
			else:
				return list(range(start_position, start_position + length))

		# Iterate through all examples
		for input_ids, attention_mask in zip(
			examples["input_ids"], examples["attention_mask"]
		):
			seq_length = len(input_ids)

			# If adding this sequence would exceed max_length, start a new packed sequence
			if current_length + seq_length + 1 > max_length:
				# Pad the current packed sequence if needed
				if current_length < max_length:
					padding_length = max_length - current_length
					current_packed_input_ids.extend([pad_token_id] * padding_length)
					current_packed_attention_mask.extend([0] * padding_length)
					current_packed_position_ids.extend([0] * padding_length)

				# Add the completed packed sequence to results
				packed_input_ids.append(current_packed_input_ids)
				packed_attention_mask.append(current_packed_attention_mask)
				packed_position_ids.append(current_packed_position_ids)

				# Start new packed sequence
				current_packed_input_ids = []
				current_packed_attention_mask = []
				current_packed_position_ids = []
				current_length = 0

			# Generate position IDs for current sequence
			position_ids = get_position_ids(seq_length, start_position=current_length)

			# Add current sequence
			current_packed_input_ids.extend(input_ids)
			current_packed_attention_mask.extend(attention_mask)
			current_packed_position_ids.extend(position_ids)

			# Add separator token
			current_packed_input_ids.append(pad_token_id)
			current_packed_attention_mask.append(0)
			current_packed_position_ids.append(
				position_ids[-1] + 1 if not reset_position_ids else 0
			)

			current_length += seq_length + 1

		# Handle the last packed sequence
		if current_packed_input_ids:
			# Pad if needed
			if current_length < max_length:
				padding_length = max_length - current_length
				current_packed_input_ids.extend([pad_token_id] * padding_length)
				current_packed_attention_mask.extend([0] * padding_length)
				current_packed_position_ids.extend([0] * padding_length)

			packed_input_ids.append(current_packed_input_ids)
			packed_attention_mask.append(current_packed_attention_mask)
			packed_position_ids.append(current_packed_position_ids)

		return {
			"input_ids": packed_input_ids,
			"attention_mask": packed_attention_mask,
			"position_ids": packed_position_ids,
		}

	# Process the dataset in batches
	packed_dataset = dataset.map(
		pack_examples,
		batched=True,
		remove_columns=dataset.column_names,
		desc="Packing sequences",
		num_proc=num_proc,
	)

	return packed_dataset


def solve_heun(self, x, t_span, mu, mask, spks, cond, training=False, guidance_scale=0.0):
        """
        Fixed heun solver for ODEs.
        Args:
            x (torch.Tensor): random noise
            t_span (torch.Tensor): n_timesteps interpolated
                shape: (n_timesteps + 1,)
            mu (torch.Tensor): output of encoder
                shape: (batch_size, n_feats, mel_timesteps)
            mask (torch.Tensor): output_mask
                shape: (batch_size, 1, mel_timesteps)
            cond: Not used but kept for future purposes
        """
        t, _, dt = t_span[0], t_span[-1], t_span[1] - t_span[0]

        #-! : reserved space for debugger
        sol = []
        steps = 1

        while steps <= len(t_span) - 1:
            dphi_dt = self.func_dphi_dt(x, mask, mu, t, spks, cond, training=training, guidance_scale=guidance_scale)
            dphi_dt_2 = self.func_dphi_dt(x + dt * dphi_dt, mask, mu, t+dt, spks, cond, training=training, guidance_scale=guidance_scale)

            #- Euler's -> Y'n+1' = Y'n' + h * F(X'n', Y'n')
            # x = x + dt * dphi_dt

            #- Heun's -> Y'n+1' = Y'n' + h * 0.5( F(X'n', Y'n') + F(X'n' + h, Y'n' + h * F(X'n', Y'n') ) )
            x = x + dt * 0.5 * (dphi_dt + dphi_dt_2)
            t = t + dt

            sol.append(x)
            if steps < len(t_span) - 1:
                dt = t_span[steps + 1] - t
            steps += 1

        return sol[-1]

def solve_heun(self, x, t_span, mu, mask, spks, cond, training=False, guidance_scale=0.0):

        t, _, dt = t_span[0], t_span[-1], t_span[1] - t_span[0]

        #-! : reserved space for debugger
        sol = []
        steps = 1

        while steps <= len(t_span) - 1:
            dphi_dt = self.func_dphi_dt(x, mask, mu, t, spks, cond, training=training, guidance_scale=guidance_scale)
            dphi_dt_2 = self.func_dphi_dt(x + dt * dphi_dt, mask, mu, t+dt, spks, cond, training=training, guidance_scale=guidance_scale)

            #- Euler's -> Y'n+1' = Y'n' + h * F(X'n', Y'n')
            # x = x + dt * dphi_dt

            #- Heun's -> Y'n+1' = Y'n' + h * 0.5( F(X'n', Y'n') + F(X'n' + h, Y'n' + h * F(X'n', Y'n') ) )
            x = x + dt * 0.5 * (dphi_dt + dphi_dt_2)
            t = t + dt

            sol.append(x)
            if steps < len(t_span) - 1:
                dt = t_span[steps + 1] - t
            steps += 1

        return sol[-1]


def reset(self) -> None:
        """
        Clear buffer and reset internal counters.
        """

        self._buffer.clear()
        self._has_item.clear()
        self._prefill = self.prefill
        self._last_tx = self._last_rx = self._generation = self._generation_ts = 0

def reset(self) -> None:

        self._buffer.clear()
        self._has_item.clear()
        self._prefill = self.prefill
        self._last_tx = self._last_rx = self._generation = self._generation_ts = 0


def has_data_stream(arg_type: Type) -> bool:
    """Recursive check for a DataStream container in a type annotation"""
    if _is_data_stream(arg_type):
        return True

    typing_args = get_args(arg_type)
    if len(typing_args) > 0:
        for typ in typing_args:
            if has_data_stream(typ):
                return True

    return False

def has_data_stream(arg_type: Type) -> bool:

    if _is_data_stream(arg_type):
        return True

    typing_args = get_args(arg_type)
    if len(typing_args) > 0:
        for typ in typing_args:
            if has_data_stream(typ):
                return True

    return False


@property
    def app_name(self):
        """Return the app name."""
        return self._app_name

@property
    def app_name(self):

        return self._app_name


@classmethod
    def langchain2doc(cls, document):
        """Transform Langchain to Document format."""
        metadata = document.metadata or {}
        return cls(content=document.page_content, metadata=metadata)

@classmethod
    def langchain2doc(cls, document):

        metadata = document.metadata or {}
        return cls(content=document.page_content, metadata=metadata)


def get_table_summary(self, table_name):
        """Get table summary for table.

        example:
            table_name(column1(column1 comment),column2(column2 comment),
            column3(column3 comment) and index keys, and table comment: {table_comment})
        """
        return _parse_table_summary(self.db, self.summary_template, table_name)

def get_table_summary(self, table_name):

        return _parse_table_summary(self.db, self.summary_template, table_name)


def __init__(self):
        """Client for vis protocol."""
        self._vis_tag: Dict[str, Vis] = {}

def __init__(self):

        self._vis_tag: Dict[str, Vis] = {}


async def bin_search(
        self, blocks: List[str], model_nam: str, max_new_token: int
    ) -> int:
        """Binary search to find the split point."""
        l, r = 0, len(blocks) - 1
        while l < r:
            mid = l + r + 1 >> 1
            current_tokens = await self._llm_client.count_token(
                model_nam, "".join(blocks[: mid + 1])
            )
            if current_tokens <= max_new_token:
                l = mid
            else:
                r = mid - 1
        return r

async def bin_search(
        self, blocks: List[str], model_nam: str, max_new_token: int
    ) -> int:

        l, r = 0, len(blocks) - 1
        while l < r:
            mid = l + r + 1 >> 1
            current_tokens = await self._llm_client.count_token(
                model_nam, "".join(blocks[: mid + 1])
            )
            if current_tokens <= max_new_token:
                l = mid
            else:
                r = mid - 1
        return r


def stop_recording(self) -> None:
        """Stop recording."""
        Thread(target=stop_record).start()

def stop_recording(self) -> None:

        Thread(target=stop_record).start()


def get_detection_dataset_dicts(
    names,
    filter_empty=True,
    min_keypoints=0,
    proposal_files=None,
    check_consistency=True,
):
    """
    Load and prepare dataset dicts for instance detection/segmentation and semantic segmentation.

    Args:
        names (str or list[str]): a dataset name or a list of dataset names
        filter_empty (bool): whether to filter out images without instance annotations
        min_keypoints (int): filter out images with fewer keypoints than
            `min_keypoints`. Set to 0 to do nothing.
        proposal_files (list[str]): if given, a list of object proposal files
            that match each dataset in `names`.
        check_consistency (bool): whether to check if datasets have consistent metadata.

    Returns:
        list[dict]: a list of dicts following the standard dataset dict format.
    """
    if isinstance(names, str):
        names = [names]
    assert len(names), names
    dataset_dicts = [DatasetCatalog.get(dataset_name) for dataset_name in names]
    for dataset_name, dicts in zip(names, dataset_dicts):
        assert len(dicts), "Dataset '{}' is empty!".format(dataset_name)

    if proposal_files is not None:
        assert len(names) == len(proposal_files)
        # load precomputed proposals from proposal files
        dataset_dicts = [
            load_proposals_into_dataset(dataset_i_dicts, proposal_file)
            for dataset_i_dicts, proposal_file in zip(dataset_dicts, proposal_files)
        ]

    if isinstance(dataset_dicts[0], torchdata.Dataset):
        return torchdata.ConcatDataset(dataset_dicts)

    dataset_dicts = list(itertools.chain.from_iterable(dataset_dicts))

    has_instances = "annotations" in dataset_dicts[0]
    if filter_empty and has_instances:
        dataset_dicts = filter_images_with_only_crowd_annotations(dataset_dicts)
    if min_keypoints > 0 and has_instances:
        dataset_dicts = filter_images_with_few_keypoints(dataset_dicts, min_keypoints)

    if check_consistency and has_instances:
        try:
            class_names = MetadataCatalog.get(names[0]).thing_classes
            check_metadata_consistency("thing_classes", names)
            print_instances_class_histogram(dataset_dicts, class_names)
        except AttributeError:  # class names are not available for this dataset
            pass

    assert len(dataset_dicts), "No valid data found in {}.".format(",".join(names))
    return dataset_dicts

def get_detection_dataset_dicts(
    names,
    filter_empty=True,
    min_keypoints=0,
    proposal_files=None,
    check_consistency=True,
):

    if isinstance(names, str):
        names = [names]
    assert len(names), names
    dataset_dicts = [DatasetCatalog.get(dataset_name) for dataset_name in names]
    for dataset_name, dicts in zip(names, dataset_dicts):
        assert len(dicts), "Dataset '{}' is empty!".format(dataset_name)

    if proposal_files is not None:
        assert len(names) == len(proposal_files)
        # load precomputed proposals from proposal files
        dataset_dicts = [
            load_proposals_into_dataset(dataset_i_dicts, proposal_file)
            for dataset_i_dicts, proposal_file in zip(dataset_dicts, proposal_files)
        ]

    if isinstance(dataset_dicts[0], torchdata.Dataset):
        return torchdata.ConcatDataset(dataset_dicts)

    dataset_dicts = list(itertools.chain.from_iterable(dataset_dicts))

    has_instances = "annotations" in dataset_dicts[0]
    if filter_empty and has_instances:
        dataset_dicts = filter_images_with_only_crowd_annotations(dataset_dicts)
    if min_keypoints > 0 and has_instances:
        dataset_dicts = filter_images_with_few_keypoints(dataset_dicts, min_keypoints)

    if check_consistency and has_instances:
        try:
            class_names = MetadataCatalog.get(names[0]).thing_classes
            check_metadata_consistency("thing_classes", names)
            print_instances_class_histogram(dataset_dicts, class_names)
        except AttributeError:  # class names are not available for this dataset
            pass

    assert len(dataset_dicts), "No valid data found in {}.".format(",".join(names))
    return dataset_dicts


def number_of_rotatable_bonds(mol: Mol) -> int:
    """
    Number of rotatable bonds.

    Calculates the total number of rotatable bonds in the molecule.

    Parameters
    ----------
    mol : RDKit ``Mol`` object
        The molecule for which the number of rotatable bonds is to be calculated.

    Examples
    --------
    >>> from rdkit.Chem import MolFromSmiles
    >>> from skfp.descriptors import number_of_rotatable_bonds
    >>> mol = MolFromSmiles("C=CC=C")  # Butadiene
    >>> number_of_rotatable_bonds(mol)
    1
    """
    return CalcNumRotatableBonds(mol)

def number_of_rotatable_bonds(mol: Mol) -> int:

    return CalcNumRotatableBonds(mol)


def _reg_loss(self, regr, gt_regr, mask):
        """L1 regression loss
        Arguments:
        regr (batch x max_objects x dim)
        gt_regr (batch x max_objects x dim)
        mask (batch x max_objects)
        """
        num = mask.float().sum()
        mask = mask.unsqueeze(2).expand_as(gt_regr).clone().float()
        isnotnan = (~torch.isnan(gt_regr)).float()
        mask *= isnotnan
        regr = regr * mask
        gt_regr = gt_regr * mask

        loss = torch.abs(regr - gt_regr)
        loss = loss.transpose(2, 0).contiguous()

        loss = torch.sum(loss, dim=2)
        loss = torch.sum(loss, dim=1)

        num = reduce_mean(num)
        loss = loss / (num + 1e-4)
        return loss

def _reg_loss(self, regr, gt_regr, mask):

        num = mask.float().sum()
        mask = mask.unsqueeze(2).expand_as(gt_regr).clone().float()
        isnotnan = (~torch.isnan(gt_regr)).float()
        mask *= isnotnan
        regr = regr * mask
        gt_regr = gt_regr * mask

        loss = torch.abs(regr - gt_regr)
        loss = loss.transpose(2, 0).contiguous()

        loss = torch.sum(loss, dim=2)
        loss = torch.sum(loss, dim=1)

        num = reduce_mean(num)
        loss = loss / (num + 1e-4)
        return loss


def run_train(
    cfg: DefaultConfigProblemBase,
    model: torch.nn.Module,
    optimizer,
    scheduler,
    epoch_steps,
    train_dataloader,
    val_dataloader,
    val_df: pd.DataFrame,
):
    """Runs the training loop.

    Args:
        cfg: DefaultConfigProblemBase config object
        model: model
        train_dataloader: custom training Dataloader
        train_df: train DataFrame
        val_dataloader: custom validation Dataloader
        val_df: validation DataFrame

    Returns:
        Validation prediction output
        Validation loss
        Validation metric
        Last train batch
    """
    if (
        hasattr(cfg.augmentation, "neftune_noise_alpha")
        and cfg.augmentation.neftune_noise_alpha > 0
    ):
        activate_neftune(model, cfg.augmentation.neftune_noise_alpha)

    scaler: GradScaler | None = None
    if cfg.environment.mixed_precision:
        scaler = GradScaler(
            enabled=(cfg.environment.mixed_precision_dtype == "float16")
        )

    optimizer.zero_grad(set_to_none=True)

    # Prepare NLP Augmentation
    nlp_augment = None
    if hasattr(cfg.augmentation, "nlp_augmentations_class"):
        nlp_augment = cfg.augmentation.nlp_augmentations_class(cfg=cfg)

    start_epoch = 0

    _, metric_mode, _ = cfg.prediction.metric_class.get(cfg.prediction.metric)
    objective_op: Callable[[float, float], bool]
    if metric_mode == "max":
        best_val_metric = -np.inf
        objective_op = np.greater
    else:
        best_val_metric = np.inf
        objective_op = np.less

    if cfg.training.evaluate_before_training:
        val_loss, val_metric = run_eval(
            cfg=cfg, model=model, val_dataloader=val_dataloader, val_df=val_df
        )

    for epoch in range(start_epoch, cfg.training.epochs):
        set_seed(
            cfg.environment._seed
            + epoch * cfg.environment._world_size * cfg.environment.number_of_workers
            + cfg.environment._local_rank * cfg.environment.number_of_workers
        )
        logger.info(f"Training Epoch: {epoch + 1} / {cfg.training.epochs}")

        if (
            cfg.environment._distributed
            and not cfg.environment.use_deepspeed
            and hasattr(train_dataloader.sampler, "set_epoch")
        ):
            train_dataloader.sampler.set_epoch(epoch)  # type: ignore

        tqdm_out = TqdmToLogger(logger, level=logging.INFO)
        progress_bar = tqdm(
            total=epoch_steps,
            disable=cfg.environment._local_rank != 0,
            file=tqdm_out,
            ascii=True,
            desc="train loss",
            mininterval=0,
        )
        tr_it = iter(train_dataloader)

        losses = []
        model.train()

        log_update_steps = max(epoch_steps // 20, 1)
        evaluation_step = max(int(epoch_steps * cfg.training.evaluation_epochs), 1)
        logger.info(f"Evaluation step: {evaluation_step}")

        for itr, data in enumerate(tr_it):
            cfg.environment._curr_step += (
                cfg.training.batch_size * cfg.environment._world_size
            )

            # Batch to device
            batch = cfg.dataset.dataset_class.batch_to_device(
                data, cfg.environment._device
            )

            # NLP augmentation
            if nlp_augment is not None:
                batch = nlp_augment(batch)

            # Plot first batch
            if epoch == 0 and itr == 0 and cfg.environment._local_rank == 0:
                plot = cfg.logging.plots_class.plot_batch(batch=batch, cfg=cfg)
                log_plot(cfg, plot, "train_data")

            # only need to sync gradients at last step of grad accumulation
            model.require_backward_grad_sync = itr % cfg.training.grad_accumulation == 0

            # Forward pass
            with autocast(
                enabled=cfg.environment.mixed_precision,
                dtype=get_torch_dtype(cfg.environment.mixed_precision_dtype),
            ):
                output_dict = model.forward(batch)

            loss = output_dict["loss"]
            if ~np.isfinite(loss.item()) and (epoch > start_epoch or itr > 20):
                raise LLMTrainingException(
                    "NaN caught in loss during training. "
                    "Please, reduce learning rate, change dtype, "
                    "or disable mixed precision. Alternatively, "
                    "gradient clipping may help to stabilize training."
                )
            losses.append(loss.item())

            # loss is a mean loss per batch/sample
            # as grad_accumulations sums up the gradients, this loss must be scaled
            # by the number of grad_accumulations, to have similar behavior for
            # BS * grad_accumulations = const.
            if cfg.training.grad_accumulation != 1:
                loss = loss / cfg.training.grad_accumulation

            # Backward pass
            if (
                cfg.environment.mixed_precision
                and len(cfg.environment.gpus)
                and not cfg.environment.use_deepspeed
            ):
                scaler.scale(loss).backward()  # type: ignore
                if itr % cfg.training.grad_accumulation == 0:
                    if cfg.training.gradient_clip > 0:
                        scaler.unscale_(optimizer)  # type: ignore
                        torch.nn.utils.clip_grad_norm_(
                            model.parameters(), cfg.training.gradient_clip
                        )
                    scaler.step(optimizer)  # type: ignore
                    scaler.update()
                    optimizer.zero_grad(set_to_none=True)
            else:
                if cfg.environment.use_deepspeed:
                    model.backward(loss)  # type: ignore[operator]
                else:
                    loss.backward()
                if itr % cfg.training.grad_accumulation == 0:
                    if cfg.training.gradient_clip > 0:
                        torch.nn.utils.clip_grad_norm_(
                            model.parameters(), cfg.training.gradient_clip
                        )
                    optimizer.step()
                    optimizer.zero_grad(set_to_none=True)

            if cfg.environment._distributed:
                torch.cuda.synchronize(device=cfg.environment._local_rank)

            if scheduler is not None:
                scheduler.step()

            if cfg.environment._local_rank == 0:
                cfg.logging._logger.log(
                    "train",
                    "loss",
                    losses[-1],
                    step=cfg.environment._curr_step
                    / cfg.environment._step_log_denominator,
                )
                cfg.logging._logger.log(
                    "meta",
                    "lr",
                    optimizer.param_groups[0]["lr"],
                    step=cfg.environment._curr_step
                    / cfg.environment._step_log_denominator,
                )
                if cfg.training.differential_learning_rate_layers:
                    cfg.logging._logger.log(
                        "meta",
                        "lr_diff",
                        optimizer.param_groups[2]["lr"],
                        step=cfg.environment._curr_step
                        / cfg.environment._step_log_denominator,
                    )

                cfg.logging._logger.log(
                    "internal",
                    "current_step",
                    cfg.environment._curr_step,
                )
                for key in output_dict:
                    if key.startswith("additional_log_"):
                        cfg.logging._logger.log(
                            "train",
                            key.replace("additional_log_", ""),
                            output_dict[key].item(),
                            step=cfg.environment._curr_step
                            / cfg.environment._step_log_denominator,
                        )

                # Show logs each 5% of the epoch (only if doing per epoch evaluation)
                if (itr + 1) % log_update_steps == 0 or itr == epoch_steps - 1:
                    progress_bar.set_description(
                        f"train loss: {np.mean(losses[-10:]):.2f}", refresh=False
                    )
                    if (itr + 1) % log_update_steps == 0:
                        progress_bar.update(log_update_steps)
                    else:
                        progress_bar.update(epoch_steps % log_update_steps)

                del output_dict

            # Validation loop
            if (itr + 1) % evaluation_step == 0:
                # TODO: Move back after fixing slow generation of deepspeed.
                if cfg.training.save_checkpoint == "last":
                    logger.info(
                        f"Saving last model checkpoint to {cfg.output_directory}"
                    )
                    save_checkpoint(model=model, path=cfg.output_directory, cfg=cfg)
                elif cfg.training.save_checkpoint == "each_evaluation_epoch":
                    checkpoint_path = os.path.join(
                        cfg.output_directory, f"epoch_{epoch}_step_{itr}"
                    )
                    logger.info(f"Saving model checkpoint to {checkpoint_path}")
                    save_checkpoint(model=model, path=checkpoint_path, cfg=cfg)
                    create_symlinks_in_parent_folder(checkpoint_path)

                val_loss, val_metric = run_eval(
                    cfg=cfg, model=model, val_dataloader=val_dataloader, val_df=val_df
                )

                if cfg.training.save_checkpoint == "best":
                    if objective_op(val_metric, best_val_metric):
                        logger.info(
                            f"Saving best model checkpoint: "
                            f"val_{cfg.prediction.metric} {best_val_metric:.5} -> "
                            f"{val_metric:.5} to {cfg.output_directory}"
                        )
                        save_checkpoint(model=model, path=cfg.output_directory, cfg=cfg)
                        best_val_metric = val_metric

                model.train()

        progress_bar.close()
        del progress_bar

        if cfg.environment._distributed:
            torch.cuda.synchronize(device=cfg.environment._local_rank)
            torch.distributed.barrier()

        if cfg.environment._local_rank == 0:
            cfg.logging._logger.log("internal", "epoch", epoch + 1)

    if cfg.environment._distributed:
        torch.distributed.barrier()

    return val_loss, val_metric

def run_train(
    cfg: DefaultConfigProblemBase,
    model: torch.nn.Module,
    optimizer,
    scheduler,
    epoch_steps,
    train_dataloader,
    val_dataloader,
    val_df: pd.DataFrame,
):

    if (
        hasattr(cfg.augmentation, "neftune_noise_alpha")
        and cfg.augmentation.neftune_noise_alpha > 0
    ):
        activate_neftune(model, cfg.augmentation.neftune_noise_alpha)

    scaler: GradScaler | None = None
    if cfg.environment.mixed_precision:
        scaler = GradScaler(
            enabled=(cfg.environment.mixed_precision_dtype == "float16")
        )

    optimizer.zero_grad(set_to_none=True)

    # Prepare NLP Augmentation
    nlp_augment = None
    if hasattr(cfg.augmentation, "nlp_augmentations_class"):
        nlp_augment = cfg.augmentation.nlp_augmentations_class(cfg=cfg)

    start_epoch = 0

    _, metric_mode, _ = cfg.prediction.metric_class.get(cfg.prediction.metric)
    objective_op: Callable[[float, float], bool]
    if metric_mode == "max":
        best_val_metric = -np.inf
        objective_op = np.greater
    else:
        best_val_metric = np.inf
        objective_op = np.less

    if cfg.training.evaluate_before_training:
        val_loss, val_metric = run_eval(
            cfg=cfg, model=model, val_dataloader=val_dataloader, val_df=val_df
        )

    for epoch in range(start_epoch, cfg.training.epochs):
        set_seed(
            cfg.environment._seed
            + epoch * cfg.environment._world_size * cfg.environment.number_of_workers
            + cfg.environment._local_rank * cfg.environment.number_of_workers
        )
        logger.info(f"Training Epoch: {epoch + 1} / {cfg.training.epochs}")

        if (
            cfg.environment._distributed
            and not cfg.environment.use_deepspeed
            and hasattr(train_dataloader.sampler, "set_epoch")
        ):
            train_dataloader.sampler.set_epoch(epoch)  # type: ignore

        tqdm_out = TqdmToLogger(logger, level=logging.INFO)
        progress_bar = tqdm(
            total=epoch_steps,
            disable=cfg.environment._local_rank != 0,
            file=tqdm_out,
            ascii=True,
            desc="train loss",
            mininterval=0,
        )
        tr_it = iter(train_dataloader)

        losses = []
        model.train()

        log_update_steps = max(epoch_steps // 20, 1)
        evaluation_step = max(int(epoch_steps * cfg.training.evaluation_epochs), 1)
        logger.info(f"Evaluation step: {evaluation_step}")

        for itr, data in enumerate(tr_it):
            cfg.environment._curr_step += (
                cfg.training.batch_size * cfg.environment._world_size
            )

            # Batch to device
            batch = cfg.dataset.dataset_class.batch_to_device(
                data, cfg.environment._device
            )

            # NLP augmentation
            if nlp_augment is not None:
                batch = nlp_augment(batch)

            # Plot first batch
            if epoch == 0 and itr == 0 and cfg.environment._local_rank == 0:
                plot = cfg.logging.plots_class.plot_batch(batch=batch, cfg=cfg)
                log_plot(cfg, plot, "train_data")

            # only need to sync gradients at last step of grad accumulation
            model.require_backward_grad_sync = itr % cfg.training.grad_accumulation == 0

            # Forward pass
            with autocast(
                enabled=cfg.environment.mixed_precision,
                dtype=get_torch_dtype(cfg.environment.mixed_precision_dtype),
            ):
                output_dict = model.forward(batch)

            loss = output_dict["loss"]
            if ~np.isfinite(loss.item()) and (epoch > start_epoch or itr > 20):
                raise LLMTrainingException(
                    "NaN caught in loss during training. "
                    "Please, reduce learning rate, change dtype, "
                    "or disable mixed precision. Alternatively, "
                    "gradient clipping may help to stabilize training."
                )
            losses.append(loss.item())

            # loss is a mean loss per batch/sample
            # as grad_accumulations sums up the gradients, this loss must be scaled
            # by the number of grad_accumulations, to have similar behavior for
            # BS * grad_accumulations = const.
            if cfg.training.grad_accumulation != 1:
                loss = loss / cfg.training.grad_accumulation

            # Backward pass
            if (
                cfg.environment.mixed_precision
                and len(cfg.environment.gpus)
                and not cfg.environment.use_deepspeed
            ):
                scaler.scale(loss).backward()  # type: ignore
                if itr % cfg.training.grad_accumulation == 0:
                    if cfg.training.gradient_clip > 0:
                        scaler.unscale_(optimizer)  # type: ignore
                        torch.nn.utils.clip_grad_norm_(
                            model.parameters(), cfg.training.gradient_clip
                        )
                    scaler.step(optimizer)  # type: ignore
                    scaler.update()
                    optimizer.zero_grad(set_to_none=True)
            else:
                if cfg.environment.use_deepspeed:
                    model.backward(loss)  # type: ignore[operator]
                else:
                    loss.backward()
                if itr % cfg.training.grad_accumulation == 0:
                    if cfg.training.gradient_clip > 0:
                        torch.nn.utils.clip_grad_norm_(
                            model.parameters(), cfg.training.gradient_clip
                        )
                    optimizer.step()
                    optimizer.zero_grad(set_to_none=True)

            if cfg.environment._distributed:
                torch.cuda.synchronize(device=cfg.environment._local_rank)

            if scheduler is not None:
                scheduler.step()

            if cfg.environment._local_rank == 0:
                cfg.logging._logger.log(
                    "train",
                    "loss",
                    losses[-1],
                    step=cfg.environment._curr_step
                    / cfg.environment._step_log_denominator,
                )
                cfg.logging._logger.log(
                    "meta",
                    "lr",
                    optimizer.param_groups[0]["lr"],
                    step=cfg.environment._curr_step
                    / cfg.environment._step_log_denominator,
                )
                if cfg.training.differential_learning_rate_layers:
                    cfg.logging._logger.log(
                        "meta",
                        "lr_diff",
                        optimizer.param_groups[2]["lr"],
                        step=cfg.environment._curr_step
                        / cfg.environment._step_log_denominator,
                    )

                cfg.logging._logger.log(
                    "internal",
                    "current_step",
                    cfg.environment._curr_step,
                )
                for key in output_dict:
                    if key.startswith("additional_log_"):
                        cfg.logging._logger.log(
                            "train",
                            key.replace("additional_log_", ""),
                            output_dict[key].item(),
                            step=cfg.environment._curr_step
                            / cfg.environment._step_log_denominator,
                        )

                # Show logs each 5% of the epoch (only if doing per epoch evaluation)
                if (itr + 1) % log_update_steps == 0 or itr == epoch_steps - 1:
                    progress_bar.set_description(
                        f"train loss: {np.mean(losses[-10:]):.2f}", refresh=False
                    )
                    if (itr + 1) % log_update_steps == 0:
                        progress_bar.update(log_update_steps)
                    else:
                        progress_bar.update(epoch_steps % log_update_steps)

                del output_dict

            # Validation loop
            if (itr + 1) % evaluation_step == 0:
                # TODO: Move back after fixing slow generation of deepspeed.
                if cfg.training.save_checkpoint == "last":
                    logger.info(
                        f"Saving last model checkpoint to {cfg.output_directory}"
                    )
                    save_checkpoint(model=model, path=cfg.output_directory, cfg=cfg)
                elif cfg.training.save_checkpoint == "each_evaluation_epoch":
                    checkpoint_path = os.path.join(
                        cfg.output_directory, f"epoch_{epoch}_step_{itr}"
                    )
                    logger.info(f"Saving model checkpoint to {checkpoint_path}")
                    save_checkpoint(model=model, path=checkpoint_path, cfg=cfg)
                    create_symlinks_in_parent_folder(checkpoint_path)

                val_loss, val_metric = run_eval(
                    cfg=cfg, model=model, val_dataloader=val_dataloader, val_df=val_df
                )

                if cfg.training.save_checkpoint == "best":
                    if objective_op(val_metric, best_val_metric):
                        logger.info(
                            f"Saving best model checkpoint: "
                            f"val_{cfg.prediction.metric} {best_val_metric:.5} -> "
                            f"{val_metric:.5} to {cfg.output_directory}"
                        )
                        save_checkpoint(model=model, path=cfg.output_directory, cfg=cfg)
                        best_val_metric = val_metric

                model.train()

        progress_bar.close()
        del progress_bar

        if cfg.environment._distributed:
            torch.cuda.synchronize(device=cfg.environment._local_rank)
            torch.distributed.barrier()

        if cfg.environment._local_rank == 0:
            cfg.logging._logger.log("internal", "epoch", epoch + 1)

    if cfg.environment._distributed:
        torch.distributed.barrier()

    return val_loss, val_metric


def create_custom_chat_agent(
    name: str, llm_config: OpenAIGPTConfig, system_message: str
) -> ChatAgent:
    """creates a ChatAgent with the given parameters.

    Args:
        name (str): The name of the agent.
        llm_config (OpenAIGPTConfig): The LLM configuration for the agent.
        system_message (str): The system message to guide the agent's LLM.

    Returns:
        ChatAgent: A configured ChatAgent instance.
    """
    # Modify the system message to include instructions for the agent
    additional_system_message = """**Response format (strictly follow this structure):**
    Pro:
    - [First key point]
    - [Second key point]
    - [Third key point]
    **Limit responses to exactly 3 points expressed as single sentences.**"
    """
    system_message = f"""
       Start your response with '{name}: ' and then follow the instructions below.
        {system_message} {additional_system_message}
        """
    return ChatAgent(
        ChatAgentConfig(
            llm=llm_config,
            name=name,
            system_message=system_message,
        )
    )

def create_custom_chat_agent(
    name: str, llm_config: OpenAIGPTConfig, system_message: str
) -> ChatAgent:

    # Modify the system message to include instructions for the agent
    additional_system_message = """**Response format (strictly follow this structure):**
    Pro:
    - [First key point]
    - [Second key point]
    - [Third key point]
    **Limit responses to exactly 3 points expressed as single sentences.**"
    """
    system_message = f"""
       Start your response with '{name}: ' and then follow the instructions below.
        {system_message} {additional_system_message}
        """
    return ChatAgent(
        ChatAgentConfig(
            llm=llm_config,
            name=name,
            system_message=system_message,
        )
    )


def clear_all_collections(self, really: bool = False, prefix: str = "") -> int:
        """Clear all collections with the given prefix."""

        if not really:
            logger.warning("Not deleting all collections, set really=True to confirm")
            return 0
        coll_names = self.list_collections(empty=False)
        coll_names = [name for name in coll_names if name.startswith(prefix)]
        if len(coll_names) == 0:
            logger.warning(f"No collections found with prefix {prefix}")
            return 0
        for name in coll_names:
            self.delete_collection(name)
        logger.warning(
            f"""
            Deleted {len(coll_names)} indices from Momento VI
            """
        )
        return len(coll_names)

def clear_all_collections(self, really: bool = False, prefix: str = "") -> int:

        if not really:
            logger.warning("Not deleting all collections, set really=True to confirm")
            return 0
        coll_names = self.list_collections(empty=False)
        coll_names = [name for name in coll_names if name.startswith(prefix)]
        if len(coll_names) == 0:
            logger.warning(f"No collections found with prefix {prefix}")
            return 0
        for name in coll_names:
            self.delete_collection(name)
        logger.warning(
            f"""
            Deleted {len(coll_names)} indices from Momento VI
            """
        )
        return len(coll_names)


@staticmethod
    def verify_onnx(
            model: T5EncoderDecoderInit,
            ort_session: InferenceSession,
            device: torch.device,
            max_cases=4,
    ):
        """Compare the result from PyTorch and OnnxRuntime to verify the ONNX model is good."""
        ort_inputs = ort_session.get_inputs()
        use_decoder_input_ids = len(ort_inputs) == 3

        test_cases = [(4, 11), (1, 2), (3, 1), (8, 5)]
        test_cases_max_diff = []
        for (batch_size, encode_sequence_length) in test_cases[:max_cases]:
            inputs = T5EncoderDecoderInitInputs.create_dummy(
                model.config,
                batch_size,
                encode_sequence_length,
                use_decoder_input_ids=use_decoder_input_ids,
                device=device,
            )

            ort_outputs = T5EncoderDecoderInitHelper.onnxruntime_inference(ort_session, inputs)

            # Run inference of PyTorch model
            input_list = inputs.to_list()
            torch_outputs = model(*input_list)

            assert torch_outputs[0].cpu().numpy().shape == ort_outputs['logits'].shape
            max_diff = numpy.amax(numpy.abs(torch_outputs[0].cpu().numpy() - ort_outputs['logits'].cpu().numpy()))
            logger.debug(f"logits max_diff={max_diff}")
            max_diff_all = max_diff

            assert torch_outputs[1].cpu().numpy().shape == ort_outputs['encoder_hidden_states'].shape
            max_diff = numpy.amax(
                numpy.abs(torch_outputs[1].cpu().numpy() - ort_outputs['encoder_hidden_states'].cpu().numpy()))
            logger.debug(f"encoder_hidden_states max_diff={max_diff}")
            max_diff_all = max(max_diff_all, max_diff)

            for i in range(2 * model.config.num_layers):
                if i % 2 == 0:
                    ort_outputs_i = ort_outputs[f'present_key_self_{i//2}']
                else:
                    ort_outputs_i = ort_outputs[f'present_value_self_{i//2}']
                max_diff = numpy.amax(numpy.abs(torch_outputs[2][i].cpu().numpy() - ort_outputs_i.cpu().numpy()))
                logger.debug(f"self attention past state {i} max_diff={max_diff}")

            for i in range(2 * model.config.num_layers):
                if i % 2 == 0:
                    ort_outputs_i = ort_outputs[f'present_key_cross_{i//2}']
                else:
                    ort_outputs_i = ort_outputs[f'present_value_cross_{i//2}']
                max_diff = numpy.amax(
                    numpy.abs(torch_outputs[3][i].cpu().numpy() - ort_outputs_i.cpu().numpy())
                )
                logger.debug(f"cross attention past state {i} max_diff={max_diff}")
                max_diff_all = max(max_diff_all, max_diff)

            test_cases_max_diff.append(max_diff_all)
            logger.info(
                f"batch_size={batch_size} encode_sequence_length={encode_sequence_length}, max_diff={max_diff_all}"
            )

        return max(test_cases_max_diff)

@staticmethod
    def verify_onnx(
            model: T5EncoderDecoderInit,
            ort_session: InferenceSession,
            device: torch.device,
            max_cases=4,
    ):

        ort_inputs = ort_session.get_inputs()
        use_decoder_input_ids = len(ort_inputs) == 3

        test_cases = [(4, 11), (1, 2), (3, 1), (8, 5)]
        test_cases_max_diff = []
        for (batch_size, encode_sequence_length) in test_cases[:max_cases]:
            inputs = T5EncoderDecoderInitInputs.create_dummy(
                model.config,
                batch_size,
                encode_sequence_length,
                use_decoder_input_ids=use_decoder_input_ids,
                device=device,
            )

            ort_outputs = T5EncoderDecoderInitHelper.onnxruntime_inference(ort_session, inputs)

            # Run inference of PyTorch model
            input_list = inputs.to_list()
            torch_outputs = model(*input_list)

            assert torch_outputs[0].cpu().numpy().shape == ort_outputs['logits'].shape
            max_diff = numpy.amax(numpy.abs(torch_outputs[0].cpu().numpy() - ort_outputs['logits'].cpu().numpy()))
            logger.debug(f"logits max_diff={max_diff}")
            max_diff_all = max_diff

            assert torch_outputs[1].cpu().numpy().shape == ort_outputs['encoder_hidden_states'].shape
            max_diff = numpy.amax(
                numpy.abs(torch_outputs[1].cpu().numpy() - ort_outputs['encoder_hidden_states'].cpu().numpy()))
            logger.debug(f"encoder_hidden_states max_diff={max_diff}")
            max_diff_all = max(max_diff_all, max_diff)

            for i in range(2 * model.config.num_layers):
                if i % 2 == 0:
                    ort_outputs_i = ort_outputs[f'present_key_self_{i//2}']
                else:
                    ort_outputs_i = ort_outputs[f'present_value_self_{i//2}']
                max_diff = numpy.amax(numpy.abs(torch_outputs[2][i].cpu().numpy() - ort_outputs_i.cpu().numpy()))
                logger.debug(f"self attention past state {i} max_diff={max_diff}")

            for i in range(2 * model.config.num_layers):
                if i % 2 == 0:
                    ort_outputs_i = ort_outputs[f'present_key_cross_{i//2}']
                else:
                    ort_outputs_i = ort_outputs[f'present_value_cross_{i//2}']
                max_diff = numpy.amax(
                    numpy.abs(torch_outputs[3][i].cpu().numpy() - ort_outputs_i.cpu().numpy())
                )
                logger.debug(f"cross attention past state {i} max_diff={max_diff}")
                max_diff_all = max(max_diff_all, max_diff)

            test_cases_max_diff.append(max_diff_all)
            logger.info(
                f"batch_size={batch_size} encode_sequence_length={encode_sequence_length}, max_diff={max_diff_all}"
            )

        return max(test_cases_max_diff)