Merge branch 'eason/refactor' of github.com:project-kxkg/project-t into eason/refactor

.gitignore

@@ -5,3 +5,5 @@
 test.py
 test.srt
 test.txt
+log_*.csv
+log.csv

SRT.py

@@ -3,6 +3,7 @@ from csv import reader
 from datetime import datetime
 import re
 import openai
+import os
 from collections import deque
 
 class SRT_segment(object):
@@ -50,9 +51,18 @@ class SRT_segment(object):
         self.source_text += seg.source_text
         self.translation += seg.translation
         self.end_time_str = seg.end_time_str
+        self.end = seg.end
+        self.end_ms = seg.end_ms
         self.duration = f"{self.start_time_str} --> {self.end_time_str}"
         pass
-    
+
+    def remove_trans_punc(self):
+        # remove punctuations in translation text
+        self.translation = self.translation.replace('，', ' ')
+        self.translation = self.translation.replace('。', ' ')
+        self.translation = self.translation.replace('！', ' ')
+        self.translation = self.translation.replace('？', ' ')
+
     def __str__(self) -> str:
         return  f'{self.duration}\n{self.source_text}\n\n'
     
@@ -62,16 +72,6 @@ class SRT_segment(object):
     def get_bilingual_str(self) -> str:
         return f'{self.duration}\n{self.source_text}\n{self.translation}\n\n'
 
-    # def set_translation(self, trans):
-    #     if trans[0] == '，':
-    #         trans = trans[1:]
-    #     self.translation = trans
-
-    # def set_src_text(self, src_text):
-    #     if src_text[0] == ',':
-    #         src_text = src_text[1:]
-    #     self.source_text = src_text
-
 class SRT_script():
     def __init__(self, segments) -> None:
         self.segments = []
@@ -105,7 +105,7 @@ class SRT_script():
         merge_list = [] # a list of indices that should be merged e.g. [[0], [1, 2, 3, 4], [5, 6], [7]]
         sentence = []
         for i, seg in enumerate(self.segments):
-            if seg.source_text[-1] == '.':
+            if seg.source_text[-1] in ['.', '!', '?']:
                 sentence.append(i)
                 merge_list.append(sentence)
                 sentence = []
@@ -117,56 +117,88 @@ class SRT_script():
             segments.append(self.merge_segs(idx_list))
 
         self.segments = segments # need memory release?
-        
     
+    def remove_trans_punctuation(self):
+        # Post-process: remove all punc after translation and split
+        for i, seg in enumerate(self.segments):
+            seg.remove_trans_punc()
 
-    def set_translation(self, translate:str, id_range:tuple, model):
+    def set_translation(self, translate:str, id_range:tuple, model, video_name, video_link=None):
         start_seg_id = id_range[0]
         end_seg_id = id_range[1]
         
-        def inner_func(input_str):
+        def inner_func(target,input_str):
             response = openai.ChatCompletion.create(
-                    model=model,
+                    #model=model,
+                    model = "gpt-3.5-turbo",
                     messages = [
-                        {"role": "system", "content": "You are a helpful assistant that help calibrates English to Chinese subtitle translations in starcraft2."},
-                        {"role": "system", "content": "You are provided with a translated Chinese transcript; you must modify or split the Chinese sentence to match the meaning and the number of the English transcript exactly one by one. You must not merge ANY Chinese lines, you can only split them but the total Chinese lines MUST equals to number of English lines."},
-                        {"role": "system", "content": "There is no need for you to add any comments or notes, and do not modify the English transcript."},
-                        {"role": "user", "content": 'You are given the English transcript and line number, your task is to merge or split the Chinese to match the exact number of lines in English transcript, no more no less. For example, if there are more Chinese lines than English lines, merge some the Chinese lines to match the number of English lines. If Chinese lines is less than English lines, split some Chinese lines to match the english lines: "{}"'.format(input_str)}
+                        #{"role": "system", "content": "You are a helpful assistant that help calibrates English to Chinese subtitle translations in starcraft2."},
+                        #{"role": "system", "content": "You are provided with a translated Chinese transcript; you must modify or split the Chinese sentence to match the meaning and the number of the English transcript exactly one by one. You must not merge ANY Chinese lines, you can only split them but the total Chinese lines MUST equals to number of English lines."},
+                        #{"role": "system", "content": "There is no need for you to add any comments or notes, and do not modify the English transcript."},
+                        #{"role": "user", "content": 'You are given the English transcript and line number, your task is to merge or split the Chinese to match the exact number of lines in English transcript, no more no less. For example, if there are more Chinese lines than English lines, merge some the Chinese lines to match the number of English lines. If Chinese lines is less than English lines, split some Chinese lines to match the english lines: "{}"'.format(input_str)}
+                        {"role": "system", "content": "你的任务是按照要求合并或拆分句子到指定行数，你需要尽可能保证句意，但必要时可以将一句话分为两行输出"},
+                        {"role": "system", "content": "注意：你只需要输出处理过的中文句子，如果你要输出序号，请使用冒号隔开"},
+                        {"role": "user", "content": '请将下面的句子拆分或组合为{}句:\n{}'.format(target,input_str)}
                     ],
-                    temperature=0.7
+                    #temperature=0.7
+                    temperature = 0.15
                 )
             return response['choices'][0]['message']['content'].strip()
         
