toy_generator.py 文件源码

def get_neighbors(words, word, window_size=2):
  if type(word) == str:
    idx = words.index(word)
  elif type(word) == int:
    idx = word
  else:
    raise Exception(" [!] Invalid type for word: %s" % type(word))

  if idx < window_size:
    ans = words[-(window_size - idx):] + words[:idx + window_size + 1]
  elif idx >= len(words) - window_size:
    ans = words[idx-window_size:] + words[:window_size + idx - len(words) + 1]
  else:
    ans = words[idx-window_size:idx+window_size+1]

  for _ in xrange(15):
    if random.random() < 0.1:
      ans.append(random.choice(ans))

  random.shuffle(ans)
  return ans