Update configurations

examples/conformer/save_conformer_from_weights.py

@@ -46,23 +46,20 @@
 
 setup_devices([args.device], cpu=args.cpu)
 
-from tensorflow_asr.configs.user_config import UserConfig
+from tensorflow_asr.configs.config import Config
 from tensorflow_asr.featurizers.speech_featurizers import TFSpeechFeaturizer
 from tensorflow_asr.featurizers.text_featurizers import CharFeaturizer
 from tensorflow_asr.models.conformer import Conformer
 
-config = UserConfig(DEFAULT_YAML, args.config, learning=True)
-speech_featurizer = TFSpeechFeaturizer(config["speech_config"])
-text_featurizer = CharFeaturizer(config["decoder_config"])
+config = Config(args.config, learning=True)
+speech_featurizer = TFSpeechFeaturizer(config.speech_config)
+text_featurizer = CharFeaturizer(config.decoder_config)
 
 tf.random.set_seed(0)
 assert args.saved
 
 # build model
-conformer = Conformer(
-    vocabulary_size=text_featurizer.num_classes,
-    **config["model_config"]
-)
+conformer = Conformer(**config.model_config, vocabulary_size=text_featurizer.num_classes)
 conformer._build(speech_featurizer.shape)
 conformer.load_weights(args.saved, by_name=True)
 conformer.summary(line_length=150)

examples/conformer/test_conformer.py

@@ -52,48 +52,45 @@
 
 setup_devices([args.device], cpu=args.cpu)
 
-from tensorflow_asr.configs.user_config import UserConfig
+from tensorflow_asr.configs.config import Config
 from tensorflow_asr.datasets.asr_dataset import ASRTFRecordDataset, ASRSliceDataset
 from tensorflow_asr.featurizers.speech_featurizers import TFSpeechFeaturizer
 from tensorflow_asr.featurizers.text_featurizers import CharFeaturizer
 from tensorflow_asr.runners.base_runners import BaseTester
 from tensorflow_asr.models.conformer import Conformer
 
-config = UserConfig(DEFAULT_YAML, args.config, learning=True)
-speech_featurizer = TFSpeechFeaturizer(config["speech_config"])
-text_featurizer = CharFeaturizer(config["decoder_config"])
+config = Config(args.config, learning=True)
+speech_featurizer = TFSpeechFeaturizer(config.speech_config)
+text_featurizer = CharFeaturizer(config.decoder_config)
 
 tf.random.set_seed(0)
 assert args.saved
 
 if args.tfrecords:
     test_dataset = ASRTFRecordDataset(
-        data_paths=config["learning_config"]["dataset_config"]["test_paths"],
-        tfrecords_dir=config["learning_config"]["dataset_config"]["tfrecords_dir"],
+        data_paths=config.learning_config.dataset_config.test_paths,
+        tfrecords_dir=config.learning_config.dataset_config.tfrecords_dir,
         speech_featurizer=speech_featurizer,
         text_featurizer=text_featurizer,
         stage="test", shuffle=False
     )
 else:
     test_dataset = ASRSliceDataset(
-        data_paths=config["learning_config"]["dataset_config"]["test_paths"],
+        data_paths=config.learning_config.dataset_config.test_paths,
         speech_featurizer=speech_featurizer,
         text_featurizer=text_featurizer,
         stage="test", shuffle=False
     )
 
 # build model
-conformer = Conformer(
-    vocabulary_size=text_featurizer.num_classes,
-    **config["model_config"]
-)
+conformer = Conformer(**config.model_config, vocabulary_size=text_featurizer.num_classes)
 conformer._build(speech_featurizer.shape)
 conformer.load_weights(args.saved, by_name=True)
 conformer.summary(line_length=120)
 conformer.add_featurizers(speech_featurizer, text_featurizer)
 
 conformer_tester = BaseTester(
-    config=config["learning_config"]["running_config"],
+    config=config.learning_config.running_config,
     output_name=args.output_name
 )
 conformer_tester.compile(conformer)

examples/conformer/test_subword_conformer.py

@@ -55,19 +55,19 @@
 
 setup_devices([args.device], cpu=args.cpu)
 
-from tensorflow_asr.configs.user_config import UserConfig
+from tensorflow_asr.configs.config import Config
 from tensorflow_asr.datasets.asr_dataset import ASRTFRecordDataset, ASRSliceDataset
 from tensorflow_asr.featurizers.speech_featurizers import TFSpeechFeaturizer
 from tensorflow_asr.featurizers.text_featurizers import SubwordFeaturizer
 from tensorflow_asr.runners.base_runners import BaseTester
 from tensorflow_asr.models.conformer import Conformer
 
