Using encoder/decoder registries

ludwig/schema/features/audio_feature.py

@@ -2,6 +2,8 @@
 
 from marshmallow_dataclass import dataclass
 
+from ludwig.encoders.registry import get_encoder_classes
+
 from ludwig.schema import utils as schema_utils
 from ludwig.schema.preprocessing import BasePreprocessingConfig, PreprocessingDataclassField
 
@@ -15,7 +17,7 @@ class AudioInputFeatureConfig(schema_utils.BaseMarshmallowConfig):
     )
 
     encoder: Optional[str] = schema_utils.StringOptions(
-        ["passthrough", "parallel_cnn", "stacked_cnn", "stacked_parallel_cnn", "rnn", "cnnrnn"],
+        list(get_encoder_classes('audio').keys()),
         default="parallel_cnn",
         description="Encoder to use for this audio feature.",
     )

ludwig/schema/features/bag_feature.py

@@ -2,6 +2,8 @@
 
 from marshmallow_dataclass import dataclass
 
+from ludwig.encoders.registry import get_encoder_classes
+
 from ludwig.schema import utils as schema_utils
 from ludwig.schema.preprocessing import BasePreprocessingConfig, PreprocessingDataclassField
 
@@ -17,7 +19,7 @@ class BagInputFeatureConfig(schema_utils.BaseMarshmallowConfig):
     )
 
     encoder: Optional[str] = schema_utils.StringOptions(
-        ["embed"],
+        list(get_encoder_classes('bag').keys()),
         default="embed",
         description="Encoder to use for this bag feature.",
     )

ludwig/schema/features/binary_feature.py

@@ -2,6 +2,9 @@
 
 from marshmallow_dataclass import dataclass
 
+from ludwig.encoders.registry import get_encoder_classes
+from ludwig.decoders.registry import get_decoder_classes
+
 from ludwig.schema import utils as schema_utils
 from ludwig.schema.preprocessing import BasePreprocessingConfig, PreprocessingDataclassField
 
@@ -15,7 +18,7 @@ class BinaryInputFeatureConfig(schema_utils.BaseMarshmallowConfig):
     )
 
     encoder: Optional[str] = schema_utils.StringOptions(
-        ["passthrough", "dense"],
+        list(get_encoder_classes('binary').keys()),
         default="passthrough",
         description="Encoder to use for this binary feature.",
     )
@@ -34,7 +37,7 @@ class BinaryOutputFeatureConfig(schema_utils.BaseMarshmallowConfig):
     """BinaryOutputFeature is a dataclass that configures the parameters used for a binary output feature."""
 
     decoder: Optional[str] = schema_utils.StringOptions(
-        ["regressor"],
+        list(get_decoder_classes('binary').keys()),
         default="regressor",
         allow_none=True,
         description="Decoder to use for this binary feature.",

ludwig/schema/features/category_feature.py

@@ -2,6 +2,9 @@
 
 from marshmallow_dataclass import dataclass
 
+from ludwig.encoders.registry import get_encoder_classes
+from ludwig.decoders.registry import get_decoder_classes
+
 from ludwig.schema import utils as schema_utils
 from ludwig.schema.preprocessing import BasePreprocessingConfig, PreprocessingDataclassField
 
@@ -15,7 +18,7 @@ class CategoryInputFeatureConfig(schema_utils.BaseMarshmallowConfig):
     )
 
     encoder: Optional[str] = schema_utils.StringOptions(
-        ["passthrough", "dense", "sparse"],
+        list(get_encoder_classes('category').keys()),
         default="dense",
         description="Encoder to use for this category feature.",
     )
@@ -34,7 +37,7 @@ class CategoryOutputFeatureConfig(schema_utils.BaseMarshmallowConfig):
     """CategoryOutputFeature is a dataclass that configures the parameters used for a category output feature."""
 
     decoder: Optional[str] = schema_utils.StringOptions(
-        ["classifier"],
+        list(get_decoder_classes('category').keys()),
         default="classifier",
         allow_none=True,
         description="Decoder to use for this category feature.",

ludwig/schema/features/date_feature.py

@@ -2,6 +2,8 @@
 
 from marshmallow_dataclass import dataclass
 
+from ludwig.encoders.registry import get_encoder_classes
+
 from ludwig.schema import utils as schema_utils
 from ludwig.schema.preprocessing import BasePreprocessingConfig, PreprocessingDataclassField
 
@@ -17,7 +19,7 @@ class DateInputFeatureConfig(schema_utils.BaseMarshmallowConfig):
     )
 
     encoder: Optional[str] = schema_utils.StringOptions(
-        ["embed", "wave"],
+        list(get_encoder_classes('date').keys()),
         default="embed",
         description="Encoder to use for this date feature.",
     )

ludwig/schema/features/h3_feature.py

@@ -2,6 +2,8 @@
 
 from marshmallow_dataclass import dataclass
 
+from ludwig.encoders.registry import get_encoder_classes
+
 from ludwig.schema import utils as schema_utils
 from ludwig.schema.preprocessing import BasePreprocessingConfig, PreprocessingDataclassField
 
@@ -17,7 +19,7 @@ class H3InputFeatureConfig(schema_utils.BaseMarshmallowConfig):
     )
 
     encoder: Optional[str] = schema_utils.StringOptions(
-        ["embed", "weighted_sum", "rnn"],
+        list(get_encoder_classes('h3').keys()),
         default="embed",
         description="Encoder to use for this h3 feature.",
     )

ludwig/schema/features/image_feature.py

@@ -2,6 +2,8 @@
 
 from marshmallow_dataclass import dataclass
 
+from ludwig.encoders.registry import get_encoder_classes
+
 from ludwig.schema import utils as schema_utils
 from ludwig.schema.preprocessing import BasePreprocessingConfig, PreprocessingDataclassField
 
@@ -17,7 +19,7 @@ class ImageInputFeatureConfig(schema_utils.BaseMarshmallowConfig):
     )
 
     encoder: Optional[str] = schema_utils.StringOptions(
-        ["stacked_cnn", "resnet", "mlp_mixer", "vit"],
+        list(get_encoder_classes('image').keys()),
         default="stacked_cnn",
         description="Encoder to use for this image feature.",
     )

ludwig/schema/features/number_feature.py

@@ -2,6 +2,9 @@
 
 from marshmallow_dataclass import dataclass
 
+from ludwig.encoders.registry import get_encoder_classes
+from ludwig.decoders.registry import get_decoder_classes
+
 from ludwig.schema import utils as schema_utils
 from ludwig.schema.preprocessing import BasePreprocessingConfig, PreprocessingDataclassField
 
@@ -15,7 +18,7 @@ class NumberInputFeatureConfig(schema_utils.BaseMarshmallowConfig):
     )
 
     encoder: Optional[str] = schema_utils.StringOptions(
-        ["passthrough", "dense"],
+        list(get_encoder_classes('number').keys()),
         default="passthrough",
         description="Encoder to use for this number feature.",
     )
@@ -33,7 +36,7 @@ class NumberInputFeatureConfig(schema_utils.BaseMarshmallowConfig):
 class NumberOutputFeatureConfig(schema_utils.BaseMarshmallowConfig):
 
     decoder: Optional[str] = schema_utils.StringOptions(
-        ["regressor"],
+        list(get_decoder_classes('number').keys()),
         default="regressor",
         allow_none=True,
         description="Decoder to use for this number feature.",

ludwig/schema/features/sequence_feature.py

@@ -2,6 +2,9 @@
 
 from marshmallow_dataclass import dataclass
 
+from ludwig.encoders.registry import get_encoder_classes
+from ludwig.decoders.registry import get_decoder_classes
+
 from ludwig.schema import utils as schema_utils
 from ludwig.schema.preprocessing import BasePreprocessingConfig, PreprocessingDataclassField
 
@@ -17,7 +20,7 @@ class SequenceInputFeatureConfig(schema_utils.BaseMarshmallowConfig):
     )
 
     encoder: Optional[str] = schema_utils.StringOptions(
-        ["passthrough", "embed", "parallel_cnn", "stacked_cnn", "stacked_parallel_cnn", "rnn", "cnnrnn", "transformer"],
+        list(get_encoder_classes('sequence').keys()),
         default="embed",
         description="Encoder to use for this sequence feature.",
     )
