change default model; short sentence merge; edit translation prompt

SRT.py

@@ -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] in ['.', '!', '?']:
+            if seg.source_text[-1] in ['.', '!', '?'] and len(seg.source_text) > 10:
                 sentence.append(i)
                 merge_list.append(sentence)
                 sentence = []
@@ -127,6 +127,11 @@ class SRT_script():
         start_seg_id = id_range[0]
         end_seg_id = id_range[1]
         
+        src_text = ""
+        for i, seg in enumerate(self.segments[start_seg_id-1:end_seg_id]):
+            src_text+=seg.source_text
+            src_text+='\n\n'
+
         def inner_func(target,input_str):
             response = openai.ChatCompletion.create(
                     #model=model,
@@ -136,13 +141,20 @@ class SRT_script():
                         #{"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)}
+                        # {"role": "system", "content": "请将以下中文与其英文句子一一对应并输出："},
+                        # {"role": "system", "content": "英文：{}".format(src_text)},
+                        # {"role": "user", "content": "中文：{}\n\n".format(input_str)},
                     ],
-                    #temperature=0.7
                     temperature = 0.15
                 )
+            # print(src_text)
+            # print(input_str)
+            # print(response['choices'][0]['message']['content'].strip())
+            # exit()
             return response['choices'][0]['message']['content'].strip()
         
         
@@ -206,7 +218,6 @@ class SRT_script():
             #print(translate)
             
         
-        
         for i, seg in enumerate(self.segments[start_seg_id-1:end_seg_id]):
             # naive way to due with merge translation problem
             # TODO: need a smarter solution

pipeline.py

@@ -19,7 +19,7 @@ parser.add_argument("--srt_file", help="srt file input path here", default=None,
 parser.add_argument("--download", help="download path", default='./downloads', type=str, required=False)
 parser.add_argument("--output_dir", help="translate result path", default='./results', type=str, required=False)
 parser.add_argument("--video_name", help="video name, if use video link as input, the name will auto-filled by youtube video name", default='placeholder', type=str, required=False)
-parser.add_argument("--model_name", help="model name only support gpt-4 and gpt-3.5-turbo", type=str, required=False, default="gpt-3.5-turbo")
+parser.add_argument("--model_name", help="model name only support gpt-4 and gpt-3.5-turbo", type=str, required=False, default="gpt-4") # default change to gpt-4
 parser.add_argument("-only_srt", help="set script output to only .srt file", action='store_true')
 parser.add_argument("-v", help="auto encode script with video", action='store_true')
 args = parser.parse_args()
@@ -136,9 +136,8 @@ else:
             .merge_by_punctuation([' '])
             .split_by_punctuation(['.', '。', '?'])
         )
-        # transcript.to_srt_vtt(srt_file_en)
+
         transcript = transcript.to_dict()
-        # print(transcript)
         srt = SRT_script(transcript['segments']) # read segments to SRT class
 
     else:
@@ -158,7 +157,7 @@ if not args.only_srt:
 
 
 # Split the video script by sentences and create chunks within the token limit
-def script_split(script_in, chunk_size = 1000):
+def script_split(script_in, chunk_size = 500):
     script_split = script_in.split('\n\n')
     script_arr = []
     range_arr = []
@@ -224,6 +223,7 @@ def get_response(model_name, sentence):
             messages = [
                 {"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": "Your translation needs to be consistent with the number of sentences in the original."},
                 {"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(sentence)}
             ],