Update modules/whisper/whisper_base.py

modules/whisper/whisper_base.py

@@ -424,7 +424,7 @@ class WhisperBase(ABC):
             result_file_path = [info['path'] for info in files_to_download.values()]
             #total_info += f"Total processing time:\t{self.format_time((time_end-time_start).total_seconds())}"
 
-            return [gr.update(value={"data": self.transform_text_to_list(total_result)}),result_file_path,total_info]
+            return [gr.update(value={"data": self.transform_text_to_list(total_result,add_timestamp_preview,diarize_speakers)}),result_file_path,total_info]
             #return [result_str,result_file_path,total_info]
         
         except Exception as e:
@@ -618,17 +618,18 @@ class WhisperBase(ABC):
         gc.collect()
 
     @staticmethod
-    def transform_text_to_list(inputdata: str) -> list:
+    def transform_text_to_list(inputdata: str,timestamps: bool, diarize: bool) -> list:
         outputdata = []
         temp_inputdata = inputdata.split("#NEWLINE#")
         for temp_line in temp_inputdata:
             temp_line_list = []	
             temp_line_items = temp_line.split("#TAB#")
 
-            # Add tabs to match Dataframe columns
-            count_items = len(temp_line_items)
-            while len(temp_line_items)<3:
-                temp_line_items.insert(0,"")
+            # Add columns to match Dataframe
+            if timestamps==True and diarize==False:
+                temp_line_items.insert(1,"")
+            if timestamps==False and diarize==True:
+                temp_line_items.insert(0,"")     
             
             for temp_line_item in temp_line_items:
                 temp_line_list.append(temp_line_item)