Update dataloaders/beat_sep_single.py

dataloaders/beat_sep_single.py

@@ -340,31 +340,41 @@ class CustomDataset(Dataset):
         if self.args.word_rep is not None:
             logger.info(f"# ---- Building cache for Word   {id_pose} and Pose {id_pose} ---- #")
             word_file = self.textgrid_file_path
-            if not os.path.exists(word_file):
-                logger.warning(f"# ---- file not found for Word   {id_pose}, skip all files with the same id ---- #")
-                self.selected_file = self.selected_file.drop(self.selected_file[self.selected_file['id'] == id_pose].index)
             word_save_path = f"{self.data_dir}{self.args.t_pre_encoder}/{id_pose}.npy"
 
-            tgrid = tg.TextGrid.fromFile(word_file)
+            def _fallback_word_tokens(length: int) -> np.ndarray:
+                token = getattr(self.lang_model, "PAD_token", 0) if self.lang_model else 0
+                return np.full((length,), token, dtype=np.int64)
 
-            for i in range(pose_each_file.shape[0]):
-                found_flag = False
-                current_time = i/self.args.pose_fps + time_offset
-                j_last = 0
-                for j, word in enumerate(tgrid[0]): 
-                    word_n, word_s, word_e = word.mark, word.minTime, word.maxTime
-                    if word_s<=current_time and current_time<=word_e:
-                        if word_n == " ":
-                            word_each_file.append(self.lang_model.PAD_token)
-                        else:
-                            word_each_file.append(self.lang_model.get_word_index(word_n))
-                        found_flag = True
-                        j_last = j
-                        break
-                    else: continue   
-                if not found_flag: 
-                    word_each_file.append(self.lang_model.UNK_token)
-            word_each_file = np.array(word_each_file)
+            if not os.path.exists(word_file):
+                logger.warning(
+                    f"# ---- TextGrid not found for Word   {id_pose}; using fallback tokens ---- #"
+                )
+                word_each_file = _fallback_word_tokens(pose_each_file.shape[0])
+            else:
+                try:
+                    tgrid = tg.TextGrid.fromFile(word_file)
+                except Exception as exc:
+                    logger.warning(
+                        f"# ---- Failed to load TextGrid for Word   {id_pose}: {exc}; using fallback tokens ---- #"
+                    )
+                    word_each_file = _fallback_word_tokens(pose_each_file.shape[0])
+                else:
+                    for i in range(pose_each_file.shape[0]):
+                        found_flag = False
+                        current_time = i/self.args.pose_fps + time_offset
+                        for word in tgrid[0]:
+                            word_n, word_s, word_e = word.mark, word.minTime, word.maxTime
+                            if word_s <= current_time <= word_e:
+                                if word_n == " ":
+                                    word_each_file.append(self.lang_model.PAD_token)
+                                else:
+                                    word_each_file.append(self.lang_model.get_word_index(word_n))
+                                found_flag = True
+                                break
+                        if not found_flag:
+                            word_each_file.append(self.lang_model.UNK_token)
+                    word_each_file = np.array(word_each_file)