-
Notifications
You must be signed in to change notification settings - Fork 6
/
Copy pathword2vec2tensor.py
executable file
·88 lines (74 loc) · 3.5 KB
/
word2vec2tensor.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
76
77
78
79
80
81
82
83
84
85
86
87
88
#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright (C) 2016 Loreto Parisi <[email protected]>
# Copyright (C) 2016 Silvio Olivastri <[email protected]>
# Copyright (C) 2016 Radim Rehurek <[email protected]>
"""
USAGE: $ python -m gensim.scripts.word2vec2tensor --input <Word2Vec model file> --output <TSV tensor filename prefix> [--binary] <Word2Vec binary flag>
Where:
<Word2Vec model file>: Input Word2Vec model
<TSV tensor filename prefix>: 2D tensor TSV output file name prefix
<Word2Vec binary flag>: Set True if Word2Vec model is binary. Defaults to False.
Output:
The script will create two TSV files. A 2d tensor format file, and a Word Embedding metadata file. Both files will
us the --output file name as prefix
This script is used to convert the word2vec format to Tensorflow 2D tensor and metadata formats for Embedding Visualization
To use the generated TSV 2D tensor and metadata file in the Projector Visualizer, please
1) Open http://projector.tensorflow.org/.
2) Choose "Load Data" from the left menu.
3) Select "Choose file" in "Load a TSV file of vectors." and choose you local "_tensor.tsv" file
4) Select "Choose file" in "Load a TSV file of metadata." and choose you local "_metadata.tsv" file
For more information about TensorBoard TSV format please visit:
https://www.tensorflow.org/versions/master/how_tos/embedding_viz/
"""
import os
import sys
import random
import logging
import argparse
import gensim
logger = logging.getLogger(__name__)
def word2vec2tensor(word2vec_model_path,tensor_filename, binary=False):
'''
Convert Word2Vec mode to 2D tensor TSV file and metadata file
Args:
param1 (str): word2vec model file path
param2 (str): filename prefix
param2 (bool): set True to use a binary Word2Vec model, defaults to False
'''
#model = gensim.models.Word2Vec.load_word2vec_format(word2vec_model_path, binary=binary)
model = gensim.models.Word2Vec.load(word2vec_model_path)
outfiletsv = tensor_filename + '_tensor.tsv'
outfiletsvmeta = tensor_filename + '_metadata.tsv'
with open(outfiletsv, 'w+') as file_vector:
with open(outfiletsvmeta, 'w+') as file_metadata:
for word in model.wv.index2word:
print(word)
file_metadata.write(str(word) + '\n')
vector_row = '\t'.join(map(str, model[word]))
file_vector.write(vector_row + '\n')
logger.info("2D tensor file saved to %s" % outfiletsv)
logger.info("Tensor metadata file saved to %s" % outfiletsvmeta)
if __name__ == "__main__":
logging.basicConfig(format='%(asctime)s : %(threadName)s : %(levelname)s : %(message)s', level=logging.INFO)
logging.root.setLevel(level=logging.INFO)
logger.info("running %s", ' '.join(sys.argv))
# check and process cmdline input
program = os.path.basename(sys.argv[0])
if len(sys.argv) < 2:
print(globals()['__doc__'] % locals())
sys.exit(1)
parser = argparse.ArgumentParser()
parser.add_argument(
"-i", "--input", required=True,
help="Input word2vec model")
parser.add_argument(
"-o", "--output", required=True,
help="Output tensor file name prefix")
parser.add_argument( "-b", "--binary",
required=False,
help="If word2vec model in binary format, set True, else False")
args = parser.parse_args()
word2vec2tensor(args.input, args.output, args.binary)
logger.info("finished running %s", program)