使用AI翻译电影字幕

def _guess_eng_subtitle_index(video_path):
    probe = ffmpeg.probe(video_path)
    streams = probe['streams']
    for index, stream in enumerate(streams):
        if stream.get('codec_type') == 'subtitle' and stream.get('tags', {}).get('language') == 'eng':
            return index
    for index, stream in enumerate(streams):
        if stream['codec_type'] == 'subtitle' and stream.get('tags', {}).get('title', "").lower().find("english")!=-1 :
            return index
    return -1

def _extract_subtitle_by_index(video_path, output_path, index):
    return ffmpeg.input(video_path).output(output_path, map='0:'+str(index)).run()

def extract_subtitle(video_path, en_subtitle_path):
    # get the streams from video with ffprobe
    index = _guess_eng_subtitle_index(video_path)
    if index == -1:
        return -1
    
    return _extract_subtitle_by_index(video_path, en_subtitle_path, index)

class UpSubs:
    def __init__(self, subs_path):
        self.subs = pysubs2.load(subs_path)

    def get_subtitle_text(self):
        text = ""
        for sub in self.subs:
            text += sub.text + "\n\n"
        return text

    def get_subtitle_text_with_index(self):
        text = ""
        for i, sub in enumerate(self.subs):
            text += "chunk-"+str(i) + ":\n" + sub.text.replace("\\N", " ") + "\n\n"
        return text
    
    def save(self, output_path):
        self.subs.save(output_path)

    def clean(self):
        indexes = []
        for i, sub in enumerate(self.subs):
            # remove xml tag and line change in sub text
            sub.text = re.sub(r"<[^>]+>", "", sub.text)
            sub.text = sub.text.replace("\\N", " ")

    def fill(self, text):
        text = text.strip()
        pattern = r"\n\s*\n"
        paragraphs = re.split(pattern, text)
        for para in paragraphs:
            try:
                firtline = para.split("\n")[0]
                countstr = firtline[6:len(firtline)-1]
                # print(countstr)
                index = int(countstr)
                p = "\n".join(para.split("\n")[1:])
                self.subs[index].text = p
            except Exception as e:
                print(f"Error merge paragraph : \n {para} \n with exception: \n {e}")
                raise(e)
    
    def merge_dual(self, subspath):
        second_subs = pysubs2.load(subspath)
        merged_subs = SSAFile()
        if len(self.subs.events) == len(second_subs.events):            
            for i, first_event in enumerate(self.subs.events):
                second_event = second_subs[i]
                if first_event.text == second_event.text:
                    merged_event = SSAEvent(first_event.start, first_event.end, first_event.text)
                else:
                    merged_event = SSAEvent(first_event.start, first_event.end, first_event.text + '\n' + second_event.text)
                merged_subs.append(merged_event)
            return merged_subs
        
        return None

def complete(prompt, max_tokens=32760):
    prompt = prompt.strip()
    if not prompt:
        return ""
    
    safety_settings = [
        {
            "category": "HARM_CATEGORY_HARASSMENT",
            "threshold": "BLOCK_NONE"
        },
        {
            "category": "HARM_CATEGORY_HATE_SPEECH",
            "threshold": "BLOCK_NONE"
        },
        {
            "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
            "threshold": "BLOCK_NONE"
        },
        {
            "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
            "threshold": "BLOCK_NONE"
        },
    ]

    retries = 3
    for _ in range(retries):
        try:
            return Gemini(max_tokens=max_tokens, safety_settings=safety_settings, temperature = 0.01).complete(prompt).text
        except Exception as e:
            print(f"Error completing prompt: {prompt} \n with error: \n ")
            traceback.print_exc()
    return ""

def tran_subtitles(fixed_subtitle, zh_subtitle=None, cncf = False, chunk_size=3000):
    subtitle_base = os.path.splitext(fixed_subtitle)[0]
    video_base = os.path.splitext(subtitle_base)[0]
    if zh_subtitle is None:
        zh_subtitle = video_base + ".zh-fixed.vtt"
    if os.path.exists(zh_subtitle):
        print(f"zh subtitle {zh_subtitle} already translated, skip to translate.")
        return 1

    prompt_tpl = MOVIE_TRAN_PROMPT_TPL
    opts = { }

    srtp = UpSubs(fixed_subtitle)
    text = srtp.get_subtitle_text_with_index()

    process_text(srtp, text, prompt_tpl, opts, chunk_size = chunk_size)
    srtp.save(zh_subtitle)

    return zh_subtitle

MOVIE_TRAN_PROMPT_TPL = """你是个专业电影字幕翻译，你需要将一份英文字幕翻译成中文。
[需要翻译的英文字幕]:

{content}

# [中文字幕]:"""

def process_text(subs, text, prompt_tpl, opts, chunk_size=2500):
    # ret = ""
    chunks = _split_subtitles(text, chunk_size)
    for(i, chunk) in enumerate(chunks):
        print("process chunk ({}/{})".format(i+1,len(chunks)))
        # if i==4:
        #     break
        # format string with all the field in a dict 
        opts["content"] = chunk
        prompt = prompt_tpl.format(**opts)

        print(prompt)
        out = complete(prompt, max_tokens=32760)
        subs.fill(out)
        print(out)