add log system; fix correct_with_force_term()

Former-commit-id: 9b687962fba78e4b1ad59bd9c4c84bdf4757f00b

.gitignore

@@ -3,7 +3,9 @@ __pycache__/
 .idea/
 downloads/
 results/
+logs/
 test/
+.vscode/
 test.py
 test.srt
 test.txt

SRT.py

@@ -3,7 +3,7 @@ import re
 from copy import copy, deepcopy
 from csv import reader
 from datetime import timedelta
-
+import logging
 import openai
 
 
@@ -72,7 +72,7 @@ class SRT_segment(object):
         :param other: Another segment that is strictly next to added segment.
         :return: new segment of the two sub-segments
         """
-        # assert other.start_ms == self.end_ms, f"cannot merge discontinuous segments."
+
         result = deepcopy(self)
         result.source_text += f' {other.source_text}'
         result.translation += f' {other.translation}'
@@ -143,6 +143,7 @@ class SRT_script():
         improper segmentation from openai-whisper.
         :return: None
         """
+        logging.info("Forming whole sentences...")
         merge_list = []  # a list of indices that should be merged e.g. [[0], [1, 2, 3, 4], [5, 6], [7]]
         sentence = []
         for i, seg in enumerate(self.segments):
@@ -155,6 +156,8 @@ class SRT_script():
 
         segments = []
         for idx_list in merge_list:
+            if len(idx_list) > 1:
+                logging.info("merging segments: %s", idx_list)
             segments.append(self.merge_segs(idx_list))
 
         self.segments = segments
