implement txt file format in `faster_whisper_inference.py`

modules/faster_whisper_inference.py

@@ -13,7 +13,7 @@ import torch
 import gradio as gr
 
 from .base_interface import BaseInterface
-from modules.subtitle_manager import get_srt, get_vtt, write_file, safe_filename
+from modules.subtitle_manager import get_srt, get_vtt, get_txt, write_file, safe_filename
 from modules.youtube_manager import get_ytdata, get_ytaudio
 
 
@@ -34,7 +34,7 @@ class FasterWhisperInference(BaseInterface):
                         fileobjs: list,
                         model_size: str,
                         lang: str,
-                        subformat: str,
+                        file_format: str,
                         istranslate: bool,
                         add_timestamp: bool,
                         beam_size: int,
@@ -54,8 +54,8 @@ class FasterWhisperInference(BaseInterface):
             Whisper model size from gr.Dropdown()
         lang: str
             Source language of the file to transcribe from gr.Dropdown()
-        subformat: str
-            Subtitle format to write from gr.Dropdown(). Supported format: [SRT, WebVTT]
+        file_format: str
+            File format to write from gr.Dropdown(). Supported format: [SRT, WebVTT, txt]
         istranslate: bool
             Boolean value from gr.Checkbox() that determines whether to translate to English.
             It's Whisper's feature to translate speech from another language directly into English end-to-end.
@@ -97,12 +97,13 @@ class FasterWhisperInference(BaseInterface):
 
                 file_name, file_ext = os.path.splitext(os.path.basename(fileobj.orig_name))
                 file_name = safe_filename(file_name)
-                subtitle = self.generate_and_write_subtitle(
+                subtitle = self.generate_and_write_file(
                     file_name=file_name,
                     transcribed_segments=transcribed_segments,
                     add_timestamp=add_timestamp,
-                    subformat=subformat
+                    file_format=file_format
                 )
+                print(f"{subtitle}")
                 files_info[file_name] = {"subtitle": subtitle, "time_for_task": time_for_task}
 
             total_result = ''
@@ -125,7 +126,7 @@ class FasterWhisperInference(BaseInterface):
                            youtubelink: str,
                            model_size: str,
                            lang: str,
-                           subformat: str,
+                           file_format: str,
                            istranslate: bool,
                            add_timestamp: bool,
                            beam_size: int,
@@ -145,8 +146,8 @@ class FasterWhisperInference(BaseInterface):
             Whisper model size from gr.Dropdown()
         lang: str
             Source language of the file to transcribe from gr.Dropdown()
-        subformat: str
-            Subtitle format to write from gr.Dropdown(). Supported format: [SRT, WebVTT]
+        file_format: str
+            File format to write from gr.Dropdown(). Supported format: [SRT, WebVTT, txt]
         istranslate: bool
             Boolean value from gr.Checkbox() that determines whether to translate to English.
             It's Whisper's feature to translate speech from another language directly into English end-to-end.
@@ -191,11 +192,11 @@ class FasterWhisperInference(BaseInterface):
             progress(1, desc="Completed!")
 
             file_name = safe_filename(yt.title)
-            subtitle = self.generate_and_write_subtitle(
+            subtitle = self.generate_and_write_file(
                 file_name=file_name,
                 transcribed_segments=transcribed_segments,
                 add_timestamp=add_timestamp,
-                subformat=subformat
+                file_format=file_format
             )
             return f"Done in {self.format_time(time_for_task)}! Subtitle file is in the outputs folder.\n\n{subtitle}"
         except Exception as e:
@@ -217,7 +218,7 @@ class FasterWhisperInference(BaseInterface):
                        micaudio: str,
                        model_size: str,
                        lang: str,
-                       subformat: str,
+                       file_format: str,
                        istranslate: bool,
                        beam_size: int,
                        log_prob_threshold: float,
@@ -236,8 +237,8 @@ class FasterWhisperInference(BaseInterface):
             Whisper model size from gr.Dropdown()
         lang: str
             Source language of the file to transcribe from gr.Dropdown()
-        subformat: str
-            Subtitle format to write from gr.Dropdown(). Supported format: [SRT, WebVTT]
+        file_format: str
+            File format to write from gr.Dropdown(). Supported format: [SRT, WebVTT, txt]
         istranslate: bool
             Boolean value from gr.Checkbox() that determines whether to translate to English.
             It's Whisper's feature to translate speech from another language directly into English end-to-end.
@@ -276,11 +277,11 @@ class FasterWhisperInference(BaseInterface):
             )
             progress(1, desc="Completed!")
 
-            subtitle = self.generate_and_write_subtitle(
+            subtitle = self.generate_and_write_file(
                 file_name="Mic",
                 transcribed_segments=transcribed_segments,
                 add_timestamp=True,
-                subformat=subformat
+                file_format=file_format
             )
             return f"Done in {self.format_time(time_for_task)}! Subtitle file is in the outputs folder.\n\n{subtitle}"
         except Exception as e:
@@ -378,11 +379,11 @@ class FasterWhisperInference(BaseInterface):
             )
 
     @staticmethod
-    def generate_and_write_subtitle(file_name: str,
-                                    transcribed_segments: list,
-                                    add_timestamp: bool,
-                                    subformat: str,
-                                    ) -> str:
+    def generate_and_write_file(file_name: str,
+                                transcribed_segments: list,
+                                add_timestamp: bool,
+                                file_format: str,
+                                ) -> str:
         """
         This method writes subtitle file and returns str to gr.Textbox
         """
@@ -392,13 +393,18 @@ class FasterWhisperInference(BaseInterface):
         else:
             output_path = os.path.join("outputs", f"{file_name}")
 
-        if subformat == "SRT":
-            subtitle = get_srt(transcribed_segments)
-            write_file(subtitle, f"{output_path}.srt")
-        elif subformat == "WebVTT":
-            subtitle = get_vtt(transcribed_segments)
-            write_file(subtitle, f"{output_path}.vtt")
-        return subtitle
+        if file_format == "SRT":
+            content = get_srt(transcribed_segments)
+            write_file(content, f"{output_path}.srt")
+
+        elif file_format == "WebVTT":
+            content = get_vtt(transcribed_segments)
+            write_file(content, f"{output_path}.vtt")
+
+        elif file_format == "txt":
+            content = get_txt(transcribed_segments)
+            write_file(content, f"{output_path}.txt")
+        return content
 
     @staticmethod
     def format_time(elapsed_time: float) -> str:

modules/subtitle_manager.py

@@ -44,6 +44,15 @@ def get_vtt(segments):
     return output
 
 
+def get_txt(segments):
+    output = ""
+    for i, segment in enumerate(segments):
+        if segment['text'].startswith(' '):
+            segment['text'] = segment['text'][1:]
+        output += f"{segment['text']}\n"
+    return output
+
+
 def parse_srt(file_path):
     """Reads SRT file and returns as dict"""
     with open(file_path, 'r', encoding='utf-8') as file: