This repository has been archived by the owner on May 15, 2024. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 6
/
main.py
75 lines (62 loc) · 2.23 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
# coding:utf-8
import logging
import json
import os
from cotk.wordvector import WordVector, Glove
from myCoTK.dataloader import WizardOfWiki, HollE
from utils import debug, try_cache, cuda_init, Storage
from seq2seq import Seq2seq
def main(args):
logging.basicConfig(filename=0,
level=logging.DEBUG,
format='%(asctime)s %(filename)s[line:%(lineno)d] %(message)s',
datefmt='%H:%M:%S')
if args.debug:
debug()
logging.info(json.dumps(args, indent=2))
cuda_init(args.cuda_num, args.cuda)
volatile = Storage()
volatile.load_exclude_set = args.load_exclude_set
volatile.restoreCallback = args.restoreCallback
if args.dataset == 'WizardOfWiki':
data_class = WizardOfWiki
elif args.dataset == 'HollE':
data_class = HollE
else:
raise ValueError
wordvec_class = WordVector.load_class(args.wvclass)
if wordvec_class is None:
wordvec_class = Glove
if not os.path.exists(args.cache_dir):
os.mkdir(args.cache_dir)
args.cache_dir = os.path.join(args.cache_dir, args.dataset)
if not os.path.exists(args.out_dir):
os.mkdir(args.out_dir)
args.out_dir = os.path.join(args.out_dir, args.dataset)
if not os.path.exists(args.model_dir):
os.mkdir(args.model_dir)
if args.dataset not in args.model_dir:
args.model_dir = os.path.join(args.model_dir, args.dataset)
if args.cache:
dm = try_cache(data_class, (args.datapath,), args.cache_dir)
volatile.wordvec = try_cache(
lambda wv, ez, vl: wordvec_class(wv).load_matrix(ez, vl),
(args.wvpath, args.embedding_size, dm.vocab_list),
args.cache_dir, wordvec_class.__name__)
else:
dm = data_class(args.datapath)
wv = wordvec_class(args.wvpath)
volatile.wordvec = wv.load_matrix(args.embedding_size, dm.vocab_list)
volatile.dm = dm
param = Storage()
param.args = args
param.volatile = volatile
model = Seq2seq(param)
if args.mode == "train":
model.train_process()
elif args.mode == "test":
model.test_process()
elif args.mode == 'dev':
model.test_dev()
else:
raise ValueError("Unknown mode")