src/translators/translation.py

from os import getenv
import logging
from time import sleep
from tqdm import tqdm
from src.srt_util.srt import split_script
from .LLM_task import LLM_task

def get_translation(srt, model, video_name, prompt, chunk_size = 1000):
    script_arr, range_arr = split_script(srt.get_source_only(),chunk_size)
    translate(srt, script_arr, range_arr, model, video_name, task=prompt)
    pass

def check_translation(sentence, translation):
    """
    check merge sentence issue from openai translation
    """
    sentence_count = sentence.count('\n\n') + 1
    translation_count = translation.count('\n\n') + 1

    if sentence_count != translation_count:
        # print("sentence length: ", len(sentence), sentence_count)
        # print("translation length: ",  len(translation), translation_count)
        return False
    else:
        return True

# TODO{david}: prompts selector
def prompt_selector(src_lang, tgt_lang, domain):
    language_map = {
        "EN": "English",
        "ZH": "Chinese",
    }
    src_lang = language_map[src_lang]
    tgt_lang = language_map[tgt_lang]
    prompt = f"""
        you are a translation assistant, your job is to translate a video in domain of {domain} from {src_lang} to {tgt_lang}, 
        you will be provided with a segement in {[src_lang]} parsed by line, where your translation text should keep the original 
        meaning and the number of lines.
        """
    return prompt

def translate(srt, script_arr, range_arr, model_name, video_name=None, attempts_count=5, task=None, temp = 0.15):
    """
    Translates the given script array into another language using the chatgpt and writes to the SRT file.

    This function takes a script array, a range array, a model name, a video name, and a video link as input. It iterates
    through sentences and range in the script and range arrays. If the translation check fails for five times, the function
    will attempt to resolve merge sentence issues and split the sentence into smaller tokens for a better translation.

    :param srt: An instance of the Subtitle class representing the SRT file.
    :param script_arr: A list of strings representing the original script sentences to be translated.
    :param range_arr: A list of tuples representing the start and end positions of sentences in the script.
    :param model_name: The name of the translation model to be used.
    :param video_name: The name of the video.
    :param attempts_count: Number of attemps of failures for unmatched sentences.
    :param task: Prompt.
    :param temp: Model temperature.
    """
    #logging.info("Start translating...")
    if input is None: 
        raise Exception("Warning! No Input have passed to LLM!")
    if task is None:
        task = "你是一个翻译助理，你的任务是翻译星际争霸视频，你会被提供一个按行分割的英文段落，你需要在保证句意和行数的情况下输出翻译后的文本。"
    logging.info(f"translation prompt: {task}")
    previous_length = 0
    for sentence, range_ in tqdm(zip(script_arr, range_arr)):
        # update the range based on previous length
        range_ = (range_[0] + previous_length, range_[1] + previous_length)
        # using chatgpt model
        print(f"now translating sentences {range_}")
        logging.info(f"now translating sentences {range_}")
        flag = True
        while flag:
            flag = False
            try:
                translate = LLM_task(model_name, sentence, task, temp)
                # detect merge sentence issue and try to solve for five times:
                while not check_translation(sentence, translate) and attempts_count > 0:
                    translate = LLM_task(model_name, sentence, task, temp)
                    attempts_count -= 1

                # if failure still happen, split into smaller tokens
                if attempts_count == 0:
                    single_sentences = sentence.split("\n\n")
                    logging.info("merge sentence issue found for range", range_)
                    translate = ""
                    for i, single_sentence in enumerate(single_sentences):
                        if i == len(single_sentences) - 1:
                            translate += LLM_task(model_name,sentence,task,temp)
                        else:
                            translate += LLM_task(model_name,sentence,task,temp) + "\n\n"
                    logging.info("solved by individually translation!")

            except Exception as e:
                logging.debug("An error has occurred during translation:", e)
                print("An error has occurred during translation:", e)
                print("Retrying... the script will continue after 30 seconds.")
                sleep(30)
                flag = True

        srt.set_translation(translate, range_, model_name, video_name)