@@ -38,7 +41,7 @@ class SequenceOutputFeatureConfig(schema_utils.BaseMarshmallowConfig):
     """
 
     decoder: Optional[str] = schema_utils.StringOptions(
-        ["generator", "tagger"],
+        list(get_decoder_classes('sequence').keys()),
         default="generator",
         allow_none=True,
         description="Decoder to use for this sequence feature.",

ludwig/schema/features/set_feature.py

@@ -2,6 +2,9 @@
 
 from marshmallow_dataclass import dataclass
 
+from ludwig.encoders.registry import get_encoder_classes
+from ludwig.decoders.registry import get_decoder_classes
+
 from ludwig.schema import utils as schema_utils
 from ludwig.schema.preprocessing import BasePreprocessingConfig, PreprocessingDataclassField
 
@@ -17,7 +20,7 @@ class SetInputFeatureConfig(schema_utils.BaseMarshmallowConfig):
     )
 
     encoder: Optional[str] = schema_utils.StringOptions(
-        ["embed"],
+        list(get_encoder_classes('set').keys()),
         default="embed",
         description="Encoder to use for this set feature.",
     )
@@ -38,7 +41,7 @@ class SetOutputFeatureConfig(schema_utils.BaseMarshmallowConfig):
     """
 
     decoder: Optional[str] = schema_utils.StringOptions(
-        ["classifier"],
+        list(get_decoder_classes('set').keys()),
         default="classifier",
         allow_none=True,
         description="Decoder to use for this set feature.",

ludwig/schema/features/text_feature.py

@@ -2,6 +2,9 @@
 
 from marshmallow_dataclass import dataclass
 
+from ludwig.encoders.registry import get_encoder_classes
+from ludwig.decoders.registry import get_decoder_classes
+
 from ludwig.schema import utils as schema_utils
 from ludwig.schema.preprocessing import BasePreprocessingConfig, PreprocessingDataclassField
 
@@ -17,9 +20,7 @@ class TextInputFeatureConfig(schema_utils.BaseMarshmallowConfig):
     )
 
     encoder: Optional[str] = schema_utils.StringOptions(
-        ["passthrough", "embed", "parallel_cnn", "stacked_cnn", "stacked_parallel_cnn", "rnn", "cnnrnn", "transformer",
-         "albert", "mt5", "xlmroberta", "bert", "xlm", "gpt", "gpt2", "roberta", "transformer_xl", "xlnet",
-         "distilbert", "ctrl", "camembert", "t5", "flaubert", "electra", "longformer", "auto_transformer"],
+        list(get_encoder_classes('text').keys()),
         default="parallel_cnn",
         description="Encoder to use for this text feature.",
     )
@@ -40,7 +41,7 @@ class TextOutputFeatureConfig(schema_utils.BaseMarshmallowConfig):
     """
 
     decoder: Optional[str] = schema_utils.StringOptions(
-        ["tagger", "generator"],
+        list(get_decoder_classes('text').keys()),
         default="generator",
         description="Decoder to use for this text output feature.",
     )

ludwig/schema/features/timeseries_feature.py

@@ -2,6 +2,8 @@
 
 from marshmallow_dataclass import dataclass
 
+from ludwig.encoders.registry import get_encoder_classes
+
 from ludwig.schema import utils as schema_utils
 from ludwig.schema.preprocessing import BasePreprocessingConfig, PreprocessingDataclassField
 
@@ -17,7 +19,7 @@ class TimeseriesInputFeatureConfig(schema_utils.BaseMarshmallowConfig):
     )
 
     encoder: Optional[str] = schema_utils.StringOptions(
-        ["passthrough", "parallel_cnn", "stacked_cnn", "stacked_parallel_cnn", "rnn", "cnnrnn", "transformer"],
+        list(get_encoder_classes('timeseries').keys()),
         default="parallel_cnn",
         description="Encoder to use for this timeseries feature.",
     )

ludwig/schema/features/vector_feature.py

@@ -2,6 +2,9 @@
 
 from marshmallow_dataclass import dataclass
 
+from ludwig.encoders.registry import get_encoder_classes
+from ludwig.decoders.registry import get_decoder_classes
+
 from ludwig.schema import utils as schema_utils
 from ludwig.schema.preprocessing import BasePreprocessingConfig, PreprocessingDataclassField
 
@@ -17,7 +20,7 @@ class VectorInputFeatureConfig(schema_utils.BaseMarshmallowConfig):
     )
 
     encoder: Optional[str] = schema_utils.StringOptions(
-        ["passthrough", "dense"],
+        list(get_encoder_classes('vector').keys()),
         default="dense",
         description="Encoder to use for this vector feature.",
     )
@@ -38,7 +41,7 @@ class VectorOutputFeatureConfig(schema_utils.BaseMarshmallowConfig):
     """
 
     decoder: Optional[str] = schema_utils.StringOptions(
-        ["projector"],
+        list(get_decoder_classes('vector').keys()),
         default="projector",
         description="Decoder to use for this vector feature.",
     )