Update modules/whisper/whisper_base.py

modules/whisper/whisper_base.py

@@ -428,11 +428,32 @@ 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())}"
 
-            result_str = self.transform_text_to_list(total_result) #Transform to gr.Dataframe format
-            new_headers = ["New Header 1", "New Header 2"]
-            new_data_rows = [[1, 2], [3, 4]]
-            updated_data = {"headers": new_headers, "data": new_data_rows}
-    
+            # Format Dataframe data    
+            new_headers = []
+            new_column_widths = []
+            new_datatype = []
+
+            if add_timestamp_preview==True and diarize_speakers==True:
+                new_headers = ["Time","Speaker","Text"]
+                new_column_widths = ["10%","15%","70%"]
+                new_datatype = ["str","str","markdown"]
+            elif add_timestamp_preview==True and diarize_speakers==False:
+                new_headers = ["Time","Text"]
+                new_column_widths = ["10%","85%"]
+                new_datatype = ["str","markdown"]
+            elif add_timestamp_preview==False and diarize_speakers==True:
+                new_headers = ["Speaker","Text"]
+                new_column_widths = ["15%","80%"]
+                new_datatype = ["str","markdown"]
+            elif add_timestamp_preview==False and diarize_speakers==False:
+                new_headers = ["Text"]
+                new_column_widths = ["95%"]
+                new_datatype = ["markdown"]
+
+            new_data_rows = self.transform_text_to_list(total_result)
+            #updated_data = {"headers": new_headers, "data": new_data_rows}
+            updated_data = {"headers": new_headers, "data": new_data_rows, "column_widths": new_column_widths, "datatype": new_datatype}
+            
             return [(gr.update(value=updated_data)),result_file_path,total_info]
             #return [result_str,result_file_path,total_info]