Update modules/whisper/whisper_base.py

modules/whisper/whisper_base.py

@@ -236,7 +236,7 @@ class WhisperBase(ABC):
     def transcribe_file(self,
                         files: Optional[List] = None,
                         input_folder_path: Optional[str] = None,
-                        file_format: str = "SRT",
+                        file_format: list = ["CSV"],
                         add_timestamp: bool = True,
                         translate_output: bool = False,
                         translate_model: str = "",
@@ -359,39 +359,18 @@ class WhisperBase(ABC):
                     subtitle = get_plaintext(transcribed_segments)
                 files_info[file_name] = {"subtitle": subtitle, "time_for_task": time_for_task, "lang": file_language, "lang_prob": file_lang_probs, "input_source_file": (file_name+file_ext), "translation": translation_note, "transcription": transcription_note}
                 
-                ## Add output file as txt
-                file_name, file_ext = os.path.splitext(os.path.basename(file))
-                subtitle, file_path = self.generate_and_write_file(
-                    file_name=file_name,
-                    transcribed_segments=transcribed_segments,
-                    add_timestamp=add_timestamp,
-                    file_format="txt",
-                    output_dir=self.output_dir
-                )
-                files_to_download[file_name+"_txt"] = {"path": file_path}
+                ## Add output file as txt, srt and/or csv
+                for format in file_format:
+                    file_name, file_ext = os.path.splitext(os.path.basename(file))
+                    subtitle, file_path = self.generate_and_write_file(
+                        file_name=file_name,
+                        transcribed_segments=transcribed_segments,
+                        add_timestamp=add_timestamp,
+                        file_format=format.lower(),
+                        output_dir=self.output_dir
+                    )
+                    files_to_download[file_name+"_"+format.lower()] = {"path": file_path}
                 
-                ## Add output file as srt
-                file_name, file_ext = os.path.splitext(os.path.basename(file))
-                subtitle, file_path = self.generate_and_write_file(
-                    file_name=file_name,
-                    transcribed_segments=transcribed_segments,
-                    add_timestamp=add_timestamp,
-                    file_format="srt",
-                    output_dir=self.output_dir
-                )
-                files_to_download[file_name+"_srt"] = {"path": file_path}
-
-                ## Add output file as csv
-                file_name, file_ext = os.path.splitext(os.path.basename(file))
-                subtitle, file_path = self.generate_and_write_file(
-                    file_name=file_name,
-                    transcribed_segments=transcribed_segments,
-                    add_timestamp=add_timestamp,
-                    file_format="csv",
-                    output_dir=self.output_dir
-                )
-                files_to_download[file_name+"_csv"] = {"path": file_path}
-            
             total_result = ""
             total_info = ""
             total_time = 0