-config = UserConfig(DEFAULT_YAML, args.config, learning=True)
-speech_featurizer = TFSpeechFeaturizer(config["speech_config"])
+config = Config(args.config, learning=True)
+speech_featurizer = TFSpeechFeaturizer(config.speech_config)
 
 if args.subwords and os.path.exists(args.subwords):
     print("Loading subwords ...")
-    text_featurizer = SubwordFeaturizer.load_from_file(config["decoder_config"], args.subwords)
+    text_featurizer = SubwordFeaturizer.load_from_file(config.decoder_config, args.subwords)
 else:
     raise ValueError("subwords must be set")
 
@@ -76,32 +76,29 @@
 
 if args.tfrecords:
     test_dataset = ASRTFRecordDataset(
-        data_paths=config["learning_config"]["dataset_config"]["test_paths"],
-        tfrecords_dir=config["learning_config"]["dataset_config"]["tfrecords_dir"],
+        data_paths=config.learning_config.dataset_config.test_paths,
+        tfrecords_dir=config.learning_config.dataset_config.tfrecords_dir,
         speech_featurizer=speech_featurizer,
         text_featurizer=text_featurizer,
         stage="test", shuffle=False
     )
 else:
     test_dataset = ASRSliceDataset(
-        data_paths=config["learning_config"]["dataset_config"]["test_paths"],
+        data_paths=config.learning_config.dataset_config.test_paths,
         speech_featurizer=speech_featurizer,
         text_featurizer=text_featurizer,
         stage="test", shuffle=False
     )
 
 # build model
-conformer = Conformer(
-    vocabulary_size=text_featurizer.num_classes,
-    **config["model_config"]
-)
+conformer = Conformer(**config.model_config, vocabulary_size=text_featurizer.num_classes)
 conformer._build(speech_featurizer.shape)
 conformer.load_weights(args.saved, by_name=True)
 conformer.summary(line_length=120)
 conformer.add_featurizers(speech_featurizer, text_featurizer)
 
 conformer_tester = BaseTester(
-    config=config["learning_config"]["running_config"],
+    config=config.learning_config.running_config,
     output_name=args.output_name
 )
 conformer_tester.compile(conformer)

examples/conformer/tflite_conformer.py

@@ -19,7 +19,7 @@
 setup_environment()
 import tensorflow as tf
 
-from tensorflow_asr.configs.user_config import UserConfig
+from tensorflow_asr.configs.config import Config
 from tensorflow_asr.featurizers.speech_featurizers import TFSpeechFeaturizer
 from tensorflow_asr.featurizers.text_featurizers import CharFeaturizer
 from tensorflow_asr.models.conformer import Conformer
@@ -43,15 +43,12 @@
 
 assert args.saved and args.output
 
-config = UserConfig(DEFAULT_YAML, args.config, learning=True)
-speech_featurizer = TFSpeechFeaturizer(config["speech_config"])
-text_featurizer = CharFeaturizer(config["decoder_config"])
+config = Config(args.config, learning=True)
+speech_featurizer = TFSpeechFeaturizer(config.speech_config)
+text_featurizer = CharFeaturizer(config.decoder_config)
 
 # build model
-conformer = Conformer(
-    **config["model_config"],
-    vocabulary_size=text_featurizer.num_classes
-)
+conformer = Conformer(**config.model_config, vocabulary_size=text_featurizer.num_classes)
 conformer._build(speech_featurizer.shape)
 conformer.load_weights(args.saved)
 conformer.summary(line_length=150)

examples/conformer/tflite_subword_conformer.py

@@ -19,7 +19,7 @@
 setup_environment()
 import tensorflow as tf
 
-from tensorflow_asr.configs.user_config import UserConfig
+from tensorflow_asr.configs.config import Config
 from tensorflow_asr.featurizers.speech_featurizers import TFSpeechFeaturizer
 from tensorflow_asr.featurizers.text_featurizers import SubwordFeaturizer
 from tensorflow_asr.models.conformer import Conformer
@@ -46,20 +46,17 @@
 
 assert args.saved and args.output
 
-config = UserConfig(DEFAULT_YAML, args.config, learning=True)
-speech_featurizer = TFSpeechFeaturizer(config["speech_config"])
+config = Config(args.config, learning=True)
+speech_featurizer = TFSpeechFeaturizer(config.speech_config)
 
 if args.subwords and os.path.exists(args.subwords):
     print("Loading subwords ...")
-    text_featurizer = SubwordFeaturizer.load_from_file(config["decoder_config"], args.subwords)
+    text_featurizer = SubwordFeaturizer.load_from_file(config.decoder_config, args.subwords)
 else:
     raise ValueError("subwords must be set")
 
 # build model
