adapt different languages for srt.py

Former-commit-id: 197fea8adb5bd35d60d2fc0d09ddc1af21cac117

src/srt_util/srt.py

@@ -8,9 +8,22 @@ import logging
 import openai
 from tqdm import tqdm
 
+punctuation_dict = {
+    "EN": ". , ? ! : ; - ( ) [ ] { } ' \"",
+    "ES": ". , ? ! : ; - ( ) [ ] { } ' \" ¡ ¿",
+    "FR": ". , ? ! : ; - ( ) [ ] { } ' \" « » —",
+    "DE": ". , ? ! : ; - ( ) [ ] { } ' \" „ “ –",
+    "RU": ". , ? ! : ; - ( ) [ ] { } ' \" « » —",
+    "ZH": "。 ， ？ ！ ： ； — （ ） ​``【oaicite:1】``​ 《 》 “ ”",
+    "JA": "。 、 ？ ！ ： ； ー （ ） ​``【oaicite:0】``​ 「 」 『 』",
+    "AR": ". , ? ! : ; - ( ) [ ] { } ، ؛ ؟ « »",
+}
 
 class SrtSegment(object):
-    def __init__(self, *args) -> None:
+    def __init__(self, src_lang, tgt_lang, *args) -> None:
+        self.src_lang = src_lang
+        self.tgt_lang = tgt_lang
+
         if isinstance(args[0], dict):
             segment = args[0]
             self.start = segment['start']
@@ -83,11 +96,11 @@ class SrtSegment(object):
 
     def remove_trans_punc(self) -> None:
         """
-        remove CN punctuations in translation text
+        remove punctuations in translation text
         :return: None
         """
-        punc_cn = "，。！？"
-        translator = str.maketrans(punc_cn, ' ' * len(punc_cn))
+        punc = punctuation_dict[self.tgt_lang]
+        translator = str.maketrans(punc, ' ' * len(punc))
         self.translation = self.translation.translate(translator)
 
     def __str__(self) -> str:
@@ -101,11 +114,13 @@ class SrtSegment(object):
 
 
 class SrtScript(object):
-    def __init__(self, segments) -> None:
-        self.segments = [SrtSegment(seg) for seg in segments]
+    def __init__(self, src_lang, tgt_lang, segments) -> None:
+        self.src_lang = src_lang
+        self.tgt_lang = tgt_lang
+        self.segments = [SrtSegment(self.src_lang, self.tgt_lang, seg) for seg in segments]
 
     @classmethod
-    def parse_from_srt_file(cls, path: str):
+    def parse_from_srt_file(cls, src_lang, tgt_lang, path: str):
         with open(path, 'r', encoding="utf-8") as f:
             script_lines = [line.rstrip() for line in f.readlines()]
         bilingual = False
@@ -119,7 +134,7 @@ class SrtScript(object):
             for i in range(0, len(script_lines), 4):
                 segments.append(list(script_lines[i:i + 4]))
 
-        return cls(segments)
+        return cls(src_lang, tgt_lang, segments)
 
     def merge_segs(self, idx_list) -> SrtSegment:
         """
@@ -309,14 +324,14 @@ class SrtScript(object):
         seg1_dict['text'] = src_seg1
         seg1_dict['start'] = start_seg1
         seg1_dict['end'] = end_seg1
-        seg1 = SrtSegment(seg1_dict)
+        seg1 = SrtSegment(self.src_lang, self.tgt_lang, seg1_dict)
         seg1.translation = trans_seg1
 
         seg2_dict = {}
         seg2_dict['text'] = src_seg2
         seg2_dict['start'] = start_seg2
         seg2_dict['end'] = end_seg2
-        seg2 = SrtSegment(seg2_dict)
+        seg2 = SrtSegment(self.src_lang, self.tgt_lang, seg2_dict)
         seg2.translation = trans_seg2
 
         result_list = []

src/task.py

@@ -129,10 +129,10 @@ class Task:
         # 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")
+        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")
@@ -157,7 +157,7 @@ class Task:
             # after get the transcript, release the gpu resource
             torch.cuda.empty_cache()
 
-        self.SRT_Script = SrtScript(transcript['segments'])
+        self.SRT_Script = SrtScript(self.source_lang, self.target_lang, transcript['segments'])
         # save the srt script to local
         self.SRT_Script.write_srt_file_src(src_srt_path)
 

src/translators/translation.py

@@ -18,8 +18,6 @@ def check_translation(sentence, translation):
     translation_count = translation.count('\n\n') + 1
 
     if sentence_count != translation_count:
-        # print("sentence length: ", len(sentence), sentence_count)
-        # print("translation length: ",  len(translation), translation_count)
         return False
     else:
         return True
@@ -34,7 +32,7 @@ def prompt_selector(src_lang, tgt_lang, domain):
     tgt_lang = language_map[tgt_lang]
     prompt = f"""
         you are a translation assistant, your job is to translate a video in domain of {domain} from {src_lang} to {tgt_lang}, 
-        you will be provided with a segement in {[src_lang]} parsed by line, where your translation text should keep the original 
+        you will be provided with a segement in {src_lang} parsed by line, where your translation text should keep the original 
         meaning and the number of lines.
         """
     return prompt
@@ -56,11 +54,11 @@ def translate(srt, script_arr, range_arr, model_name, video_name=None, attempts_
     :param task: Prompt.
     :param temp: Model temperature.
     """
-    #logging.info("Start translating...")
+
     if input is None: 
         raise Exception("Warning! No Input have passed to LLM!")
     if task is None:
-        task = "你是一个翻译助理，你的任务是翻译星际争霸视频，你会被提供一个按行分割的英文段落，你需要在保证句意和行数的情况下输出翻译后的文本。"
+        task = "你是一个翻译助理，你的任务是翻译视频，你会被提供一个按行分割的英文段落，你需要在保证句意和行数的情况下输出翻译后的文本。"
     logging.info(f"translation prompt: {task}")
     previous_length = 0
     for sentence, range_ in tqdm(zip(script_arr, range_arr)):