Merge branch 'eason/main' of https://github.com/project-kxkg/project-t into eason/main

finetune_data/test.csv

@@ -0,0 +1,193 @@
+barracks,['兵营']
+zerg,['虫族']
+protoss,['神族']
+terran,['人族']
+engineering bay,['工程站']
+forge,['锻炉']
+blink,['闪现']
+evolution chamber,['进化腔']
+cybernetics core,['控制芯核']
+cybercore,['控制芯核']
+enhanced shockwaves,['EMP范围']
+gravitic boosters,['ob速度']
+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,['大龙']
+broodling,"['巢虫', '小虫子']"
+greater spire,['大龙塔']
+anabolic synthesis,['大牛速度']
+cyclone,['导弹车']
+bunker,['地堡']
+lurker,['地刺']
+seismic spines,['地刺射程']
+adaptive talons,['地刺速埋']
+lurker den,['地刺塔']
+widow mine,['地雷']
+ground carapace,['地面单位甲壳等级']
+high templar,['电兵']
+shield battery,['电池']
+observer,"['叮当', 'OB']"
+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,['炮台']
+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,['先知']
+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,['自爆球']
+protross,['神族']
+gas,['气矿']
+hatch,['基地']
+supply,['人口']
+macro,['运营']
+natural,['二矿']
+proxy,['前置']
+kite,['甩']
+cyber core,['控制芯核']
+saturate,['满采']
+nydusworm,['虫洞']
+float,['余钱']
+transfuse,['加血']
+chrono boost,['加速']
+war prism,['棱镜']

pipeline.py

@@ -9,6 +9,7 @@ import whisper
 from srt2ass import srt2ass
 import logging
 from datetime import datetime
+import torch
 
 import subprocess
 
@@ -109,7 +110,10 @@ def get_srt_class(srt_file_en, result_path, video_name, audio_path, audio_file =
 
             # use stable-whisper
             elif method == "stable":
-                model = stable_whisper.load_model(whisper_model)
+
+                # use cuda if available
+                devices = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 
+                model = stable_whisper.load_model(whisper_model, device = devices)
                 transcript = model.transcribe(audio_path, regroup = False, initial_prompt="Hello, welcome to my lecture. Are you good my friend?")
                 (
                     transcript
@@ -265,7 +269,9 @@ 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')
+    if not os.path.exists(args.log_dir):
+        os.makedirs(args.log_dir)
+    logging.basicConfig(level=logging.INFO, handlers=[logging.FileHandler("{}/{}_{}.log".format(args.log_dir, VIDEO_NAME, datetime.now().strftime("%m%d%Y_%H%M%S")), 'w', encoding='utf-8')])
     logging.info("---------------------Video Info---------------------")
     logging.info("Video name: {}, translation model: {}, video link: {}".format(VIDEO_NAME, args.model_name, args.link))