spellchecker(enchant) added, known bug(line 206)

Former-commit-id: 5dea602712d88bf2f0a49321497f5575a86beede

SRT.py

@@ -2,7 +2,6 @@ from datetime import timedelta
 import os
 import whisper
 from csv import reader
-import re
 
 class SRT_segment(object):
     def __init__(self, *args) -> None:
@@ -20,12 +19,12 @@ class SRT_segment(object):
                 self.end_time_str = str(0)+end_time.split('.')[0]+',000'
             else:
                 self.end_time_str = str(0)+end_time.split('.')[0]+','+end_time.split('.')[1][:3]
-            self.source_text = segment['text'][1:]
+            self.source_text = segment['text']
             self.duration = f"{self.start_time_str} --> {self.end_time_str}"
             self.translation = ""
 
         elif isinstance(args[0], list):
-            self.source_text = args[0][2][:-1]
+            self.source_text = args[0][2]
             self.duration = args[0][1]
             self.start_time_str = self.duration.split(" --> ")[0]
             self.end_time_str = self.duration.split(" --> ")[1]
@@ -179,29 +178,91 @@ class SRT_script():
         # TODO: variety of translation
 
         # load term dictionary
-        with open("finetune_data/dict.csv",'r', encoding='utf-8') as f:
-            csv_reader = reader(f)
-            term_dict = {rows[0]:rows[1] for rows in csv_reader}
+        # with open("dict_enzh.csv",'r', encoding='utf-8') as f:
+        #     csv_reader = reader(f)
+        #     term_enzh_dict = {rows[0]:rows[1] for rows in csv_reader}
 
         # change term
         for seg in self.segments:
-            ready_words = re.sub('\n', '\n ', seg.source_text).split(" ")
+            ready_words = seg.source_text.split(" ")
             for i in range(len(ready_words)):
                 word = ready_words[i]
-                if word[-2:] == ".\n" :
-                    if word[:-2].lower() in term_dict :
-                        new_word = word.replace(word[:-2], term_dict.get(word[:-2].lower())) + ' '
-                        ready_words[i] = new_word
-                    else:
-                        ready_words[i] = word + ' '
-                elif word.lower() in term_dict :
-                    new_word = word.replace(word,term_dict.get(word.lower())) + ' '
-                    ready_words[i] = new_word
-                else :
-                    ready_words[i]= word + ' '
-            seg.source_text = re.sub('\n ', '\n', "".join(ready_words))
+                ready_words[i] = self.spell_correction(word, 0)
+                # if word[-2:] == ".\n":
+                #     if word[:-2].lower() in term_enzh_dict:
+                #         new_word = word.replace(word[:-2], term_enzh_dict.get(word[:-2].lower()))
+                #         ready_words[i] = new_word
+                #     else:
+                #         ready_words[i] = word
+                # elif word.lower() in term_enzh_dict:
+                #     new_word = word.replace(word,term_enzh_dict.get(word.lower()))
+                #     ready_words[i] = new_word
+                # else:
+                #     ready_words[i]= word
+            seg.source_text = " ".join(ready_words)
         pass
 
-    
+    def spell_check_term(self):
+        ## known bug: I've will be replaced because i've is not in the dict
+
+        
+        # import enchant
+        # dict = enchant.Dict('en_US')
+        # term_spellDict = enchant.PyPWL('project-t/finetune_data/dict_freq.txt')
+
+        for seg in self.segments:
+            ready_words = seg.source_text.split(" ")
+            for i in range(len(ready_words)):
+                word = ready_words[i]
+                ready_words[i] = self.spell_correction(word, 1)
+                # if word[-2:] == ".\n":
+                #     real_word = word[:-2]
+                #     if not dict.check(real_word.lower()):
+                #         new_word = word.replace(word[:-2], term_spellDict.suggest(real_word.lower())[0])
+                #         ready_words[i] = new_word
+                # elif word[-1:] in [".", "\n", ","]:
+                #     real_word = word[:-1]
+                #     if not dict.check(real_word.lower()):
+                #         new_word = word.replace(word[:-1], term_spellDict.suggest(real_word.lower())[0])
+                #         ready_words[i] = new_word
+                # elif not dict.check(word.lower()):
+                #     new_word = word.replace(word,term_spellDict.suggest(word.lower())[0])
+                #     ready_words[i] = new_word
+            seg.source_text = " ".join(ready_words)
+        pass
+
+    def spell_correction(self, word:str, arg:int):
+        try:
+            arg in [0,1]
+        except ValueError:
+            print('only 0 or 1 for argument')
 
