debug piepline

Former-commit-id: 651fae66447937370d60073672ac834b0b2481de

configs/task_config.yaml

@@ -8,3 +8,4 @@ output_type:
 source_lang: EN
 target_lang: ZH
 field: SC2
+chunk_size: 1000

src/task.py

@@ -11,7 +11,7 @@ import subprocess
 from src.srt_util.srt import SrtScript
 from src.srt_util.srt2ass import srt2ass
 from time import time, strftime, gmtime, sleep
-from translation.translation import get_translation, translate
+from src.translators.translation import get_translation, translate
 
 import torch
 import stable_whisper
@@ -119,17 +119,15 @@ class Task:
         self.t_s = time()
         # self.SRT_Script = SrtScript
 
-        en_srt_path = self.task_local_dir.joinpath(f"task_{self.task_id})_en.srt")
-        if not Path.exists(en_srt_path):
+        src_srt_path = self.task_local_dir.joinpath(f"task_{self.task_id})_{self.source_lang}.srt")
+        if not Path.exists(src_srt_path):
             # extract script from audio
             logging.info("extract script from audio")
             device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-            logging.info("device: ", device)
-
-            audio_path = self.task_local_dir.joinpath(f"task_{self.task_id}.mp3")
+            # logging.info("device: ", device)
 
             if method == "api":
-                with open(audio_path, 'rb') as audio_file:
+                with open(self.audio_path, 'rb') as audio_file:
                     transcript = openai.Audio.transcribe(model="whisper-1", file=audio_file, response_format="srt")
             elif method == "stable":
                 model = stable_whisper.load_model(whisper_model, device)
@@ -147,11 +145,9 @@ class Task:
             # after get the transcript, release the gpu resource
             torch.cuda.empty_cache()
 
-        self.SRT_Script = SrtScript(transcript)
+        self.SRT_Script = SrtScript(transcript['segments'])
         # save the srt script to local
-        self.SRT_Script.write_srt_file_src(en_srt_path)
-        time.sleep(5)
-        pass
+        self.SRT_Script.write_srt_file_src(src_srt_path)
 
     # Module 2: SRT preprocess: perform preprocess steps
     # TODO: multi-lang and multi-field support according to task_cfg   
@@ -161,67 +157,67 @@ class Task:
         self.SRT_Script.form_whole_sentence()
         # self.SRT_Script.spell_check_term()
         self.SRT_Script.correct_with_force_term()
-        processed_srt_path_en = str(Path(self.task_local_dir).with_suffix('')) + '_processed.srt'
-        self.SRT_Script.write_srt_file_src(processed_srt_path_en)
+        processed_srt_path_src = str(Path(self.task_local_dir) / f'{self.task_id}_processed.srt')
+        self.SRT_Script.write_srt_file_src(processed_srt_path_src)
 
         if self.output_type["subtitle"] == "ass":
             logging.info("write English .srt file to .ass")
-            assSub_en = srt2ass(processed_srt_path_en)
-            logging.info('ASS subtitle saved as: ' + assSub_en)
+            assSub_src = srt2ass(processed_srt_path_src)
+            logging.info('ASS subtitle saved as: ' + assSub_src)
         self.script_input = self.SRT_Script.get_source_only()
         pass
     
     def update_translation_progress(self, new_progress):
         if self.progress == TaskStatus.TRANSLATING:
             self.progress = TaskStatus.TRANSLATING.value[0], new_progress
-        time.sleep(5)
 
     # Module 3: perform srt translation
     def translation(self):
         logging.info("---------------------Start Translation--------------------")
-        get_translation(self.srt,self.model, self.video_name, self.video_link)
-        time.sleep(5)
-        pass
+        get_translation(self.SRT_Script, self.model, self.task_id)
     
     # Module 4: perform srt post process steps
-    def postprocess(self, encode=False, srt_only=False):
+    def postprocess(self):
         self.status = TaskStatus.POST_PROCESSING
 
         logging.info("---------------------Start Post-processing SRT class---------------------")
         self.SRT_Script.check_len_and_split()
         self.SRT_Script.remove_trans_punctuation()
+        logging.info("---------------------Post-processing SRT class finished---------------------")
+
+    # Module 5: output module
+    def output_render(self):
+        self.status = TaskStatus.OUTPUT_MODULE
+        video_out = self.output_type["video"]
+        subtitle_type = self.output_type["subtitle"]
+        is_bilingal = self.output_type["bilingal"]
+
+        results_dir = Path(self.task_local_dir)/ "results"
 
-        base_path = Path(self.dir_result).joinpath(self.video_name).joinpath(self.video_name)
+        subtitle_path = f"{results_dir}/{self.task_id}_{self.target_lang}.srt"
+        self.SRT_Script.write_srt_file_translate(subtitle_path)
+        if is_bilingal:
+            subtitle_path = f"{results_dir}/{self.task_id}_{self.source_lang}_{self.target_lang}.srt"
+            self.SRT_Script.write_srt_file_bilingual(subtitle_path)
 