+        
         lines = translate.split('\n\n')
         if len(lines) < (end_seg_id - start_seg_id + 1):
             count = 0
+            solved = True
             while count<5 and len(lines) != (end_seg_id - start_seg_id + 1):
-                
                 count += 1
                 print("Solving Unmatched Lines|iteration {}".format(count))
-                input_str = "\n"
+                #input_str = "\n"
                 #initialize GPT input
-                for i, seg in enumerate(self.segments[start_seg_id-1:end_seg_id]):
-                    input_str += 'Sentence %d: ' %(i+1)+ seg.source_text + '\n'
-                    #Append to prompt string
-                    #Adds sentence index let GPT keep track of sentence breaks
-                input_str += translate    
+                #for i, seg in enumerate(self.segments[start_seg_id-1:end_seg_id]):
+                #    input_str += 'Sentence %d: ' %(i+1)+ seg.source_text + '\n'
+                #    #Append to prompt string
+                #    #Adds sentence index let GPT keep track of sentence breaks
+                #input_str += translate    
                 #append translate to prompt
-                
                 flag = True
                 while flag:
                     flag = False
+                    #print("translate:")
+                    #print(translate)
                     try:
-                        translate = inner_func(input_str)
+                        #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\n')
+                lines = translate.split('\n')
+                #print("result")
+                #print(len(lines))
+                
             if len(lines) < (end_seg_id - start_seg_id + 1):
+                solved = False
                 print("Failed Solving unmatched lines, Manually parse needed")
             
+            if not os.path.exists("./logs"):
+                os.mkdir("./logs")
+            if video_link:
+                log_file = "./logs/log_link.csv"
+                log_exist = os.path.exists(log_file)
+                with open(log_file,"a") as log:
+                    if not log_exist:
+                        log.write("range_of_text,iterations_solving,solved,file_length,video_link" + "\n")
+                    log.write(str(id_range)+','+str(count)+','+str(solved)+','+str(len(self.segments))+','+video_link + "\n")
+            else:
+                log_file = "./logs/log_name.csv"
+                log_exist = os.path.exists(log_file)
+                with open(log_file,"a") as log:
+                    if not log_exist:
+                        log.write("range_of_text,iterations_solving,solved,file_length,video_name" + "\n")
+                    log.write(str(id_range)+','+str(count)+','+str(solved)+','+str(len(self.segments))+','+video_name + "\n")
+            
             print(lines)
             #print(id_range)
             #for i, seg in enumerate(self.segments[start_seg_id-1:end_seg_id]):
@@ -182,23 +214,29 @@ class SRT_script():
             if i < len(lines):
                 if "Note:" in lines[i]: # to avoid note
                     lines.remove(lines[i])
+                    max_num -= 1
                     if i == len(lines) - 1:
                         break
                 try:
-                    seg.translation = lines[i].split("：" or ":")[1]
+                    seg.translation = lines[i].split("：" or ":" or ".")[1]
                 except:
                     seg.translation = lines[i]
-                    #print(lines[i])
-        pass
+
     