-    
+
+        def uncover(word:str):
+            if word[-2:] == ".\n":
+                real_word = word[:-2].lower()
+                n = -2
+            elif word[-1:] in [".", "\n", ",", "!", "?"]:
+                real_word = word[:-1].lower()
+                n = -1
+            else:
+                real_word = word.lower()
+                n = 0
+            return real_word, len(word)+n
+        
+        real_word = uncover(word)[0]
+        pos = uncover(word)[1]
+        new_word = word
+        if arg == 0:  # term translate mode
+            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)}
+            if real_word in term_enzh_dict:
+                new_word = word.replace(word[:pos], term_enzh_dict.get(real_word))
+        elif arg == 1:  # term spell check mode
+            import enchant
+            dict = enchant.Dict('en_US')
+            term_spellDict = enchant.PyPWL('finetune_data/dict_freq.txt')
+            if not dict.check(real_word):
+                if term_spellDict.suggest(real_word):  # relax spell check
+                    new_word = word.replace(word[:pos],term_spellDict.suggest(real_word)[0])
+        return new_word

finetune_data/{dict.csv → dict_enzh.csv}

@@ -1,4 +1,4 @@
-barracks,兵营
+barracks,兵营
 engineering bay,工程站
 forge,锻炉
 blink,闪现

finetune_data/dict_freq.csv

@@ -0,0 +1,176 @@
+barracks,1
+engineering bay,1
+forge,1
+blink,1
+evolution chamber,1
+cybernetics core,1
+enhanced shockwaves,1
+gravitic boosters,1
+armory,1
+robotics bay,1
+twilight council,1
+fusion core,1
+fleet beacon,1
+factory,1
+ghost academy,1
+infestation pit,1
+robotics facility,1
+stargate,1
+starport,1
+archon,1
+smart servos,1
+gateway,1
+warpgate,1
+immortal,1
+zealot,1
+nydus network,1
+nydus worm,1
+hydralisk,1
+grooved spines,1
+muscular augments,1
+hydralisk den,1
+planetary fortress,1
+battle cruiser,1
+weapon refit,1
+brood lord,1
+greater spire,1
+anabolic synthesis,1
+cyclone,1
+bunker,1
+lurker,1
+seismic spines,1
+adaptive talons,1
+lurker den,1
+widow mine,1
+ground carapace,1
+high templar,1
+shield battery,1
+observer,1
+baneling,1
+centrifugal hooks,1
+baneling nest,1
+raven,1
+combat shield,1
+shield,1
+lair,1
+missile turret,1
+spore crawler,1
+supply depot,1
+overlord,1
+pneumatized carapace,1
+mutalisk,1
+spire,1
+viper,1
+flyer attacks,1
+flyer carapace,1
+tempest,1
+tectonic destabilizers,1
+phoenix,1
+anion pulse-crystals,1
+corruptor,1
+infestor,1
+pathogen glands,1
+zergling,1
+spawning pool,1
+metabolic boost,1
+spine crawler,1
+marauder,1
+ghost,1
+arm silo with nuke,1
+carrier,1
+hellion,1
+hellbat,1
+ravager,1
+nexus,1
+hatchery,1
+command center,1
+neosteel armor,1
+hi-sec auto tracking,1
+ship weapons,1
+charge,1
+liberator,1
+advanced ballistics,1
+melee attacks,1
+colossus,1
+extended thermal lance,1
+creep tumor,1
+tech lab,1
+air armor,1
+air weapons,1
+adrenal glands,1
+mule,1
+infernal pre-igniter,1
+thor,1
+warp prism,1
+gravitic drive,1
+dragoon,1
+cocoon,1
+larva,1
+mothership,1
+burrow,1
+changeling,1
+ultralisk,1
+chitinous plating,1
+ultralisk cavern,1
+drone,1
+scv,1
+queen,1
+banshee,1
+hyperflight rotors,1
+photon cannon,1
+missile attacks,1
+assimilator,1
+extractor,1
+refinery,1
+roach,1
+marine,1
+sensor tower,1
+infantry armor,1
+infantry weapons,1
+hive,1
+psionic storm,1
+templar archives,1
+sentry,1
+ground armor,1
+ground weapons,1
+adept,1
+resonating glaives,1
+reactor,1
+pylon,1
+reaper,1
+drilling claws,1
+swarm host,1
+mag-field accelerator,1
+siege tank,1
+probe,1
+corvid reactor,1
+neural parasite,1
+viking,1
+oracle,1
+broodling,1
+locust,1
+mothership core,1
+orbital command,1
+stimpack,1
+void ray,1
+flux vanes,1
+overseer,1
+ignite afterburners,1
+dark templar,1
+shadow stride,1
+dark shrine,1
+cloaking field,1
+personal cloaking,1
+medivac dropship,1
+vehicle and ship plating,1
+vehicle weapons,1
+war hound,1
+roach warren,1
+tunneling claws,1
+glial reconstitution,1
+concussive shells,1
+stalker,1
+disruptor,1
+zerg,1
+protross,1
+terran,1

