# Copyright (c) ModelScope Contributors. All rights reserved.
from typing import TYPE_CHECKING

from swift.utils.import_utils import _LazyModule
from . import patcher

if TYPE_CHECKING:
    from .arguments import Seq2SeqTrainingArguments, TrainArgumentsMixin, TrainingArguments
    from .embedding_trainer import EmbeddingTrainer
    from .mixin import DataLoaderMixin, SwiftMixin
    from .reranker_trainer import RerankerTrainer
    from .seq2seq_trainer import Seq2SeqTrainer
    from .trainer import Trainer
    from .trainer_factory import TrainerFactory
    from .utils import (calculate_max_steps, disable_gradient_checkpointing, dynamic_gradient_checkpointing,
                        per_token_loss_func)
else:
    _import_structure = {
        'arguments': ['TrainArgumentsMixin', 'Seq2SeqTrainingArguments', 'TrainingArguments'],
        'embedding_trainer': ['EmbeddingTrainer'],
        'mixin': ['DataLoaderMixin', 'SwiftMixin'],
        'reranker_trainer': ['RerankerTrainer'],
        'seq2seq_trainer': ['Seq2SeqTrainer'],
        'trainer': ['Trainer'],
        'trainer_factory': ['TrainerFactory'],
        'utils': [
            'disable_gradient_checkpointing', 'dynamic_gradient_checkpointing', 'per_token_loss_func',
            'calculate_max_steps'
        ]
    }

    import sys

    sys.modules[__name__] = _LazyModule(
        __name__,
        globals()['__file__'],
        _import_structure,
        module_spec=__spec__,
        extra_objects={},
    )