Model/__init__.py

# -*- coding: utf-8 -*-

"""Recurrent depth transformer model package."""

from Model.config import (
    OMVTConfig,
    RDTConfig,
    TrainingConfig,
    base_config,
    pretrain_config,
    small_config,
    tiny_config,
)
from Model.model import RDTForCausalLM
from Model.training import JsonlPretrainingDataset, PretrainingCollator

__all__ = [
    "JsonlPretrainingDataset",
    "OMVTConfig",
    "PretrainingCollator",
    "RDTConfig",
    "RDTForCausalLM",
    "TrainingConfig",
    "base_config",
    "pretrain_config",
    "small_config",
    "tiny_config",
]