google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
"""Client for looking up word2vec Embeddings."""
17

18
from typing import List
19

20
from gensim.models import KeyedVectors
21
import numpy as np
22

23

24
class Word2VecClient(object):
25
  """API for looking up word2vec vectors for words."""
26

27
  def __init__(self, path):
28
    self.path = path
29

30
    self.keyed_vectors = KeyedVectors.load_word2vec_format(path, binary=True)
31

32
    test_vec = self.keyed_vectors['test']
33
    assert test_vec is not None
34

35
    self.vec_len = len(test_vec)
36
    self.empty_vec = np.array([0.0] * self.vec_len)
37

38
  def lookup(self, words):
39
    """Look up a list of words and return list of wordvecs."""
40
    res = []
41

42
    for word in words:
43
      try:
44
        vec = self.keyed_vectors[word]
45
      except KeyError:
46
        vec = self.empty_vec
47
      res.append(vec)
48

49
    return res
50