small fix

Former-commit-id: 400a737586a90592188ba1fb899464599925400e

README.md

@@ -8,11 +8,11 @@ pip install -r requirement.txt
 
 ## Usage
 ```
-usage: pipeline.py [-h] [--link LINK] [--local_path LOCAL_PATH] [--download DOWNLOAD] [--result RESULT] [--video_name VIDEO_NAME]
+usage: pipeline.py [-h] [--link LINK] [--local_path LOCAL_PATH] [--download DOWNLOAD] [--result RESULT] [--video_name VIDEO_NAME] [--model_name]
 
 quick start:
 
-example online: python3 pipeline.py --link https://www.youtube.com/watch?v=XbgFIkhMM3s --download ./downloads --result ./results --video_name uncle_roger_test
+example online: python3 pipeline.py --link https://www.youtube.com/watch?v=61c4dn6851g --download ./downloads --result ./results --video_name SO_I_CHOSE_RANDOM
 
 example offline: python3 pipeline.py --local_path test_translation.m4a --result ./results --video_name test_translation
 
@@ -27,4 +27,10 @@ options:
   --result RESULT       translate result path
   --video_name VIDEO_NAME
                         video name
+  --model_name MODEL_NAME 
+                        model name
+
+if you cannot download youtube video, please follow the link below.
+https://github.com/pytube/pytube/issues/1498
+
 ```

pipeline.py

@@ -4,8 +4,6 @@ import argparse
 import os
 import io
 
-
-
 parser = argparse.ArgumentParser()
 parser.add_argument("--link", help="youtube video link here", default=None, type=str, required=False)
 parser.add_argument("--local_path", help="local video path here", default=None, type=str, required=False)
@@ -13,40 +11,46 @@ parser.add_argument("--text_file", help="text file path here", default=None, typ
 parser.add_argument("--download", help="download path", default='./downloads', type=str, required=False)
 parser.add_argument("--result", help="translate result path", default='./results', type=str, required=False)
 parser.add_argument("--video_name", help="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", default='placeholder', type=str, required=False, default="gpt-3.5-turbo")
+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")
 args = parser.parse_args()
 
 if args.link is None and args.local_path is None and args.text_file is None:
     print("need video source or text file")
     exit()
 
+# set openai api key
 openai.api_key = os.getenv("OPENAI_API_KEY")
 DOWNLOAD_PATH = args.download
 RESULT_PATH = args.result
 VIDEO_NAME = args.video_name
-n_threshold = 1000  # Token limit for the GPT-3.5 model
-# model_name = "text-davinci-003" # replace this to our own fintune model
 model_name = args.model_name
-# model_name = "davinci"
 
 # get source audio
 if args.link is not None and args.local_path is None:
     # Download audio from YouTube
     video_link = args.link
+    video = None
+    audio = None
     try:
         video = YouTube(video_link)
         audio = video.streams.filter(only_audio=True, file_extension='mp4').first()
-        audio.download(DOWNLOAD_PATH)
-        print('Download Completed!')
+        if audio:
+            audio.download(DOWNLOAD_PATH)
+            print('Download Completed!')
+        else:
+            print("Error: Audio stream not found")
     except Exception as e:
         print("Connection Error")
-        print(e) 
-    audio_file = open('{}/{}'.format(DOWNLOAD_PATH, audio.default_filename), "rb")
-    VIDEO_NAME = audio.default_filename.split('.')[0]
+        print(e)
+    if audio:
+        audio_file = open('{}/{}'.format(DOWNLOAD_PATH, audio.default_filename), "rb")
+        VIDEO_NAME = audio.default_filename.split('.')[0]
+    else:
+        print("Error: Unable to download audio from the YouTube video")
+        exit()
 elif args.local_path is not None:
     # Read from local
-    audio_file= open(args.local_path, "rb")
-
+    audio_file = open(args.local_path, "rb")
 
 
 # Instead of using the script_en variable directly, we'll use script_input
@@ -68,7 +72,7 @@ else:
         script_input = script_en
 
 # Split the video script by sentences and create chunks within the token limit
-n_threshold = 4096  # Token limit for the GPT-3 model
+n_threshold = 1500  # Token limit for the GPT-3 model
 script_split = script_input.split('.')
 
 script_arr = []
@@ -93,11 +97,12 @@ for s in script_arr:
                 {"role": "system", "content": "You are a helpful assistant that translates English to Chinese and have decent background in starcraft2."},
                 {"role": "user", "content": 'Translate the following English text to Chinese: "{}"'.format(s)}
             ],
-            temperature=0.1
+            temperature=0.15
         )
         with open(f"{RESULT_PATH}/{VIDEO_NAME}_zh.txt", 'a+') as f:
             f.write(response['choices'][0]['message']['content'].strip())
-            f.write('\n')
+
+    # using davinci model
     if model_name == "text-davinci-003":
         prompt = f"Please help me translate this into Chinese:\n\n{s}\n\n"
         print(prompt)
@@ -112,5 +117,4 @@ for s in script_arr:
         )
 
         with open(f"{RESULT_PATH}/{VIDEO_NAME}_zh.txt", 'a+') as f:
-            f.write(response['choices'][0]['text'].strip())
-            f.write('\n')
+            f.write(response['choices'][0]['text'].strip())