debugs

Former-commit-id: 2d7d950f54b1deb8b4dd9b68c98d65384954c47e

configs/task_config.yaml

@@ -1,18 +1,35 @@
 # configuration for each task
-model: gpt-4
-# output type that user receive
-output_type: 
-  subtitle: srt
-  video: False
-  bilingal: False
 source_lang: EN
 target_lang: ZH
 field: SC2
-chunk_size: 1000
+
+# ASR config
+ASR:
+  ASR_model: whisper
+  whisper_config:
+    whisper_model: tiny
+    method: stable
+  
+# pre-process module config
 pre_process: 
-  ON: True
   sentence_form: True
   spell_check: False
   term_correct: True
+
+# Translation module config
+translation:
+  model: gpt-4
+  chunk_size: 1000
+
+# post-process module config
 post_process: 
-  ON: True
+  check_len_and_split: True
+  remove_trans_punctuation: True
+
+# output type that user receive
+output_type: 
+  subtitle: srt
+  video: False
+  bilingal: False
+
+

entries/run.py

@@ -10,6 +10,13 @@ from datetime import datetime
 import shutil
 from uuid import uuid4
 
+"""
+    Main entry for terminal environment.
+    Use it for debug and development purpose. 
+    Usage: python3 entries/run.py [-h] [--link LINK] [--video_file VIDEO_FILE] [--audio_file AUDIO_FILE] [--srt_file SRT_FILE] [--continue CONTINUE]
+              [--launch_cfg LAUNCH_CFG] [--task_cfg TASK_CFG]
+"""
+
 def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument("--link", help="youtube video link here", default=None, type=str, required=False)
@@ -42,8 +49,9 @@ if __name__ == "__main__":
     task_dir.mkdir(parents=False, exist_ok=False)
     task_dir.joinpath("results").mkdir(parents=False, exist_ok=False)
 