-conformer = Conformer(
-    **config["model_config"],
-    vocabulary_size=text_featurizer.num_classes
-)
+conformer = Conformer(**config.model_config, vocabulary_size=text_featurizer.num_classes)
 conformer._build(speech_featurizer.shape)
 conformer.load_weights(args.saved)
 conformer.summary(line_length=150)

examples/conformer/train_conformer.py

@@ -56,69 +56,66 @@
 
 strategy = setup_strategy(args.devices)
 
-from tensorflow_asr.configs.user_config import UserConfig
+from tensorflow_asr.configs.config import Config
 from tensorflow_asr.datasets.asr_dataset import ASRTFRecordDataset, ASRSliceDataset
 from tensorflow_asr.featurizers.speech_featurizers import TFSpeechFeaturizer
 from tensorflow_asr.featurizers.text_featurizers import CharFeaturizer
 from tensorflow_asr.runners.transducer_runners import TransducerTrainer
 from tensorflow_asr.models.conformer import Conformer
 from tensorflow_asr.optimizers.schedules import TransformerSchedule
 
-config = UserConfig(DEFAULT_YAML, args.config, learning=True)
-speech_featurizer = TFSpeechFeaturizer(config["speech_config"])
-text_featurizer = CharFeaturizer(config["decoder_config"])
+config = Config(args.config, learning=True)
+speech_featurizer = TFSpeechFeaturizer(config.speech_config)
+text_featurizer = CharFeaturizer(config.decoder_config)
 
 if args.tfrecords:
     train_dataset = ASRTFRecordDataset(
-        data_paths=config["learning_config"]["dataset_config"]["train_paths"],
-        tfrecords_dir=config["learning_config"]["dataset_config"]["tfrecords_dir"],
+        data_paths=config.learning_config.dataset_config.train_paths,
+        tfrecords_dir=config.learning_config.dataset_config.tfrecords_dir,
         speech_featurizer=speech_featurizer,
         text_featurizer=text_featurizer,
-        augmentations=config["learning_config"]["augmentations"],
+        augmentations=config.learning_config.augmentations,
         stage="train", cache=args.cache, shuffle=True
     )
     eval_dataset = ASRTFRecordDataset(
-        data_paths=config["learning_config"]["dataset_config"]["eval_paths"],
-        tfrecords_dir=config["learning_config"]["dataset_config"]["tfrecords_dir"],
+        data_paths=config.learning_config.dataset_config.eval_paths,
+        tfrecords_dir=config.learning_config.dataset_config.tfrecords_dir,
         speech_featurizer=speech_featurizer,
         text_featurizer=text_featurizer,
         stage="eval", cache=args.cache, shuffle=True
     )
 else:
     train_dataset = ASRSliceDataset(
-        data_paths=config["learning_config"]["dataset_config"]["train_paths"],
+        data_paths=config.learning_config.dataset_config.train_paths,
         speech_featurizer=speech_featurizer,
         text_featurizer=text_featurizer,
-        augmentations=config["learning_config"]["augmentations"],
+        augmentations=config.learning_config.augmentations,
         stage="train", cache=args.cache, shuffle=True
     )
     eval_dataset = ASRSliceDataset(
-        data_paths=config["learning_config"]["dataset_config"]["eval_paths"],
+        data_paths=config.learning_config.dataset_config.eval_paths,
         speech_featurizer=speech_featurizer,
         text_featurizer=text_featurizer,
         stage="eval", cache=args.cache, shuffle=True
     )
 
 conformer_trainer = TransducerTrainer(
-    config=config["learning_config"]["running_config"],
+    config=config.learning_config.running_config,
     text_featurizer=text_featurizer, strategy=strategy
 )
 
 with conformer_trainer.strategy.scope():
     # build model
-    conformer = Conformer(
-        **config["model_config"],
-        vocabulary_size=text_featurizer.num_classes
-    )
+    conformer = Conformer(**config.model_config, vocabulary_size=text_featurizer.num_classes)
     conformer._build(speech_featurizer.shape)
     conformer.summary(line_length=120)
 
