virtex/models/__init__.py

from .captioning import (
    ForwardCaptioningModel,
    BidirectionalCaptioningModel,
    VirTexModel
)
from .masked_lm import MaskedLMModel
from .classification import (
    MultiLabelClassificationModel,
    TokenClassificationModel,
)
from .contrastive import ImageTextContrastiveModel


__all__ = [
    "VirTexModel",
    "BidirectionalCaptioningModel",
    "ForwardCaptioningModel",
    "MaskedLMModel",
    "MultiLabelClassificationModel",
    "TokenClassificationModel",
    "ImageTextContrastiveModel",
]