Merge branch 'oop-refactor' of https://github.com/project-kxkg/project-t into oop-refactor

src/task.py

@@ -11,7 +11,7 @@ import subprocess
 from src.srt_util.srt import SrtScript
 from src.srt_util.srt2ass import srt2ass
 from time import time, strftime, gmtime, sleep
-from src.translators.translation import get_translation, translate
+from src.translators.translation import get_translation, prompt_selector
 
 import torch
 import stable_whisper
@@ -183,9 +183,10 @@ class Task:
             self.progress = TaskStatus.TRANSLATING.value[0], new_progress
 
     # Module 3: perform srt translation
-    def translation(self):
+    def translation(self,task_cfg):
         logging.info("---------------------Start Translation--------------------")
-        get_translation(self.SRT_Script, self.model, self.task_id)
+        prompt = prompt_selector(self.source_lang,self.target_lang,task_cfg['field'])
+        get_translation(self.SRT_Script, self.model, self.task_id, prompt, task_cfg['chunk_size'])
     
     # Module 4: perform srt post process steps
     def postprocess(self):

src/translators/translation.py

@@ -5,9 +5,9 @@ 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):
-    script_arr, range_arr = split_script(srt.get_source_only())
-    translate(srt, script_arr, range_arr, model, video_name)
+def get_translation(srt, model, video_name, task, 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)
     pass
 
 def check_translation(sentence, translation):
@@ -26,8 +26,18 @@ def check_translation(sentence, translation):
 
 # TODO{david}: prompts selector
 def prompt_selector(src_lang, tgt_lang, domain):
-
-    return ""
+    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, attempts_count=5, task=None, temp = 0.15):
     """
@@ -51,7 +61,7 @@ def translate(srt, script_arr, range_arr, model_name, video_name, attempts_count
         raise Exception("Warning! No Input have passed to LLM!")
     if task is None:
         task = "你是一个翻译助理，你的任务是翻译星际争霸视频，你会被提供一个按行分割的英文段落，你需要在保证句意和行数的情况下输出翻译后的文本。"
-
+    print(task)
     previous_length = 0
     for sentence, range_ in tqdm(zip(script_arr, range_arr)):
         # update the range based on previous length