[ASR] Add pretrained ASR models for Croatian (#4682)

* [ASR] Add pretrained ASR models for Croatian

Signed-off-by: Ante Jukić <[email protected]>

* Fix style for import

Signed-off-by: Ante Jukić <[email protected]>

Signed-off-by: Ante Jukić <[email protected]>
Co-authored-by: Ante Jukić <[email protected]>
Co-authored-by: Nithin Rao <[email protected]>
Co-authored-by: Eric Harper <[email protected]>
Co-authored-by: Somshubra Majumdar <[email protected]>

Author: 4 people

docs/source/asr/data/benchmark_hr.csv

@@ -0,0 +1,3 @@
+Model,Model Base Class,Model Card
+stt_hr_conformer_ctc_large,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_hr_conformer_ctc_large"
+stt_hr_conformer_transducer_large,EncDecRNNTBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_hr_conformer_transducer_large"

docs/source/asr/data/scores/hr/conformer_hr.csv

@@ -0,0 +1,3 @@
+Model Name,Language,ParlaSpeech-HR v1.0 (dev),ParlaSpeech-HR v1.0 (test)
+stt_hr_conformer_ctc_large,hr,4.43,4.70
+stt_hr_conformer_transducer_large,hr,4.56,4.69

docs/source/asr/scores.rst

@@ -169,6 +169,16 @@ FR
 
 --------------------
 
+HR
+^^
+
+.. csv-table::
+    :header-rows: 1
+    :align: left
+    :file: data/scores/hr/conformer_hr.csv
+    
+--------------------
+
 IT
 ^^
 

nemo/collections/asr/models/ctc_bpe_models.py

@@ -186,7 +186,6 @@ def _setup_transcribe_dataloader(self, config: Dict) -> 'torch.utils.data.DataLo
             'shuffle': False,
             'num_workers': config.get('num_workers', min(batch_size, os.cpu_count() - 1)),
             'pin_memory': True,
-            'channel_selector': config.get('channel_selector', None),
             'use_start_end_token': self.cfg.validation_ds.get('use_start_end_token', False),
         }
 
@@ -536,16 +535,9 @@ def list_available_models(cls) -> Optional[PretrainedModelInfo]:
         results.append(model)
 
         model = PretrainedModelInfo(
-            pretrained_model_name="stt_enes_conformer_ctc_large_codesw",
-            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_enes_conformer_ctc_large_codesw",
-            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_enes_conformer_ctc_large_codesw/versions/1.0.0/files/stt_enes_conformer_ctc_large_codesw.nemo",
-        )
-        results.append(model)
-
-        model = PretrainedModelInfo(
-            pretrained_model_name="stt_be_conformer_ctc_large",
-            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_be_conformer_ctc_large",
-            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_be_conformer_ctc_large/versions/1.12.0/files/stt_be_conformer_ctc_large.nemo",
+            pretrained_model_name="stt_hr_conformer_ctc_large",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_hr_conformer_ctc_large",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_hr_conformer_ctc_large/versions/1.11.0/files/stt_hr_conformer_ctc_large.nemo",
         )
         results.append(model)
 

nemo/collections/asr/models/rnnt_bpe_models.py

@@ -198,23 +198,9 @@ def list_available_models(cls) -> List[PretrainedModelInfo]:
         results.append(model)
 
         model = PretrainedModelInfo(
-            pretrained_model_name="stt_enes_conformer_transducer_large_codesw",
-            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_enes_conformer_transducer_large_codesw",
-            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_enes_conformer_transducer_large_codesw/versions/1.0.0/files/stt_enes_conformer_transducer_large_codesw.nemo",
-        )
-        results.append(model)
-
-        model = PretrainedModelInfo(
-            pretrained_model_name="stt_kab_conformer_transducer_large",
-            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_kab_conformer_transducer_large",
-            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_kab_conformer_transducer_large/versions/1.12.0/files/stt_kab_conformer_transducer_large.nemo",
-        )
-        results.append(model)
-
-        model = PretrainedModelInfo(
-            pretrained_model_name="stt_be_conformer_transducer_large",
-            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_be_conformer_transducer_large",
-            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_be_conformer_transducer_large/versions/1.12.0/files/stt_be_conformer_transducer_large.nemo",
+            pretrained_model_name="stt_hr_conformer_transducer_large",
+            description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_hr_conformer_transducer_large",
+            location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/stt_hr_conformer_transducer_large/versions/1.11.0/files/stt_hr_conformer_transducer_large.nemo",
         )
         results.append(model)
 