@@ -166,6 +169,7 @@ class SRT_script():
         """
         for i, seg in enumerate(self.segments):
             seg.remove_trans_punc()
+        logging.info("Removed punctuation in translation.")
 
     def set_translation(self, translate: str, id_range: tuple, model, video_name, video_link=None):
         start_seg_id = id_range[0]
@@ -351,21 +355,24 @@ class SRT_script():
         return result_list
 
     def check_len_and_split(self, text_threshold=30, time_threshold=1.0):
-        # DEPRECATED
         # if sentence length >= threshold and sentence duration > time_threshold, split this segments to two
+        logging.info("performing check_len_and_split")
         segments = []
-        for seg in self.segments:
+        for i, seg in enumerate(self.segments):
             if len(seg.translation) > text_threshold and (seg.end - seg.start) > time_threshold:
                 seg_list = self.split_seg(seg, text_threshold, time_threshold)
+                logging.info("splitting segment {} in to {} parts".format(i+1, len(seg_list)))
                 segments += seg_list
             else:
                 segments.append(seg)
 
         self.segments = segments
+        logging.info("check_len_and_split finished")
 
         pass
 
     def check_len_and_split_range(self, range, text_threshold=30, time_threshold=1.0):
+        # DEPRECATED
         # if sentence length >= text_threshold, split this segments to two
         start_seg_id = range[0]
         end_seg_id = range[1]
@@ -384,28 +391,24 @@ class SRT_script():
 
     def correct_with_force_term(self):
         ## force term correction
-
+        logging.info("performing force term correction")
         # load term dictionary
         with open("./finetune_data/dict_enzh.csv", 'r', encoding='utf-8') as f:
             term_enzh_dict = {rows[0]: rows[1] for rows in reader(f)}
+        
+        keywords = list(term_enzh_dict.keys())
+        keywords.sort(key=lambda x: len(x), reverse=True)
 
-        # change term
-        for seg in self.segments:
-            ready_words = seg.source_text.split(" ")
-            for i in range(len(ready_words)):
-                word = ready_words[i]
-                [real_word, pos] = self.get_real_word(word)
-                if real_word in term_enzh_dict:
-                    new_word = word.replace(word[:pos], term_enzh_dict.get(real_word))
-                else:
-                    new_word = word
-                ready_words[i] = new_word
-            seg.source_text = " ".join(ready_words)
-        pass
+        for word in keywords:
+            for i, seg in enumerate(self.segments):
+                if word in seg.source_text.lower():
+                    seg.source_text = seg.source_text.lower().replace(word, term_enzh_dict.get(word))
+                    logging.info("replace term: " + word + " --> " + term_enzh_dict.get(word) + " in time stamp {}".format(i+1))
+                    logging.info("source text becomes: " + seg.source_text)
 
     def spell_check_term(self):
         ## known bug: I've will be replaced because i've is not in the dict
-
+        logging.info("performing spell check")
         import enchant
         dict = enchant.Dict('en_US')
         term_spellDict = enchant.PyPWL('./finetune_data/dict_freq.txt')
@@ -419,10 +422,10 @@ class SRT_script():
                     suggest = term_spellDict.suggest(real_word)
                     if suggest and enchant.utils.levenshtein(word, suggest[0]) < (len(word)+len(suggest[0]))/4:  # relax spell check
                         
-                        with open("dislog.log","a") as log:
-                            if not os.path.exists("dislog.log"):
-                                log.write("word \t suggest \t levenshtein \n")
-                            log.write(word + "\t" + suggest[0] + "\t" + str(enchant.utils.levenshtein(word, suggest[0]))+'\n')                            
+                        # with open("dislog.log","a") as log:
+                        #     if not os.path.exists("dislog.log"):
+                        #         log.write("word \t suggest \t levenshtein \n")
+                        logging.info(word + "\t" + suggest[0] + "\t" + str(enchant.utils.levenshtein(word, suggest[0]))+'\n')                            
                         #print(word + ":" + suggest[0] + ":---:levenshtein:" + str(enchant.utils.levenshtein(word, suggest[0])))                                                
                         new_word = word.replace(word[:pos],suggest[0])
                     else:
@@ -518,11 +521,13 @@ class SRT_script():
         pass
 
     def write_srt_file_translate(self, path: str):
+        logging.info("writing to " + path)
         with open(path, "w", encoding='utf-8') as f:
             f.write(self.reform_trans_str())
         pass
 
     def write_srt_file_bilingual(self, path: str):
+        logging.info("writing to " + path)
         with open(path, "w", encoding='utf-8') as f:
             f.write(self.form_bilingual_str())
         pass

finetune_data/dict_enzh.csv

@@ -1,4 +1,7 @@
 barracks,兵营
+zerg,虫族
+protoss,神族
+terran,人族
 engineering bay,工程站
 forge,锻炉
 blink,闪现

pipeline.py

@@ -7,6 +7,8 @@ from SRT import SRT_script
 import stable_whisper
 import whisper
 from srt2ass import srt2ass
+import logging
+from datetime import datetime
 
 import subprocess
 
@@ -22,6 +24,7 @@ def parse_args():
     parser.add_argument("--output_dir", help="translate result path", default='./results', type=str, required=False)
     parser.add_argument("--video_name", help="video name, if use video link as input, the name will auto-filled by youtube video name", default='placeholder', type=str, required=False)
     parser.add_argument("--model_name", help="model name only support gpt-4 and gpt-3.5-turbo", type=str, required=False, default="gpt-4") # default change to gpt-4
+    parser.add_argument("--log_dir", help="log path", default='./logs', type=str, required=False)
     parser.add_argument("-only_srt", help="set script output to only .srt file", action='store_true')
     parser.add_argument("-v", help="auto encode script with video", action='store_true')
     args = parser.parse_args()
@@ -30,6 +33,10 @@ def parse_args():
 
 def get_sources(args, download_path, result_path, video_name):
     # get source audio
+    audio_path = None
+    audio_file = None
+    video_path = None
+
     if args.link is not None and args.video_file is None:
         # Download audio from YouTube
         video_link = args.link
@@ -198,6 +205,7 @@ def get_response(model_name, sentence):
         
 # Translate and save
 def translate(srt, script_arr, range_arr, model_name, video_name, video_link):
+    logging.info("start translating...")
     previous_length = 0
     for sentence, range in tqdm(zip(script_arr, range_arr)):
         # update the range based on previous length
@@ -205,12 +213,14 @@ def translate(srt, script_arr, range_arr, model_name, video_name, video_link):
 
         # 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 = get_response(model_name, sentence)
             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.")
                 time.sleep(30)
@@ -227,6 +237,7 @@ def main():
         exit()
 
     # set up
+    start_time = time.time()
     openai.api_key = os.getenv("OPENAI_API_KEY")
     DOWNLOAD_PATH = args.download
     if not os.path.exists(DOWNLOAD_PATH):
@@ -237,7 +248,7 @@ def main():
     RESULT_PATH = args.output_dir
     if not os.path.exists(RESULT_PATH):
         os.mkdir(RESULT_PATH)
-
+    
     # set video name as the input file name if not specified
     if args.video_name == 'placeholder' :
         # set video name to upload file name
@@ -254,25 +265,34 @@ def main():
 
     audio_path, audio_file, video_path, VIDEO_NAME = get_sources(args, DOWNLOAD_PATH, RESULT_PATH, VIDEO_NAME)
 
+    logging.basicConfig(level=logging.INFO, handlers=[logging.FileHandler("{}/{}_{}.log".format(args.log_dir, VIDEO_NAME, datetime.now().strftime("%m%d%Y_%H%M%S")))], encoding='utf-8')
+    logging.info("---------------------Video Info---------------------")
+    logging.info("Video name: {}, translation model: {}, video link: {}".format(VIDEO_NAME, args.model_name, args.link))
+
     srt_file_en, srt = get_srt_class(args.srt_file, RESULT_PATH, VIDEO_NAME, audio_path, audio_file)
 
     # SRT class preprocess
+    logging.info("---------------------Start Preprocessing SRT class---------------------")
+    srt.write_srt_file_src(srt_file_en)
     srt.form_whole_sentence()
-    srt.spell_check_term()
+    # srt.spell_check_term()
     srt.correct_with_force_term()
-    srt.write_srt_file_src(srt_file_en)
+    processed_srt_file_en = srt_file_en.split('.srt')[0] + '_processed.srt'
+    srt.write_srt_file_src(processed_srt_file_en)
     script_input = srt.get_source_only()
 
     # write ass
     if not args.only_srt:
-        assSub_en = srt2ass(srt_file_en, "default", "No", "Modest")
-        print('ASS subtitle saved as: ' + assSub_en)
+        logging.info("write English .srt file to .ass")
+        assSub_en = srt2ass(processed_srt_file_en, "default", "No", "Modest")
+        logging.info('ASS subtitle saved as: ' + assSub_en)
 
     script_arr, range_arr = script_split(script_input)
-
+    logging.info("---------------------Start Translation--------------------")
     translate(srt, script_arr, range_arr, args.model_name, VIDEO_NAME, args.link)
 
     # SRT post-processing
+    logging.info("---------------------Start Post-processing SRT class---------------------")
     srt.check_len_and_split()
     srt.remove_trans_punctuation()
     srt.write_srt_file_translate(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt")
@@ -280,16 +300,20 @@ def main():
 
     # write ass
     if not args.only_srt:
+        logging.info("write Chinese .srt file to .ass")
         assSub_zh = srt2ass(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt", "default", "No", "Modest")
-        print('ASS subtitle saved as: ' + assSub_zh)
+        logging.info('ASS subtitle saved as: ' + assSub_zh)
 
     # encode to .mp4 video file
     if args.v:
+        logging.info("encoding video file")
         if args.only_srt:
             os.system(f'ffmpeg -i {video_path} -vf "subtitles={RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt" {RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}.mp4')
         else:
             os.system(f'ffmpeg -i {video_path} -vf "subtitles={RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.ass" {RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}.mp4')
 
+    end_time = time.time()
+    logging.info("Pipeline finished, time duration:{}".format(start_time - end_time))
 
 if __name__ == "__main__":
     main()