blob: 658051fe9b09085544a61182379a6d3e3846d700 (
plain)
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
|
import heapq
import operator
import sys
class Vocab(object):
def __init__(self, words=(), unk="<unk>"):
self.words = []
self.word_index = {}
self.insert_word(unk)
self.unk = self.word_index[unk]
for word in words:
self.insert_word(word)
def from_counts(self, counts, size, unk="<unk>"):
# Keep only most frequent words
q = [(-count, word) for (word, count) in counts.iteritems()]
heapq.heapify(q)
inserted = 0
while len(self.words) < size and len(q) > 0:
_, word = heapq.heappop(q)
inserted += 1
if word not in self.word_index:
self.insert_word(word)
return inserted
def insert_word(self, word):
i = len(self.words)
self.words.append(word)
self.word_index[word] = i
def lookup_word(self, word):
return self.word_index.get(word, self.unk)
|