@@ -508,7 +494,6 @@ def _setup_transcribe_dataloader(self, config: Dict) -> 'torch.utils.data.DataLo
             'shuffle': False,
             'num_workers': config.get('num_workers', min(batch_size, os.cpu_count() - 1)),
             'pin_memory': True,
-            'channel_selector': config.get('channel_selector', None),
             'use_start_end_token': self.cfg.validation_ds.get('use_start_end_token', False),
         }
 

scripts/checkpoint_averaging/checkpoint_averaging.py

@@ -35,6 +35,7 @@
 import sys
 
 import torch
+from tqdm.auto import tqdm
 
 from nemo.core import ModelPT
 from nemo.utils import logging, model_utils
@@ -44,13 +45,14 @@ def main():
     parser = argparse.ArgumentParser()
     parser.add_argument(
         'model_fname_list',
-        metavar='N',
+        metavar='NEMO_FILE_OR_FOLDER',
         type=str,
         nargs='+',
         help='Input .nemo files (or folders who contains them) to parse',
     )
     parser.add_argument(
         '--import_fname_list',
+        metavar='FILE',
         type=str,
         nargs='+',
         default=[],
@@ -59,7 +61,7 @@ def main():
     args = parser.parse_args()
 
     logging.info(
-        f"\n\nIMPORTANT: Use --import_fname_list for all files that contain missing classes (AttributeError: Can't get attribute '???' on <module '__main__' from '???'>)\n\n"
+        f"\n\nIMPORTANT:\nIf you get the following error:\n\t(AttributeError: Can't get attribute '???' on <module '__main__' from '???'>)\nuse:\n\t--import_fname_list\nfor all files that contain missing classes.\n\n"
     )
 
     for fn in args.import_fname_list:
@@ -77,7 +79,7 @@ def main():
                 filter(lambda fn: not fn.endswith("-averaged.nemo"), glob.glob(os.path.join(model_fname, "*.nemo")))
             )
             if len(nemo_files) != 1:
-                raise RuntimeError(f"Expected only a single .nemo files but discovered {len(nemo_files)} .nemo files")
+                raise RuntimeError(f"Expected exactly one .nemo file but discovered {len(nemo_files)} .nemo files")
 
             model_fname = nemo_files[0]
 
@@ -107,23 +109,25 @@ def main():
 
         logging.info(f"Averaging {n} checkpoints ...")
 
-        for ix, path in enumerate(checkpoint_paths):
+        for ix, path in enumerate(tqdm(checkpoint_paths, total=n, desc='Averaging checkpoints')):
             checkpoint = torch.load(path, map_location=device)
 
             if 'state_dict' in checkpoint:
                 checkpoint = checkpoint['state_dict']
+            else:
+                raise RuntimeError(f"Checkpoint from {path} does not include a state_dict.")
 
             if ix == 0:
                 # Initial state
                 avg_state = checkpoint
 
-                logging.info(f"Initialized average state dict with checkpoint : {path}")
+                logging.info(f"Initialized average state dict with checkpoint:\n\t{path}")
             else:
                 # Accumulated state
                 for k in avg_state:
                     avg_state[k] = avg_state[k] + checkpoint[k]
 
-                logging.info(f"Updated average state dict with state from checkpoint : {path}")
+                logging.info(f"Updated average state dict with state from checkpoint:\n\t{path}")
 
         for k in avg_state:
             if str(avg_state[k].dtype).startswith("torch.int"):
@@ -136,7 +140,7 @@ def main():
         # restore merged weights into model
         nemo_model.load_state_dict(avg_state, strict=True)
         # Save model
-        logging.info(f"Saving average mdel to: {avg_model_fname}")
+        logging.info(f"Saving average model to:\n\t{avg_model_fname}")
         nemo_model.save_to(avg_model_fname)