code2vec/code2vec.py at master · cmusv/code2vec · GitHub

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
fromvocabulariesimportVocabType
fromconfigimportConfig
frominteractive_predictimportInteractivePredictor
frommodel_baseimportCode2VecModelBase


defload_model_dynamically(config: Config) ->Code2VecModelBase:
assertconfig.DL_FRAMEWORKin {'tensorflow', 'keras'}
ifconfig.DL_FRAMEWORK=='tensorflow':
fromtensorflow_modelimportCode2VecModel
elifconfig.DL_FRAMEWORK=='keras':
fromkeras_modelimportCode2VecModel
returnCode2VecModel(config)


if__name__=='__main__':
config=Config(set_defaults=True, load_from_args=True, verify=True)

model=load_model_dynamically(config)
config.log('Done creating code2vec model')

ifconfig.is_training:
model.train()
ifconfig.SAVE_W2VisnotNone:
model.save_word2vec_format(config.SAVE_W2V, VocabType.Token)
config.log('Origin word vectors saved in word2vec text format in: %s'%config.SAVE_W2V)
ifconfig.SAVE_T2VisnotNone:
model.save_word2vec_format(config.SAVE_T2V, VocabType.Target)
config.log('Target word vectors saved in word2vec text format in: %s'%config.SAVE_T2V)
if (config.is_testingandnotconfig.is_training) orconfig.RELEASE:
eval_results=model.evaluate()
ifeval_resultsisnotNone:
config.log(
str(eval_results).replace('topk', 'top{}'.format(config.TOP_K_WORDS_CONSIDERED_DURING_PREDICTION)))
ifconfig.PREDICT:
predictor=InteractivePredictor(config, model)
predictor.predict()
model.close_session()