Simplify model creation for SpeakerExtraction and SpeechRecognition

DigitalPhonetics · Dec 19, 2023 · a7439c5 · a7439c5
1 parent 76928c5
commit a7439c5
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 9 deletions.
diff --git a/anonymization/modules/speaker_embeddings/speaker_extraction.py b/anonymization/modules/speaker_embeddings/speaker_extraction.py
@@ -44,10 +44,7 @@ def __init__(self, devices: list, settings: dict, results_dir: Path = None, mode
             'model_path': self.embed_model_path,
         }
 
-        if self.n_processes > 1:
-            self.extractors = None
-        else:
-            self.extractors = create_extractors(hparams=self.model_hparams, device=self.devices[0])
+        self.extractors = [create_extractors(hparams=self.model_hparams, device=device) for device, process in zip(cycle(devices), range(n_processes*len(devices)))]
 
     def extract_speakers(self, dataset_path, dataset_name=None, emb_level=None):
         dataset_name = dataset_name if dataset_name is not None else dataset_path.name

diff --git a/anonymization/modules/text/speech_recognition.py b/anonymization/modules/text/speech_recognition.py
@@ -33,11 +33,8 @@ def __init__(self, devices, settings, results_dir=None, save_intermediate=True,
             if self.save_intermediate:
                 raise ValueError('Results dir must be specified in parameters or settings!')
 
-        self.asr_model = create_model_instance(hparams=self.model_hparams, device=devices[0])
-        self.is_phones = (self.asr_model.output == 'phones')
-
-        if self.n_processes > 1:
-            self.asr_model = None
+        self.asr_models = [create_model_instance(hparams=self.model_hparams, device=device) for device, process in zip(cycle(devices), range(n_processes*len(devices)))]
+        self.is_phones = (self.asr_models[0].output == 'phones')
 
     def recognize_speech(self, dataset_path, dataset_name=None, utterance_list=None):
         dataset_name = dataset_name if dataset_name else dataset_path.name