from .modeling import GECToR
from .configuration import GECToRConfig
from .dataset import load_dataset, GECToRDataset
from .predict import predict, load_verb_dict
from .predict_verbose import predict_verbose
from .vocab import (
    build_vocab,
    load_vocab_from_config,
    load_vocab_from_official
)

# Import Triton model only if tritonclient is available
try:
    from .triton_modeling import GECToRTriton
    _triton_available = True
except ImportError:
    _triton_available = False
    GECToRTriton = None

__all__ = [
    'GECToR',
    'GECToRConfig',
    'load_dataset',
    'GECToRDataset',
    'predict',
    'load_verb_dict',
    'predict_verbose',
    'build_vocab',
    'load_vocab_from_config',
    'load_vocab_from_official'
]

if _triton_available:
    __all__.append('GECToRTriton')