finetune_data/dict_freq.txt

@@ -0,0 +1,177 @@
+barracks
+engineering bay
+forge
+blink
+evolution chamber
+cybernetics core
+enhanced shockwaves
+gravitic boosters
+armory
+robotics bay
+twilight council
+fusion core
+fleet beacon
+factory
+ghost academy
+infestation pit
+robotics facility
+stargate
+starport
+archon
+smart servos
+gateway
+warpgate
+immortal
+zealot
+nydus network
+nydus worm
+hydralisk
+grooved spines
+muscular augments
+hydralisk den
+planetary fortress
+battle cruiser
+weapon refit
+brood lord
+greater spire
+anabolic synthesis
+cyclone
+bunker
+lurker
+seismic spines
+adaptive talons
+lurker den
+widow mine
+ground carapace
+high templar
+shield battery
+observer
+baneling
+centrifugal hooks
+baneling nest
+raven
+combat shield
+shield
+lair
+missile turret
+spore crawler
+supply depot
+overlord
+pneumatized carapace
+mutalisk
+spire
+viper
+flyer attacks
+flyer carapace
+tempest
+tectonic destabilizers
+phoenix
+anion pulse-crystals
+corruptor
+infestor
+pathogen glands
+zergling
+spawning pool
+metabolic boost
+spine crawler
+marauder
+ghost
+arm silo with nuke
+carrier
+hellion
+hellbat
+ravager
+nexus
+hatchery
+command center
+neosteel armor
+hi-sec auto tracking
+ship weapons
+charge
+liberator
+advanced ballistics
+melee attacks
+colossus
+extended thermal lance
+creep tumor
+tech lab
+air armor
+air weapons
+adrenal glands
+mule
+infernal pre-igniter
+thor
+warp prism
+gravitic drive
+dragoon
+cocoon
+larva
+mothership
+burrow
+changeling
+ultralisk
+chitinous plating
+ultralisk cavern
+drone
+scv
+queen
+banshee
+hyperflight rotors
+photon cannon
+missile attacks
+assimilator
+extractor
+refinery
+roach
+marine
+sensor tower
+infantry armor
+infantry weapons
+hive
+psionic storm
+templar archives
+sentry
+ground armor
+ground weapons
+adept
+resonating glaives
+reactor
+pylon
+reaper
+drilling claws
+swarm host
+mag-field accelerator
+siege tank
+probe
+corvid reactor
+neural parasite
+viking
+oracle
+broodling
+locust
+mothership core
+orbital command
+stimpack
+void ray
+flux vanes
+overseer
+ignite afterburners
+dark templar
+shadow stride
+dark shrine
+cloaking field
+personal cloaking
+medivac dropship
+vehicle and ship plating
+vehicle weapons
+war hound
+roach warren
+tunneling claws
+glial reconstitution
+concussive shells
+stalker
+disruptor
+zerg
+protross
+terran
+starcraft