Visualize Gensim word2vec model in TensorBoard

1. '''
2. Output Little-Endian bytes and labels file from gensim model
3. Also outputs necessary json config file portion
4. For use with TensorBoard
5. '''
7. import os
8. import sys
9. import struct
10. import logging
11. import argparse
12. from gensim.models import Word2Vec, Doc2Vec
13. # necessary for seeing logs
15. logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', \
16. level=logging.INFO)
18. parser = argparse.ArgumentParser(description='Convert gensim Word2Vec model to \
19. TensorBoard visualization format')
21. # Required positional argument
22. parser.add_argument('input_file', type=str,
23. help='Input word2vec model file')
24. #parser.add_argument('output_file', type=str,
25. # help='Base output filename of TensorBoard data files')
27. args = parser.parse_args()
29. w2vmodel = Word2Vec.load(args.input_file)
31. try:
32. vocab = w2vmodel.vocab
33. except:
34. vocab = w2vmodel.wv.vocab
36. num_rows = len(vocab)
38. dim = w2vmodel.vector_size
40. try:
41. base_fn = args.input_file.split('/')[-1]
42. except:
43. base_fn = args.input_file.split('\\')[-1]
45. assert len(base_fn) > 0
47. tensor_bytes_out_fn = '%s_%d_%dd_tensors.bytes' % (base_fn, num_rows, dim)
48. tensor_tsv_out_fn = '%s_%d_%dd_tensors.tsv' % (base_fn, num_rows, dim)
49. labels_out_fn = '%s_%d_%dd_labels.tsv' % (base_fn, num_rows, dim)
51. tensor_bytes_out = open(tensor_bytes_out_fn, 'wb')
53. try:
54. labels_out = open(labels_out_fn, 'w', encoding='utf-8')
55. except:
56. labels_out = open(labels_out_fn, 'w')
58. try:
59. tensor_tsv_out = open(tensor_tsv_out_fn, 'w', encoding='utf-8')
60. except:
61. tensor_tsv_out = open(tensor_tsv_out_fn, 'w')
63. labels_out.write('word\tcount\n')
65. for i, wd in enumerate(vocab):
66. if i % 100 == 0:
67. sys.stderr.write('\rDump vocabulary: %d/%d (%.2f%%)' % \
68. (i+1, num_rows, 100.0 * (i+1) / num_rows))
70. floatvals = w2vmodel[wd].tolist()
71. assert dim == len(floatvals)
72. assert '\t' not in wd
74. if i > 0:
75. tensor_tsv_out.write('\n')
77. for f_i, f in enumerate(floatvals):
78. tensor_bytes_out.write(struct.pack('<f', f))
80. if f_i > 0:
81. tensor_tsv_out.write('\t%.8f' % f)
82. else:
83. tensor_tsv_out.write('%.8f' % f)
85. try:
86. labels_out.write('%s\t%s\n' % (wd, vocab[wd].count))
87. except:
88. labels_out.write(('%s\t%s\n' % (wd, vocab[wd].count)).encode('utf-8'))
90. sys.stderr.write('\n')
91. sys.stderr.flush()
93. tensor_bytes_out.close()
94. tensor_tsv_out.close()
95. labels_out.close()
97. # projector file
99. print('''{
100. "embeddings": [
101. {
102. "tensorName": "%s",
103. "tensorShape": [%d, %d],
104. "tensorPath": "%s",
105. "metadataPath": "%s"
106. }
107. ],
108. "modelCheckpointPath": "Demo datasets"
109. }''' % (base_fn, num_rows, dim, tensor_bytes_out_fn, labels_out_fn))

'''
Output Little-Endian bytes and labels file from gensim model
Also outputs necessary json config file portion
For use with TensorBoard
'''

import os
import sys
import struct
import logging
import argparse
from gensim.models import Word2Vec, Doc2Vec
# necessary for seeing logs

logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', \
    level=logging.INFO)

parser = argparse.ArgumentParser(description='Convert gensim Word2Vec model to \
    TensorBoard visualization format')

# Required positional argument
parser.add_argument('input_file', type=str,
                    help='Input word2vec model file')
#parser.add_argument('output_file', type=str,
#                    help='Base output filename of TensorBoard data files')

args = parser.parse_args()

w2vmodel = Word2Vec.load(args.input_file)

try:
    vocab = w2vmodel.vocab
except:
    vocab = w2vmodel.wv.vocab

num_rows = len(vocab)

dim = w2vmodel.vector_size

try:
    base_fn = args.input_file.split('/')[-1]
except:
    base_fn = args.input_file.split('\\')[-1]

assert len(base_fn) > 0

tensor_bytes_out_fn = '%s_%d_%dd_tensors.bytes' % (base_fn, num_rows, dim)
tensor_tsv_out_fn = '%s_%d_%dd_tensors.tsv' % (base_fn, num_rows, dim)
labels_out_fn = '%s_%d_%dd_labels.tsv' % (base_fn, num_rows, dim)

tensor_bytes_out = open(tensor_bytes_out_fn, 'wb')

try:
    labels_out = open(labels_out_fn, 'w', encoding='utf-8')
except:
    labels_out = open(labels_out_fn, 'w')

try:
    tensor_tsv_out = open(tensor_tsv_out_fn, 'w', encoding='utf-8')
except:
    tensor_tsv_out = open(tensor_tsv_out_fn, 'w')

labels_out.write('word\tcount\n')

for i, wd in enumerate(vocab):
    if i % 100 == 0:
        sys.stderr.write('\rDump vocabulary: %d/%d (%.2f%%)' % \
            (i+1, num_rows, 100.0 * (i+1) / num_rows))

    floatvals = w2vmodel[wd].tolist()
    assert dim == len(floatvals)
    assert '\t' not in wd

    if i > 0:
        tensor_tsv_out.write('\n')

    for f_i, f in enumerate(floatvals):
        tensor_bytes_out.write(struct.pack('<f', f))
        
        if f_i > 0:
            tensor_tsv_out.write('\t%.8f' % f)
        else:
            tensor_tsv_out.write('%.8f' % f)

    try:
        labels_out.write('%s\t%s\n' % (wd, vocab[wd].count))
    except:
        labels_out.write(('%s\t%s\n' % (wd, vocab[wd].count)).encode('utf-8'))

sys.stderr.write('\n')
sys.stderr.flush()

tensor_bytes_out.close()
tensor_tsv_out.close()
labels_out.close()

# projector file

print('''{
  "embeddings": [
    {
      "tensorName": "%s",
      "tensorShape": [%d, %d],
      "tensorPath": "%s",
      "metadataPath": "%s"
    }
  ],
  "modelCheckpointPath": "Demo datasets"
}''' % (base_fn, num_rows, dim, tensor_bytes_out_fn, labels_out_fn))