Spaces:
Sleeping
Sleeping
File size: 2,824 Bytes
63ecb0d a30dacb 63ecb0d 1b7b90e 63ecb0d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
import openai
from pytube import YouTube
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)
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)
args = parser.parse_args()
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")
DOWNLOAD_PATH = args.download
RESULT_PATH = args.result
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
video_link = args.link
try:
video = YouTube(video_link)
audio = video.streams.filter(only_audio=True, file_extension='mp4').first()
audio.download(DOWNLOAD_PATH)
print('Download Completed!')
except:
print("Connection Error")
audio_file = open('{}/{}'.format(DOWNLOAD_PATH, audio.default_filename), "rb")
VIDEO_NAME = audio.default_filename.split('.')[0]
else:
# Read from local
audio_file= open(args.local_path, "rb")
# perform speech-to-text and save it in <video name>_en.txt under RESULT PATH.
if not os.path.exists("{}/{}_en.txt".format(RESULT_PATH, VIDEO_NAME)):
transcript = openai.Audio.transcribe("whisper-1", audio_file)
with open("{}/{}_en.txt".format(RESULT_PATH, VIDEO_NAME), 'w') as f:
f.write(transcript['text'])
# split the video script(open ai prompt limit: about 5000)
with open("{}/{}_en.txt".format(RESULT_PATH, VIDEO_NAME), 'r') as f:
script_en = f.read()
N = len(script_en)
script_split = script_en.split('.')
script_arr = []
script = ""
for sentence in script_split:
if len(script) <= n_threshold:
n = len(sentence)
script+=sentence
else:
script_arr.append(script)
script = ""
script_arr.append(script)
# translate and save
for s in script_arr:
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
)
with open("{}/{}_zh.txt".format(RESULT_PATH, VIDEO_NAME), 'a+') as f:
f.write(response['choices'][0]['text']) |