Fix segmentation using binary search

src/segment.py

@@ -119,8 +119,8 @@ def generate_segments(words, tokenizer, segmentation_args):
 def extract_segment(words, start, end, map_function=None):
     """Extracts all words with time in [start, end]"""
 
-    a = binary_search_below(words, 0, len(words) - 1, start)
-    b = min(binary_search_above(words, 0, len(words) - 1, end) + 1, len(words))
+    a = max(binary_search_below(words, 0, len(words), start), 0)
+    b = min(binary_search_above(words, -1, len(words) -1, end) + 1, len(words))
 
     to_transform = map_function is not None and callable(map_function)
 
@@ -153,7 +153,7 @@ def binary_search_above(transcript, start_index, end_index, time):
 
     middle_index = (start_index + end_index + 1) // 2
     middle = transcript[middle_index]
-    middle_time = avg(middle['start'], middle['end'])
+    middle_time = avg(word_start(middle), word_end(middle))
 
     if time >= middle_time:
         return binary_search_above(transcript, middle_index, end_index, time)