in progress

Former-commit-id: a417baf3dcafc23778cbd99c1315f892d9569822

SRT.py

@@ -109,13 +109,13 @@ class SRT_script():
         
     
 
-    def set_translation(self, translate:str, id_range:tuple):
+    def set_translation(self, translate:str, id_range:tuple, model):
         start_seg_id = id_range[0]
         end_seg_id = id_range[1]
         
         def inner_func(input_str):
             response = openai.ChatCompletion.create(
-                    model="gpt-3.5-turbo",
+                    model=model,
                     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."},

pipeline.py

@@ -6,6 +6,8 @@ from tqdm import tqdm
 from SRT import SRT_script
 import stable_whisper
 
+import time
+
 parser = argparse.ArgumentParser()
 parser.add_argument("--link", help="youtube video link here", default=None, type=str, required=False)
 parser.add_argument("--video_file", help="local video path here", default=None, type=str, required=False)
@@ -14,7 +16,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 text-davinci-003 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-3.5-turbo")
 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()
@@ -36,7 +38,18 @@ RESULT_PATH = args.output_dir
 if not os.path.exists(RESULT_PATH):
     os.mkdir(RESULT_PATH)
 
-VIDEO_NAME = args.video_name
+# set video name as the input file name if not specified
+if args.video_name == 'placeholder' :
+    # set video name to upload file name
+    if args.video_file is not None:
+        VIDEO_NAME = args.video_file.split('/')[-1].split('.')[0]
+    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]
+else:
+    VIDEO_NAME = args.video_name
+
 model_name = args.model_name
 
 # get source audio
@@ -83,6 +96,10 @@ elif args.video_file is not None:
 if not os.path.exists(f'{RESULT_PATH}/{VIDEO_NAME}'):
     os.mkdir(f'{RESULT_PATH}/{VIDEO_NAME}')
 
+if args.audio_file is not None:
+    audio_file= open(args.audio_file, "rb")
+    audio_path = args.audio_file
+
 # Instead of using the script_en variable directly, we'll use script_input
 srt_file_en = args.srt_file 
 
@@ -164,7 +181,7 @@ def script_split(script_in, chunk_size = 1000):
 script_arr, range_arr = script_split(script_input)
 
 def get_response(model_name):
-    if model_name == "gpt-3.5-turbo":
+    if model_name == "gpt-3.5-turbo" or model_name == "gpt-4":
         # print(s + "\n")
         response = openai.ChatCompletion.create(
             model=model_name,
@@ -207,8 +224,9 @@ for s, range in tqdm(zip(script_arr, range_arr)):
         except Exception as e:
             print("An error has occurred during translation:",e)
             print("Retrying...")
+            time.sleep(30)
             flag = True
-    srt.set_translation(translate, range)
+    srt.set_translation(translate, range, model_name)
 
 srt.check_len_and_split()
 srt.write_srt_file_translate(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt")