add batch output and continue translate function

Former-commit-id: 9163523f8449d647b6b1fa0da45ef72afa4ff0c3

SRT.py

@@ -3,6 +3,7 @@ from csv import reader
 from datetime import datetime
 import re
 import openai
+from collections import deque
 
 class SRT_segment(object):
     def __init__(self, *args) -> None:
@@ -180,7 +181,7 @@ class SRT_script():
                     #print(lines[i])
         pass
     
-    def split_seg(self, seg, threshold):
+    def split_seg(self, seg, threshold=500):
         # TODO: evenly split seg to 2 parts and add new seg into self.segments
         source_text = seg.source_text
         translation = seg.translation
@@ -193,7 +194,7 @@ class SRT_script():
             src_split_idx = src_space[len(src_space)//2] if len(src_space) % 2 == 1 else src_space[len(src_space)//2 - 1]
 
         if len(trans_commas) != 0:
-            trans_split_idx = trans_commas[len(src_commas)//2] if len(trans_commas) % 2 == 1 else trans_commas[len(trans_commas)//2 - 1]
+            trans_split_idx = trans_commas[len(trans_commas)//2] if len(trans_commas) % 2 == 1 else trans_commas[len(trans_commas)//2 - 1]
         else:
             trans_split_idx = len(translation)//2
         
@@ -246,6 +247,22 @@ class SRT_script():
 
         pass
 
+    def check_len_and_split_range(self, range, threshold=30):
+        # TODO: if sentence length >= threshold, split this segments to two
+        start_seg_id = range[0]
+        end_seg_id = range[1]
+        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)
+                segments += seg_list
+            else:
+                segments.append(seg)
+        
+        self.segments[start_seg_id-1:end_seg_id] = segments
+
+        return len(segments)
+    
     def get_source_only(self):
         # return a string with pure source text
         result = ""
@@ -383,4 +400,22 @@ class SRT_script():
         else:
             real_word = word.lower()
             n = 0
-        return real_word, len(word)+n
+        return real_word, len(word)+n
+    
+    def realtime_write_srt(self,path,range,length,idx):
+        start_seg_id = range[0]
+        end_seg_id = range[1]
+        with open(path, "a", encoding='utf-8') as f:
+            for i, seg in enumerate(self.segments[start_seg_id-1:end_seg_id+length]):
+                f.write(f'{i+idx}\n')
+                f.write(seg.get_trans_str())
+        pass
+
+    def realtime_bilingual_write_srt(self,path,range,length,idx):
+        start_seg_id = range[0]
+        end_seg_id = range[1]
+        with open(path, "a", encoding='utf-8') as f:
+            for i, seg in enumerate(self.segments[start_seg_id-1:end_seg_id+length]):
+                f.write(f'{i+idx}\n')
+                f.write(seg.get_bilingual_str())
+        pass

pipeline.py

@@ -52,6 +52,8 @@ else:
 
 model_name = args.model_name
 
+threshold = 30
+
 # get source audio
 if args.link is not None and args.video_file is None:
     # Download audio from YouTube
@@ -140,7 +142,7 @@ else:
     else:
         srt = SRT_script.parse_from_srt_file(srt_file_en)
 
-# srt preprocess
+# srt class preprocess
 srt.form_whole_sentence()
 srt.spell_check_term()
 srt.correct_with_force_term()
@@ -179,8 +181,38 @@ def script_split(script_in, chunk_size = 1000):
     return script_arr, range_arr
 
 script_arr, range_arr = script_split(script_input)
+# print(script_arr, range_arr)
+
+# check whether previous translation is done
+zh_file = "{}/{}/{}_zh.srt".format(RESULT_PATH, VIDEO_NAME, VIDEO_NAME)
+segidx = 1
+if os.path.exists(zh_file):
+    temp_file = "{}/{}/temp.srt".format(RESULT_PATH, VIDEO_NAME)
+    if os.path.exists(temp_file):
+        os.remove(temp_file)
+    with open(zh_file, "r") as f0:
+        for count, _ in enumerate(f0):
+            pass
+        count += 1
+        segidx = int(count/4)+1
+        
+    with open("{}/{}/{}_en.srt".format(RESULT_PATH, VIDEO_NAME, VIDEO_NAME), "r") as f1, open(temp_file, "a") as f2:
+        x = f1.readlines()
+        #print(len(x))
+        if count >= len(x):
+            print('Work already done! Please delete {}_zh.srt files in result directory first in order to rework'.format(VIDEO_NAME))
+            exit()
+        for i, line in enumerate(x):
+            if i >= count:
+                #print(i)
+                f2.write(line)
+                
+    srt = SRT_script.parse_from_srt_file(temp_file)
+    print('temp_contents')
+    print(srt.get_source_only())  
+
 
-def get_response(model_name):
+def get_response(model_name, sentence):
     if model_name == "gpt-3.5-turbo" or model_name == "gpt-4":
         # print(s + "\n")
         response = openai.ChatCompletion.create(
@@ -189,48 +221,53 @@ def get_response(model_name):
                 {"role": "system", "content": "You are a helpful assistant that translates English to Chinese and have decent background in starcraft2."},
                 {"role": "system", "content": "Your translation has to keep the orginal format and be as accurate as possible."},
                 {"role": "system", "content": "There is no need for you to add any comments or notes."},
-                {"role": "user", "content": 'Translate the following English text to Chinese: "{}"'.format(s)}
+                {"role": "user", "content": 'Translate the following English text to Chinese: "{}"'.format(sentence)}
             ],
             temperature=0.15
         )
 
         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()
+    
+    # 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
-for s, range in tqdm(zip(script_arr, range_arr)):
+for sentence, range in tqdm(zip(script_arr, range_arr)):
     # using chatgpt model
     print(f"now translating sentences {range}")
     flag = True
     while flag:
         flag = False
         try:
-            translate = get_response(model_name)
+            translate = get_response(model_name, sentence)
         except Exception as e:
             print("An error has occurred during translation:",e)
-            print("Retrying...")
+            print("Retrying... the script will continue after 30 seconds.")
             time.sleep(30)
             flag = True
+    # add read-time output back and modify the post-processing by using one batch as an unit.
+    print(translate)
     srt.set_translation(translate, range, model_name)
+    add_length = srt.check_len_and_split_range(range)
+    srt.realtime_write_srt(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt",range, add_length,segidx)
+    srt.realtime_bilingual_write_srt(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_bi.srt",range, add_length,segidx)
 
-srt.check_len_and_split()
-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.check_len_and_split()
+# 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")
 
 if not args.only_srt:
     assSub_zh = srt2ass(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt", "default", "No", "Modest")