-    def split_seg(self, seg, threshold):
+    def split_seg(self, seg, text_threshold, time_threshold):
         # evenly split seg to 2 parts and add new seg into self.segments
-        if seg.source_text[:2] == ', ':
-            seg.source_text = seg.source_text[2:]
+
+        # ignore the initial comma to solve the recursion problem
+        if len(seg.source_text) > 2:
+            if seg.source_text[:2] == ', ':
+                seg.source_text = seg.source_text[2:]
         if seg.translation[0] == '，':
             seg.translation = seg.translation[1:]
+
         source_text = seg.source_text
         translation = seg.translation
+
+        # split the text based on commas
         src_commas = [m.start() for m in re.finditer(',', source_text)]
         trans_commas = [m.start() for m in re.finditer('，', translation)]
         if len(src_commas) != 0:
@@ -215,13 +253,18 @@ class SRT_script():
         else:
             trans_split_idx = len(translation)//2
         
+        # split the time duration based on text length
+        time_split_ratio = trans_split_idx/(len(seg.translation) - 1)
+
         src_seg1 = source_text[:src_split_idx]
         src_seg2 = source_text[src_split_idx:]
         trans_seg1 = translation[:trans_split_idx]
         trans_seg2 = translation[trans_split_idx:]
+
         start_seg1 = seg.start
-        end_seg1 = start_seg2 = seg.start + (seg.end - seg.start)/2
+        end_seg1 = start_seg2 = seg.start + (seg.end - seg.start)*time_split_ratio
         end_seg2 = seg.end
+
         seg1_dict = {}
         seg1_dict['text'] = src_seg1
         seg1_dict['start'] = start_seg1
@@ -237,26 +280,26 @@ class SRT_script():
         seg2.translation = trans_seg2
 
         result_list = []
-        if len(seg1.translation) > threshold:
-            result_list += self.split_seg(seg1, threshold)
+        if len(seg1.translation) > text_threshold and (seg1.end - seg1.start) > time_threshold:
+            result_list += self.split_seg(seg1, text_threshold, time_threshold)
         else:
             result_list.append(seg1)
         
-        if len(seg2.translation) > threshold:
-            result_list += self.split_seg(seg2, threshold)
+        if len(seg2.translation) > text_threshold and (seg2.end - seg2.start) > time_threshold:
+            result_list += self.split_seg(seg2, text_threshold, time_threshold)
         else:
             result_list.append(seg2)
 
         return result_list
 
 
-    def check_len_and_split(self, threshold=30):
+    def check_len_and_split(self, text_threshold=30, time_threshold=1.0):
         # DEPRECATED
-        # if sentence length >= threshold, split this segments to two
+        # if sentence length >= threshold and sentence duration > time_threshold, split this segments to two
         segments = []
         for seg in self.segments:
-            if len(seg.translation) > threshold:
-                seg_list = self.split_seg(seg, threshold)
+            if len(seg.translation) > text_threshold and (seg.end - seg.start) > time_threshold:
+                seg_list = self.split_seg(seg, text_threshold, time_threshold)
                 segments += seg_list
             else:
                 segments.append(seg)
@@ -265,73 +308,25 @@ class SRT_script():
 
         pass
 
-    def check_len_and_split_range(self, range, threshold=30):
-        # if sentence length >= threshold, split this segments to two
+    def check_len_and_split_range(self, range, text_threshold=30, time_threshold=1.0):
+        # if sentence length >= text_threshold, split this segments to two
         start_seg_id = range[0]
         end_seg_id = range[1]
         extra_len = 0
         segments = []
         for i, seg in enumerate(self.segments[start_seg_id-1:end_seg_id]):
-            if len(seg.translation) > threshold:
-                seg_list = self.split_seg(seg, threshold)
+            if len(seg.translation) > text_threshold and (seg.end - seg.start) > time_threshold:
+                seg_list = self.split_seg(seg, text_threshold, time_threshold)
                 segments += seg_list
                 extra_len += len(seg_list) - 1 
             else:
                 segments.append(seg)
         
         self.segments[start_seg_id-1:end_seg_id] = segments
-
         return extra_len
