Merge pull request #74 from Camb-ai/level-bug-fix

added if clause for checking 'level-1'
m-bain · Feb 14, 2023 · 4cb167a · 4cb167a
2 parents d687cf3 + 2e30781
commit 4cb167a
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/whisperx/alignment.py b/whisperx/alignment.py
@@ -389,10 +389,11 @@ def align(
         wseg["segment-text-end"].fillna(len(text)-1, inplace=True)
 
         cseg = char_segments_arr.loc[seg_idx].loc[sub_start:sub_end-1]
-        cseg['segment-text-start'] = cseg['level_1']
-        cseg['segment-text-end'] = cseg['level_1'] + 1
-        del cseg['level_1']
-        del cseg['level_0']
+        # fixes bug for single segment in transcript
+        cseg['segment-text-start'] = cseg['level_1'] if 'level_1' in cseg else 0
+        cseg['segment-text-end'] = cseg['level_1'] + 1 if 'level_1' in cseg else 1
+        if 'level_1' in cseg: del cseg['level_1']
+        if 'level_0' in cseg: del cseg['level_0']
         cseg.reset_index(inplace=True)
         aligned_segments.append(
             {