-    optimizer_config = config["learning_config"]["optimizer_config"]
+    optimizer_config = config.learning_config.optimizer_config
     optimizer = tf.keras.optimizers.Adam(
         TransformerSchedule(
-            d_model=config["model_config"]["dmodel"],
+            d_model=config.model_config["dmodel"],
             warmup_steps=optimizer_config["warmup_steps"],
-            max_lr=(0.05 / math.sqrt(config["model_config"]["dmodel"]))
+            max_lr=(0.05 / math.sqrt(config.model_config["dmodel"]))
         ),
         beta_1=optimizer_config["beta1"],
         beta_2=optimizer_config["beta2"],

examples/conformer/train_ga_conformer.py

@@ -56,69 +56,66 @@
 
 strategy = setup_strategy(args.devices)
 
-from tensorflow_asr.configs.user_config import UserConfig
+from tensorflow_asr.configs.config import Config
 from tensorflow_asr.datasets.asr_dataset import ASRTFRecordDataset, ASRSliceDataset
 from tensorflow_asr.featurizers.speech_featurizers import TFSpeechFeaturizer
 from tensorflow_asr.featurizers.text_featurizers import CharFeaturizer
 from tensorflow_asr.runners.transducer_runners import TransducerTrainerGA
 from tensorflow_asr.models.conformer import Conformer
 from tensorflow_asr.optimizers.schedules import TransformerSchedule
 
-config = UserConfig(DEFAULT_YAML, args.config, learning=True)
-speech_featurizer = TFSpeechFeaturizer(config["speech_config"])
-text_featurizer = CharFeaturizer(config["decoder_config"])
+config = Config(args.config, learning=True)
+speech_featurizer = TFSpeechFeaturizer(config.speech_config)
+text_featurizer = CharFeaturizer(config.decoder_config)
 
 if args.tfrecords:
     train_dataset = ASRTFRecordDataset(
-        data_paths=config["learning_config"]["dataset_config"]["train_paths"],
-        tfrecords_dir=config["learning_config"]["dataset_config"]["tfrecords_dir"],
+        data_paths=config.learning_config.dataset_config.train_paths,
+        tfrecords_dir=config.learning_config.dataset_config.tfrecords_dir,
         speech_featurizer=speech_featurizer,
         text_featurizer=text_featurizer,
-        augmentations=config["learning_config"]["augmentations"],
+        augmentations=config.learning_config.augmentations,
         stage="train", cache=args.cache, shuffle=True
     )
     eval_dataset = ASRTFRecordDataset(
-        data_paths=config["learning_config"]["dataset_config"]["eval_paths"],
-        tfrecords_dir=config["learning_config"]["dataset_config"]["tfrecords_dir"],
+        data_paths=config.learning_config.dataset_config.eval_paths,
+        tfrecords_dir=config.learning_config.dataset_config.tfrecords_dir,
         speech_featurizer=speech_featurizer,
         text_featurizer=text_featurizer,
         stage="eval", cache=args.cache, shuffle=True
     )
 else:
     train_dataset = ASRSliceDataset(
-        data_paths=config["learning_config"]["dataset_config"]["train_paths"],
+        data_paths=config.learning_config.dataset_config.train_paths,
         speech_featurizer=speech_featurizer,
         text_featurizer=text_featurizer,
-        augmentations=config["learning_config"]["augmentations"],
+        augmentations=config.learning_config.augmentations,
         stage="train", cache=args.cache, shuffle=True
     )
     eval_dataset = ASRSliceDataset(
-        data_paths=config["learning_config"]["dataset_config"]["eval_paths"],
+        data_paths=config.learning_config.dataset_config.eval_paths,
         speech_featurizer=speech_featurizer,
         text_featurizer=text_featurizer,
         stage="eval", cache=args.cache, shuffle=True
     )
 
 conformer_trainer = TransducerTrainerGA(
-    config=config["learning_config"]["running_config"],
+    config=config.learning_config.running_config,
     text_featurizer=text_featurizer, strategy=strategy
 )
 
 with conformer_trainer.strategy.scope():
     # build model
-    conformer = Conformer(
-        **config["model_config"],
-        vocabulary_size=text_featurizer.num_classes
-    )
+    conformer = Conformer(**config.model_config, vocabulary_size=text_featurizer.num_classes)
     conformer._build(speech_featurizer.shape)
     conformer.summary(line_length=120)
 
-    optimizer_config = config["learning_config"]["optimizer_config"]
+    optimizer_config = config.learning_config.optimizer_config
     optimizer = tf.keras.optimizers.Adam(
         TransformerSchedule(
-            d_model=config["model_config"]["dmodel"],
+            d_model=config.model_config["dmodel"],
             warmup_steps=optimizer_config["warmup_steps"],
-            max_lr=(0.05 / math.sqrt(config["model_config"]["dmodel"]))
+            max_lr=(0.05 / math.sqrt(config.model_config["dmodel"]))
         ),
         beta_1=optimizer_config["beta1"],
         beta_2=optimizer_config["beta2"],