translation module implementation

Former-commit-id: e6b1d4b73c063aca65f510a23fb4ead6eccea1c1

src/srt_util/srt.py

@@ -532,3 +532,27 @@ class SrtScript(object):
                 f.write(f'{i + idx}\n')
                 f.write(seg.get_bilingual_str())
         pass
+
+def split_script(script_in, chunk_size=1000):
+    script_split = script_in.split('\n\n')
+    script_arr = []
+    range_arr = []
+    start = 1
+    end = 0
+    script = ""
+    for sentence in script_split:
+        if len(script) + len(sentence) + 1 <= chunk_size:
+            script += sentence + '\n\n'
+            end += 1
+        else:
+            range_arr.append((start, end))
+            start = end + 1
+            end += 1
+            script_arr.append(script.strip())
+            script = sentence + '\n\n'
+    if script.strip():
+        script_arr.append(script.strip())
+        range_arr.append((start, len(script_split) - 1))
+
+    assert len(script_arr) == len(range_arr)
+    return script_arr, range_arr

src/task.py

@@ -10,6 +10,9 @@ import logging
 import subprocess
 from src.srt_util.srt import SrtScript
 from src.srt_util.srt2ass import srt2ass
+from translation.translation import get_translation, translate
+
+
 
 """
 Youtube link
@@ -137,7 +140,8 @@ class Task:
 
     # Module 3: perform srt translation
     def translation(self):
-        time.sleep(5)
+        logging.info("---------------------Start Translation--------------------")
+        get_translation(self.srt,self.model, self.video_name, self.video_link)
         pass
     
     # Module 4: perform srt post process steps

src/translation/LLM_task.py

@@ -0,0 +1,26 @@
+
+import openai
+
+
+def LLM_task(model_name, input, task, temp = 0.15):
+    """
+    Translates input sentence with desired LLM.
+
+    :param model_name: The name of the translation model to be used.
+    :param input: Sentence for translation.
+    :param task: Prompt.
+    :param temp: Model temperature.
+    """
+    if model_name == "gpt-3.5-turbo" or model_name == "gpt-4":
+        response = openai.ChatCompletion.create(
+            model=model_name,
+            messages=[
+                {"role": "system","content": task},
+                {"role": "user", "content": input}
+            ],
+            temperature=temp
+        )
+        return response['choices'][0]['message']['content'].strip()
+    # Other LLM not implemented
+    else:
+        raise NotImplementedError

src/translation/translation.py

@@ -0,0 +1,88 @@
+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,video_link):
+    script_arr, range_arr = split_script(srt)
+    translate(srt, script_arr, range_arr, model, video_name, video_link)
+    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
+
+
+def translate(srt, script_arr, range_arr, model_name, video_name, video_link, 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 video_link: The link to 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 = "你是一个翻译助理，你的任务是翻译星际争霸视频，你会被提供一个按行分割的英文段落，你需要在保证句意和行数的情况下输出翻译后的文本。"
+
+    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_}, time: {datetime.now()}")
+        flag = True
+        while flag:
+            flag = False
+            try:
+                translate = LLM_task(model_name, sentence)
+                # 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, video_link)