new dict utils need prune

Former-commit-id: c3d204c86e52f7ef0958fd59487a12505addea3c

dict_util.py

@@ -52,4 +52,29 @@ with open("../test.csv", "w", encoding='utf-8') as w:
     export_csv_dict(term_dict_sc2,w)
 
 ## for load pickle, just:
-# pickle.load(f)
+# pickle.load(f)
+
+
+def form_dict(src_path,tgt_path) -> dict:
+    final_dict = {}
+    with open(src_path, 'r', encoding="utf-8") as file:
+        src_dict = list(csv.reader(file, delimiter=","))
+    with open(tgt_path, 'r', encoding="utf-8") as file:
+        tgt_dict = list(csv.reader(file, delimiter="," ))
+    for idx, value in enumerate(src_dict):
+        for item in value:
+            final_dict.update({item:tgt_dict[idx]})
+    return final_dict
+
+
+class term_dict(dict):
+    def __init__(self, path, src_lang, tgt_lang) -> None:
+        src_dict = f"{path}/{src_lang}.csv"
+        tgt_dict = f"{path}/{tgt_lang}.csv"
+        super().__init__(form_dict(src_dict, tgt_dict))
+
+
+    def get(self, key:str) -> str:
+        word = self[key][randint(0,len(self[key])-1)]
+        return word
+

domain_dict/SC2/EN.csv

@@ -1,43 +1,43 @@
-barracks,
-zerg,
-protoss,
-terran,
+barracks
+zerg
+protoss
+terran
 engineering bay,engin bay
-forge,
-blink,
-evolution chamber,
+forge
+blink
+evolution chamber
 cybernetics core,cybercore
-enhanced shockwaves,
-gravitic boosters,
-armory,
+enhanced shockwaves
+gravitic boosters
+armory
 robotics bay,robo bay
 twilight council,twilight
-fusion core,
-fleet beacon,
-factory,
-ghost academy,
-infestation pit,
+fusion core
+fleet beacon
+factory
+ghost academy
+infestation pit
 robotics facility,robo
-stargate,
-starport,
-archon,
-smart servos,
-gateway,
-warpgate,
-immortal,
-zealot,
-nydus network,
-nydus worm,
+stargate
+starport
+archon
+smart servos
+gateway
+warpgate
+immortal
+zealot
+nydus network
+nydus worm
 hydralisk,hydra
-grooved spines,
-muscular augments,
+grooved spines
+muscular augments
 hydralisk den,hydra den
-planetary fortress,
-battle cruiser,
-weapon refit,
-brood lord,
-broodling,
-greater spire,
-anabolic synthesis,
-cyclone,
-bunker,
+planetary fortress
+battle cruiser
+weapon refit
+brood lord
+broodling
+greater spire
+anabolic synthesis
+cyclone
+bunker

domain_dict/SC2/ZH.csv

@@ -1,43 +1,43 @@
-兵营,
-虫族,
-神族,
-人族,
+兵营
+虫族
+神族
+人族
 工程站,BE
 BF,锻炉
-闪现,
-进化腔,
+闪现
+进化腔
 BY,赛博核心
-EMP范围,
-ob速度,
-军械库,
+EMP范围
+ob速度
+军械库
 机械研究所,VB
 光影议会,VC
-聚变芯体,
-舰队航标,
-重工厂,
-幽灵军校,
-感染深渊,
+聚变芯体
+舰队航标
+重工厂
+幽灵军校
+感染深渊
 VR,机械台
 神族VS,星门
 星港,人族VS
-白球,
-变形加速,
-传送门,
-折跃门,
-不朽,
-叉叉,
-虫道网络,
-坑道虫,
-刺蛇,
-刺蛇射程,
-刺蛇速度,
-刺蛇塔,
+白球
+变形加速
+传送门
+折跃门
+不朽
+叉叉
+虫道网络
+坑道虫
+刺蛇
+刺蛇射程
+刺蛇速度
+刺蛇塔
 大地堡,行星要塞
-大和,
-大和炮,
-大龙,
-巢虫,
-大龙塔,
-大牛速度,
-导弹车,
-地堡,
+大和
+大和炮
+大龙
+巢虫
+大龙塔
+大牛速度
+导弹车
+地堡

src/srt_util/srt.py

@@ -7,6 +7,7 @@ from datetime import timedelta
 import logging
 import openai
 from tqdm import tqdm
+import dict_util
 
 # punctuation dictionary for supported languages
 punctuation_dict = {
@@ -161,9 +162,13 @@ class SrtScript(object):
         if self.domain != "General":
             if os.path.exists(f"{dict_path}/{self.domain}"):
                 # TODO: load dictionary
+                self.dict = dict_util.term_dict(f"{dict_path}/{self.domain}", src_lang, tgt_lang)
+                print(self.dict["robo"])
                 ...
             else:
-                logging.error(f"domain {self.domain} doesn't exist")
+                logging.error(f"domain {self.domain} doesn't exist, fallback to general domain, this will disable correct_with_force_term and spell_check_term")
+                self.domain = "General"
+
 
     @classmethod
     def parse_from_srt_file(cls, src_lang, tgt_lang, path: str):
@@ -444,23 +449,19 @@ class SrtScript(object):
         if self.domain == "General":
             logging.info("General domain could not perform correct_with_force_term. skip this step.")
             pass
-
-        # 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)
-
-        for word in keywords:
-            for i, seg in enumerate(self.segments):
-                if word in seg.source_text.lower():
-                    seg.source_text = re.sub(fr"({word}es|{word}s?)\b", "{}".format(term_enzh_dict.get(word)),
-                                             seg.source_text, flags=re.IGNORECASE)
-                    logging.info(
-                        "replace term: " + word + " --> " + term_enzh_dict.get(word) + " in time stamp {}".format(
-                            i + 1))
-                    logging.info("source text becomes: " + seg.source_text)
+        else:
+            keywords = list(self.dict.keys())
+            keywords.sort(key=lambda x: len(x), reverse=True)
+
+            for word in keywords:
+                for i, seg in enumerate(self.segments):
+                    if word in seg.source_text.lower():
+                        seg.source_text = re.sub(fr"({word}es|{word}s?)\b", "{}".format(self.dict.get(word)),
+                                                seg.source_text, flags=re.IGNORECASE)
+                        logging.info(
+                            "replace term: " + word + " --> " + self.dict.get(word) + " in time stamp {}".format(
+                                i + 1))
+                        logging.info("source text becomes: " + seg.source_text)
 
     comp_dict = []