Fix subset issue with multiple dictionaries

docs/source/changelog/changelog_2.0.rst

@@ -17,6 +17,7 @@ Release candidates
 - Fixed a bug where sample rate could not be specified when not using multiprocessing :github_pr:`444`
 - Fixed an incompatibility with the Kaldi version 1016 where BLAS libraries were not operating in single-threaded mode
 - Further optimized large multispeaker dictionary loading
+- Fixed a bug where subsets were not properly generated when multiple dictionaries were used
 
 2.0.0rc6
 --------

montreal_forced_aligner/corpus/base.py

@@ -706,7 +706,7 @@ def create_subset(self, subset: int) -> None:
                 subsets_per_dictionary = {}
                 utts_per_dictionary = {}
                 subsetted = 0
-                for dict_id in getattr(self, "dictionary_mapping", {}).keys():
+                for dict_id in getattr(self, "dictionary_lookup", {}).values():
                     num_utts = (
                         session.query(Utterance)
                         .join(Utterance.speaker)
@@ -720,7 +720,7 @@ def create_subset(self, subset: int) -> None:
                 remaining_subset = subset - sum(subsets_per_dictionary.values())
                 remaining_dicts = num_dictionaries - subsetted
                 remaining_subset_per_dictionary = int(remaining_subset / remaining_dicts)
-                for dict_id in getattr(self, "dictionary_mapping", {}).keys():
+                for dict_id in getattr(self, "dictionary_lookup", {}).values():
                     num_utts = utts_per_dictionary[dict_id]
                     if dict_id in subsets_per_dictionary:
                         subset_per_dictionary = subsets_per_dictionary[dict_id]

montreal_forced_aligner/dictionary/multispeaker.py

@@ -218,10 +218,7 @@ def reversed_word_mapping(self, dictionary_id: int = 1) -> Dict[int, str]:
     @property
     def num_dictionaries(self) -> int:
         """Number of pronunciation dictionaries"""
-        if self._num_dictionaries is None:
-            with self.session() as session:
-                self._num_dictionaries = session.query(Dictionary).count()
-        return self._num_dictionaries
+        return len(self.dictionary_lookup)
 
     @property
     def sanitize_function(self) -> MultispeakerSanitizationFunction:
@@ -297,19 +294,23 @@ def dictionary_setup(self):
         self._speaker_ids = getattr(self, "_speaker_ids", {})
         dictionary_id_cache = {}
         with self.session() as session:
-            for speaker_id, speaker_name, dictionary_id, path in (
-                session.query(Speaker.id, Speaker.name, Dictionary.id, Dictionary.path)
+            for speaker_id, speaker_name, dictionary_id, dict_name, path in (
+                session.query(
+                    Speaker.id, Speaker.name, Dictionary.id, Dictionary.name, Dictionary.path
+                )
                 .join(Speaker.dictionary)
                 .filter(Dictionary.default == False)  # noqa
             ):
                 self._speaker_ids[speaker_name] = speaker_id
                 dictionary_id_cache[path] = dictionary_id
+                self.dictionary_lookup[dict_name] = dictionary_id
             dictionary = (
                 session.query(Dictionary).filter(Dictionary.default == True).first()  # noqa
             )
             if dictionary:
                 self._default_dictionary_id = dictionary.id
                 dictionary_id_cache[dictionary.path] = self._default_dictionary_id
+                self.dictionary_lookup[dictionary.name] = dictionary.id
             word_primary_key = 1
             pronunciation_primary_key = 1
             word_objs = []