-    # logging
-    logging.basicConfig(level=logging.INFO, handlers=[
+    # logging setting
+    logfmt = "%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s"
+    logging.basicConfig(level=logging.INFO, format=logfmt, handlers=[
         logging.FileHandler(
             "{}/{}_{}.log".format(task_dir, f"task_{task_id}", datetime.now().strftime("%m%d%Y_%H%M%S")),
             'w', encoding='utf-8')])

src/srt_util/srt.py

@@ -185,7 +185,6 @@ class SrtScript(object):
 
         def inner_func(target, input_str):
             response = openai.ChatCompletion.create(
-                # model=model,
                 model="gpt-4",
                 messages=[
                     {"role": "system",
@@ -208,19 +207,13 @@ class SrtScript(object):
                 flag = True
                 while flag:
                     flag = False
-                    # print("translate:")
-                    # print(translate)
                     try:
-                        # print("target")
-                        # print(end_seg_id - start_seg_id + 1)
                         translate = inner_func(end_seg_id - start_seg_id + 1, translate)
                     except Exception as e:
                         print("An error has occurred during solving unmatched lines:", e)
                         print("Retrying...")
                         flag = True
                 lines = translate.split('\n')
-                # print("result")
-                # print(len(lines))
 
             if len(lines) < (end_seg_id - start_seg_id + 1):
                 solved = False
@@ -264,6 +257,7 @@ class SrtScript(object):
         # evenly split seg to 2 parts and add new seg into self.segments
 
         # ignore the initial comma to solve the recursion problem
+        # FIXME: accomodate multilingual setting
         if len(seg.source_text) > 2:
             if seg.source_text[:2] == ', ':
                 seg.source_text = seg.source_text[2:]

src/task.py

@@ -55,7 +55,6 @@ class TaskStatus(str, Enum):
     OUTPUT_MODULE = 'OUTPUT_MODULE'
 
 
-
 class Task:
     @property
     def status(self):
@@ -70,69 +69,74 @@ class Task:
     def __init__(self, task_id, task_local_dir, task_cfg):
         self.__status_lock = threading.Lock()
         self.__status = TaskStatus.CREATED
+        self.gpu_status = 0
         openai.api_key = getenv("OPENAI_API_KEY")
-        self.launch_info = task_cfg # do not use, just for fallback
+        self.task_id = task_id
+        
         self.task_local_dir = task_local_dir
-        self.model = task_cfg["model"]
-        self.gpu_status = 0
+        self.ASR_setting = task_cfg["ASR"]
+        self.translation_setting = task_cfg["translation"]
+        self.translation_model = self.translation_setting["model"]
+        
         self.output_type = task_cfg["output_type"]
         self.target_lang = task_cfg["target_lang"]
         self.source_lang = task_cfg["source_lang"]
         self.field = task_cfg["field"]
         self.pre_setting = task_cfg["pre_process"]
         self.post_setting = task_cfg["post_process"]
-        self.task_id = task_id
+        
         self.audio_path = None
         self.SRT_Script = None
         self.result = None
         self.s_t = None
         self.t_e = None
 
-        print(f" Task ID: {self.task_id}")
-        logging.info(f" Task ID: {self.task_id}")
-        logging.info(f" {self.source_lang} -> {self.target_lang} task in {self.field}")
-        logging.info(f" Model: \t\t\t{self.model}")
-        logging.info(f" subtitle_type: \t\t{self.output_type['subtitle']}")
-        logging.info(f" video_ouput: \t\t{self.output_type['video']}")
-        logging.info(f" bilingal_ouput: \t{self.output_type['bilingal']}")
-        logging.info(" PREprocess setting:")
-        for key, value in self.pre_setting:
-            logging.info(f"  {key}: {value}")
-        logging.info(" POSTprocess setting:")
-        for key, value in self.post_setting:
-            logging.info(f"  {key}: {value}")
+        print(f"Task ID: {self.task_id}")
+        logging.info(f"Task ID: {self.task_id}")
+        logging.info(f"{self.source_lang} -> {self.target_lang} task in {self.field}")
+        logging.info(f"Translation Model: {self.translation_model}")
+        logging.info(f"subtitle_type: {self.output_type['subtitle']}")
+        logging.info(f"video_ouput: {self.output_type['video']}")
+        logging.info(f"bilingal_ouput: {self.output_type['bilingal']}")
+        logging.info("Pre-process setting:")
+        for key in self.pre_setting:
+            logging.info(f"{key}: {self.pre_setting[key]}")
+        logging.info("Post-process setting:")
+        for key in self.post_setting:
+            logging.info(f"{key}: {self.post_setting[key]}")
 
     @staticmethod
     def fromYoutubeLink(youtube_url, task_id, task_dir, task_cfg):
         # convert to audio
-        logging.info(" Task Creation method: Youtube Link")
+        logging.info("Task Creation method: Youtube Link")
         return YoutubeTask(task_id, task_dir, task_cfg, youtube_url)
 
     @staticmethod
     def fromAudioFile(audio_path, task_id, task_dir, task_cfg):
         # get audio path
-        logging.info(" Task Creation method: Audio File")
+        logging.info("Task Creation method: Audio File")
         return AudioTask(task_id, task_dir, task_cfg, audio_path)
     
     @staticmethod
     def fromVideoFile(video_path, task_id, task_dir, task_cfg):
         # get audio path
-        logging.info(" Task Creation method: Video File")
+        logging.info("Task Creation method: Video File")
         return VideoTask(task_id, task_dir, task_cfg, video_path)
     
     # Module 1 ASR: audio --> SRT_script
-    def get_srt_class(self, whisper_model='tiny', method="stable"):
+    def get_srt_class(self):
         # Instead of using the script_en variable directly, we'll use script_input
+        # TODO: setup ASR module like translator
         self.status = TaskStatus.INITIALIZING_ASR
         self.t_s = time()
         # self.SRT_Script = SrtScript
-
+        method = self.ASR_setting["whisper_config"]["method"]
+        whisper_model = self.ASR_setting["whisper_config"]["whisper_model"]
         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)
 
             if method == "api":
                 with open(self.audio_path, 'rb') as audio_file:
@@ -158,7 +162,6 @@ class Task:
         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   
     def preprocess(self):
         self.status = TaskStatus.PRE_PROCESSING
         logging.info("--------------------Start Preprocessing SRT class--------------------")
@@ -183,18 +186,20 @@ class Task:
             self.progress = TaskStatus.TRANSLATING.value[0], new_progress
 
     # Module 3: perform srt translation
-    def translation(self,task_cfg):
+    def translation(self):
         logging.info("---------------------Start Translation--------------------")
-        prompt = prompt_selector(self.source_lang,self.target_lang,task_cfg['field'])
-        get_translation(self.SRT_Script, self.model, self.task_id, prompt, task_cfg['chunk_size'])
+        prompt = prompt_selector(self.source_lang, self.target_lang, self.field)
+        get_translation(self.SRT_Script, self.translation_model, self.task_id, prompt, self.translation_setting['chunk_size'])
     
     # Module 4: perform srt post process steps
     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()
+        if self.post_setting["check_len_and_split"]:
+            self.SRT_Script.check_len_and_split()
+        if self.post_setting["remove_trans_punctuation"]:
+            self.SRT_Script.remove_trans_punctuation()
         logging.info("---------------------Post-processing SRT class finished---------------------")
 
     # Module 5: output module
@@ -233,11 +238,9 @@ class Task:
     
     def run_pipeline(self):
         self.get_srt_class()
-        if self.pre_setting["ON"]:
-            self.preprocess()
+        self.preprocess()
         self.translation()
-        if self.post_setting["ON"]:
-            self.postprocess()
+        self.postprocess()
         self.result = self.output_render()
         print(self.result)
 
@@ -259,7 +262,6 @@ class YoutubeTask(Task):
         audio = yt.streams.filter(only_audio=True).first()
         if audio:
             audio.download(str(self.task_local_dir), filename=f"task_{self.task_id}.mp3")
-            # logging.info(f'Audio download completed to {self.task_local_dir}!')
         else:
             logging.info(" download audio failed, using ffmpeg to extract audio")
             subprocess.run(

src/translators/translation.py

@@ -5,9 +5,9 @@ from tqdm import tqdm
 from src.srt_util.srt import split_script
 from .LLM_task import LLM_task
 
-def get_translation(srt, model, video_name, task, chunk_size = 1000):
+def get_translation(srt, model, video_name, prompt, chunk_size = 1000):
     script_arr, range_arr = split_script(srt.get_source_only(),chunk_size)
-    translate(srt, script_arr, range_arr, model, video_name, task)
+    translate(srt, script_arr, range_arr, model, video_name, task=prompt)
     pass
 
 def check_translation(sentence, translation):
@@ -39,7 +39,7 @@ def prompt_selector(src_lang, tgt_lang, domain):
         """
     return prompt
 
-def translate(srt, script_arr, range_arr, model_name, video_name, attempts_count=5, task=None, temp = 0.15):
+def translate(srt, script_arr, range_arr, model_name, video_name=None, 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.
 
@@ -61,14 +61,14 @@ def translate(srt, script_arr, range_arr, model_name, video_name, attempts_count
         raise Exception("Warning! No Input have passed to LLM!")
     if task is None:
         task = "你是一个翻译助理，你的任务是翻译星际争霸视频，你会被提供一个按行分割的英文段落，你需要在保证句意和行数的情况下输出翻译后的文本。"
-    print(task)
+    logging.info(f"translation prompt: {task}")
     previous_length = 0
     for sentence, range_ in tqdm(zip(script_arr, range_arr)):
         # update the range based on previous length
         range_ = (range_[0] + previous_length, range_[1] + previous_length)
         # using chatgpt model
         print(f"now translating sentences {range_}")
-        #logging.info(f"now translating sentences {range_}, time: {datetime.now()}")
+        logging.info(f"now translating sentences {range_}")
         flag = True
         while flag:
             flag = False