small_fix increase token size

Former-commit-id: 98d31138827b8cd27577205c3c21155a6f686772

README.md

@@ -10,6 +10,12 @@ pip install -r requirement.txt
 ```
 usage: pipeline.py [-h] [--link LINK] [--local_path LOCAL_PATH] [--download DOWNLOAD] [--result RESULT] [--video_name VIDEO_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 offline: python3 pipeline.py --local_path test_translation.m4a --result ./results --video_name test_translation
+
 options:
   -h, --help            show this help message and exit
   --link LINK           youtube video link here

pipeline.py

@@ -16,8 +16,8 @@ if args.link is None and args.local_path is None:
     print("need video source")
     exit()
 
-# openai.api_key = "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"  
-openai.api_key = os.getenv("OPENAI_API_KEY")
+openai.api_key = "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"  
+# openai.api_key = os.getenv("OPENAI_API_KEY")
 
 DOWNLOAD_PATH = args.download
 RESULT_PATH = args.result
@@ -25,6 +25,7 @@ VIDEO_NAME = args.video_name
 n_threshold = 5000
 model_name = "text-davinci-003" # replace this to our own fintune model
 
+
 # get source audio
 if args.link is not None:
     # Download audio from YouTube
@@ -34,8 +35,9 @@ if args.link is not None:
         audio = video.streams.filter(only_audio=True, file_extension='mp4').first()
         audio.download(DOWNLOAD_PATH)
         print('Download Completed!')
-    except:
-        print("Connection Error") 
+    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]
 else:
@@ -54,28 +56,34 @@ with open("{}/{}_en.txt".format(RESULT_PATH, VIDEO_NAME), 'r') as f:
     N = len(script_en)
     script_split = script_en.split('.')
 
+# Split the video script by sentences and create chunks within the token limit
+n_threshold = 4096  # Token limit for the GPT-3 model
+script_split = script_en.split('.')
+
 script_arr = []
 script = ""
 for sentence in script_split:
-    if len(script) <= n_threshold:
-        n = len(sentence)
-        script+=sentence
+    if len(script) + len(sentence) + 1 <= n_threshold:
+        script += sentence + '.'
     else:
-        script_arr.append(script)
-        script = ""
-script_arr.append(script)
+        script_arr.append(script.strip())
+        script = sentence + '.'
+if script.strip():
+    script_arr.append(script.strip())
 
-# translate and save
+# Translate and save
 for s in script_arr:
+    prompt = f"Please help me translate this into Chinese:\n\n{s}\n\n"
     response = openai.Completion.create(
-    model=model_name,
-    prompt="Please healp me translate this into Chinese:\n\n{}\n\n".format(s),
-    temperature=0.1,
-    max_tokens=2000,
-    top_p=1.0,
-    frequency_penalty=0.0,
-    presence_penalty=0.0
+        model=model_name,
+        prompt=prompt,
+        temperature=0.1,
+        max_tokens=2000,
+        top_p=1.0,
+        frequency_penalty=0.0,
+        presence_penalty=0.0
     )
 
-    with open("{}/{}_zh.txt".format(RESULT_PATH, VIDEO_NAME), 'a+') as f:
-        f.write(response['choices'][0]['text'])
+    with open(f"{RESULT_PATH}/{VIDEO_NAME}_zh.txt", 'a+') as f:
+        f.write(response['choices'][0]['text'].strip())
+        f.write('\n')

requirement.txt

@@ -13,7 +13,6 @@ panda==0.3.1
 pandas==1.5.3
 python-dateutil==2.8.2
 pytube==12.1.2
-pytube3==9.6.4
 pytz==2022.7.1
 requests==2.28.2
 six==1.16.0