-    
-    def get_source_only(self):
-        # return a string with pure source text
-        result = ""
-        for i, seg in enumerate(self.segments):
-            result+=f'SENTENCE {i+1}: {seg.source_text}\n\n\n'
-        
-        return result
-    
-    def reform_src_str(self):
-        result = ""
-        for i, seg in enumerate(self.segments):
-            result += f'{i+1}\n'
-            result += str(seg)
-        return result
-
-    def reform_trans_str(self):
-        result = ""
-        for i, seg in enumerate(self.segments):
-            result += f'{i+1}\n'
-            result += seg.get_trans_str()
-        return result
-    
-    def form_bilingual_str(self):
-        result = ""
-        for i, seg in enumerate(self.segments):
-            result += f'{i+1}\n'
-            result += seg.get_bilingual_str()
-        return result
-
-    def write_srt_file_src(self, path:str):
-        # write srt file to path
-        with open(path, "w", encoding='utf-8') as f:
-            f.write(self.reform_src_str())
-        pass
-
-    def write_srt_file_translate(self, path:str):
-        with open(path, "w", encoding='utf-8') as f:
-            f.write(self.reform_trans_str())
-        pass
-
-    def write_srt_file_bilingual(self, path:str):
-        with open(path, "w", encoding='utf-8') as f:
-            f.write(self.form_bilingual_str())
-        pass
 
     def correct_with_force_term(self):
         ## force term correction
-        # TODO: shortcut translation i.e. VA, ob
-        # TODO: variety of translation
 
         # load term dictionary
         with open("./finetune_data/dict_enzh.csv",'r', encoding='utf-8') as f:
@@ -420,8 +415,57 @@ class SRT_script():
             real_word = word.lower()
             n = 0
         return real_word, len(word)+n
+
+
+    ## WRITE AND READ FUNCTIONS ##
+
+    def get_source_only(self):
+        # return a string with pure source text
+        result = ""
+        for i, seg in enumerate(self.segments):
+            result+=f'SENTENCE {i+1}: {seg.source_text}\n\n\n'
+        
+        return result
     
+    def reform_src_str(self):
+        result = ""
+        for i, seg in enumerate(self.segments):
+            result += f'{i+1}\n'
+            result += str(seg)
+        return result
+
+    def reform_trans_str(self):
+        result = ""
+        for i, seg in enumerate(self.segments):
+            result += f'{i+1}\n'
+            result += seg.get_trans_str()
+        return result
+    
+    def form_bilingual_str(self):
+        result = ""
+        for i, seg in enumerate(self.segments):
+            result += f'{i+1}\n'
+            result += seg.get_bilingual_str()
+        return result
+
+    def write_srt_file_src(self, path:str):
+        # write srt file to path
+        with open(path, "w", encoding='utf-8') as f:
+            f.write(self.reform_src_str())
+        pass
+
+    def write_srt_file_translate(self, path:str):
+        with open(path, "w", encoding='utf-8') as f:
+            f.write(self.reform_trans_str())
+        pass
+
+    def write_srt_file_bilingual(self, path:str):
+        with open(path, "w", encoding='utf-8') as f:
+            f.write(self.form_bilingual_str())
+        pass
+
     def realtime_write_srt(self,path,range,length, idx):
+        # DEPRECATED
         start_seg_id = range[0]
         end_seg_id = range[1]
         with open(path, "a", encoding='utf-8') as f:
@@ -436,6 +480,7 @@ class SRT_script():
         pass
 
     def realtime_bilingual_write_srt(self,path,range, length,idx):
+        # DEPRECATED
         start_seg_id = range[0]
         end_seg_id = range[1]
         with open(path, "a", encoding='utf-8') as f:
@@ -444,4 +489,6 @@ class SRT_script():
                 if i>=range[1] + length :break
                 f.write(f'{i+idx}\n')
                 f.write(seg.get_bilingual_str())
-        pass
+        pass
+    
+

pipeline.py

@@ -5,6 +5,8 @@ import os
 from tqdm import tqdm
 from SRT import SRT_script
 import stable_whisper
+import whisper
+
 import subprocess
 
 import time
@@ -47,7 +49,7 @@ if args.video_name == 'placeholder' :
     elif args.audio_file is not None:
         VIDEO_NAME = args.audio_file.split('/')[-1].split('.')[0]
     elif args.srt_file is not None:
-        VIDEO_NAME = args.srt_file.split('/')[-1].split('.')[0]
+        VIDEO_NAME = args.srt_file.split('/')[-1].split('.')[0].split("_")[0]
     else:
         VIDEO_NAME = args.video_name
 else:
@@ -95,14 +97,7 @@ elif args.video_file is not None:
         audio_file= open(args.audio_file, "rb")
         audio_path = args.audio_file
     else:
-        # escaped_video_path = args.video_file.replace('(', '\(').replace(')', '\)').replace(' ', '\ ')
-        # print(escaped_video_path)
-        # os.system(f'ffmpeg -i {escaped_video_path} -f mp3 -ab 192000 -vn {DOWNLOAD_PATH}/audio/{VIDEO_NAME}.mp3')
-        # audio_file= open(f'{DOWNLOAD_PATH}/audio/{VIDEO_NAME}.mp3', "rb")
-        # audio_path = f'{DOWNLOAD_PATH}/audio/{VIDEO_NAME}.mp3'
         output_audio_path = f'{DOWNLOAD_PATH}/audio/{VIDEO_NAME}.mp3'
-        # print(video_path)
-        # print(output_audio_path)
         subprocess.run(['ffmpeg', '-i', video_path, '-f', 'mp3', '-ab', '192000', '-vn', output_audio_path])
         audio_file = open(output_audio_path, "rb")
         audio_path = output_audio_path
@@ -133,7 +128,7 @@ else:
 
         # use stable-whisper
         model = stable_whisper.load_model('base')
-        transcript = model.transcribe(audio_path, regroup = False)
+        transcript = model.transcribe(audio_path, regroup = False, initial_prompt="Hello, welcome to my lecture. Are you good my friend?")
         (
             transcript
             .split_by_punctuation(['.', '。', '?'])
@@ -143,14 +138,9 @@ else:
         )
         # transcript.to_srt_vtt(srt_file_en)
         transcript = transcript.to_dict()
+        # print(transcript)
         srt = SRT_script(transcript['segments']) # read segments to SRT class
 
-        #Write SRT file
-        
-        # from whisper.utils import WriteSRT
-        # with open(srt_file_en, 'w', encoding="utf-8") as f:
-        #     writer = WriteSRT(RESULT_PATH)
-        #     writer.write_result(transcript, f)
     else:
         srt = SRT_script.parse_from_srt_file(srt_file_en)
 
@@ -241,21 +231,6 @@ def get_response(model_name, sentence):
         )
 
         return response['choices'][0]['message']['content'].strip()
-    
-    # if model_name == "text-davinci-003":
-    #     prompt = f"Please help me translate this into Chinese:\n\n{s}\n\n"
-    #     # print(prompt)
-    #     response = openai.Completion.create(
-    #         model=model_name,
-    #         prompt=prompt,
-    #         temperature=0.1,
-    #         max_tokens=2000,
-    #         top_p=1.0,
-    #         frequency_penalty=0.0,
-    #         presence_penalty=0.0
-    #     )
-    #     return response['choices'][0]['text'].strip()
-    pass
         
         
 # Translate and save
@@ -277,16 +252,14 @@ for sentence, range in tqdm(zip(script_arr, range_arr)):
             time.sleep(30)
             flag = True
     # add read-time output back and modify the post-processing by using one batch as an unit.
-    srt.set_translation(translate, range, model_name)
-    # add_length = srt.check_len_and_split_range(range, threshold)
-    # srt.realtime_write_srt(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt",range, add_length ,segidx)
-    # # save current length as previous length
-    # previous_length = add_length
+    srt.set_translation(translate, range, model_name, VIDEO_NAME, args.link)
+
     # srt.realtime_bilingual_write_srt(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_bi.srt",range, add_length,segidx)
 
 srt.check_len_and_split()
+srt.remove_trans_punctuation()
 srt.write_srt_file_translate(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt")
-# srt.write_srt_file_bilingual(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_bi.srt")
+srt.write_srt_file_bilingual(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_bi.srt")
 
 if not args.only_srt:
     assSub_zh = srt2ass(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt", "default", "No", "Modest")