-
Notifications
You must be signed in to change notification settings - Fork 516
/
__init__.py
51 lines (39 loc) · 1.71 KB
/
__init__.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
#
# For licensing see accompanying LICENSE file.
# Copyright (C) 2024 Apple Inc. All Rights Reserved.
#
import argparse
from corenet.modeling.text_encoders.base_text_encoder import BaseTextEncoder
from corenet.utils import logger
from corenet.utils.registry import Registry
TEXT_ENCODER_REGISTRY = Registry(
"text_encoder",
base_class=BaseTextEncoder,
lazy_load_dirs=["corenet/modeling/text_encoders"],
internal_dirs=["corenet/internal", "corenet/internal/projects/*"],
)
def arguments_text_encoder(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
"""Register arguments of all text encoders."""
# add arguments for text_encoder
parser = BaseTextEncoder.add_arguments(parser)
# add class specific arguments
parser = TEXT_ENCODER_REGISTRY.all_arguments(parser)
return parser
def build_text_encoder(opts, projection_dim: int, *args, **kwargs) -> BaseTextEncoder:
"""Helper function to build the text encoder from command-line arguments.
Args:
opts: Command-line arguments
projection_dim: The dimensionality of the projection head after text encoder.
Returns:
Text encoder module.
"""
text_encoder_name = getattr(opts, "model.text.name")
# We registered the base class using a special `name` (i.e., `__base__`)
# in order to access the arguments defined inside those classes. However, these classes are not supposed to
# be used. Therefore, we raise an error for such cases
if text_encoder_name == "__base__":
logger.error("__base__ can't be used as a projection name. Please check.")
text_encoder = TEXT_ENCODER_REGISTRY[text_encoder_name](
opts, projection_dim, *args, **kwargs
)
return text_encoder