-        self.SRT_Script.write_srt_file_translate(f"{base_path}_zh.srt")
-        self.SRT_Script.write_srt_file_bilingual(f"{base_path}_bi.srt")
+        if subtitle_type == "ass":
+            logging.info("write .srt file to .ass")
+            subtitle_path = srt2ass(subtitle_path, "default", "No", "Modest")
+            logging.info('ASS subtitle saved as: ' + subtitle_path)
 
-        logging.info("write Chinese .srt file to .ass")
-        assSub_zh = srt2ass(f"{base_path}_zh.srt", "default", "No", "Modest")
-        logging.info('ASS subtitle saved as: ' + assSub_zh)
+        final_res = subtitle_path
 
         # encode to .mp4 video file
-        if encode:
+        if video_out and self.video_path is not None:
             logging.info("encoding video file")
-            if srt_only:
-                subprocess.run(
-                    f'ffmpeg -i {self.video_path} -vf "subtitles={base_path}_zh.srt" {base_path}.mp4')
-            else:
-                subprocess.run(
-                    f'ffmpeg -i {self.video_path} -vf "subtitles={base_path}_zh.ass" {base_path}.mp4')
+            subprocess.run(
+                f'ffmpeg -i {self.video_path} -vf "subtitles={subtitle_path}" {results_dir}/{self.task_id}.mp4')
+            final_res = f"{results_dir}/{self.task_id}.mp4"
 
         self.t_e = time()
         logging.info(
             "Pipeline finished, time duration:{}".format(strftime("%H:%M:%S", gmtime(self.t_e - self.t_s))))
-
-
-        time.sleep(5)
-        pass
-
-    # Module 5: output module
-    def output_render(self):
-        self.status = TaskStatus.OUTPUT_MODULE
-        return "TODO"
+        return final_res
     
     def run_pipeline(self):
         self.get_srt_class()
@@ -229,6 +225,7 @@ class Task:
         self.translation()
         self.postprocess()
         self.result = self.output_render()
+        print(self.result)
 
 class YoutubeTask(Task):
     def __init__(self, task_id, task_local_dir, task_cfg, youtube_url):

src/{translation → translators}/translation.py

@@ -3,11 +3,11 @@ import logging
 from time import sleep
 from tqdm import tqdm
 from src.srt_util.srt import split_script
-from LLM_task import LLM_task
+from .LLM_task import LLM_task
 
-def get_translation(srt,model,video_name,video_link):
-    script_arr, range_arr = split_script(srt)
-    translate(srt, script_arr, range_arr, model, video_name, video_link)
+def get_translation(srt, model, video_name):
+    script_arr, range_arr = split_script(srt.get_source_only())
+    translate(srt, script_arr, range_arr, model, video_name)
     pass
 
 def check_translation(sentence, translation):
@@ -25,7 +25,7 @@ def check_translation(sentence, translation):
         return True
 
 
-def translate(srt, script_arr, range_arr, model_name, video_name, video_link, attempts_count=5, task=None, temp = 0.15):
+def translate(srt, script_arr, range_arr, model_name, video_name, attempts_count=5, task=None, temp = 0.15):
     """
     Translates the given script array into another language using the chatgpt and writes to the SRT file.
 
@@ -38,7 +38,6 @@ def translate(srt, script_arr, range_arr, model_name, video_name, video_link, at
     :param range_arr: A list of tuples representing the start and end positions of sentences in the script.
     :param model_name: The name of the translation model to be used.
     :param video_name: The name of the video.
-    :param video_link: The link to the video.
     :param attempts_count: Number of attemps of failures for unmatched sentences.
     :param task: Prompt.
     :param temp: Model temperature.
@@ -60,10 +59,10 @@ def translate(srt, script_arr, range_arr, model_name, video_name, video_link, at
         while flag:
             flag = False
             try:
-                translate = LLM_task(model_name, sentence)
+                translate = LLM_task(model_name, sentence, task, temp)
                 # detect merge sentence issue and try to solve for five times:
                 while not check_translation(sentence, translate) and attempts_count > 0:
-                    translate = LLM_task(model_name,sentence,task,temp)
+                    translate = LLM_task(model_name, sentence, task, temp)
                     attempts_count -= 1
 
                 # if failure still happen, split into smaller tokens
@@ -85,4 +84,4 @@ def translate(srt, script_arr, range_arr, model_name, video_name, video_link, at
                 sleep(30)
                 flag = True
 
-        srt.set_translation(translate, range_, model_name, video_name, video_link)
+        srt.set_translation(translate, range_, model_name, video_name)