text: """
  A script to embed every phrase in a dataset as a dense vector, then
  to find the top-k neighbors of each phrase according to cosine
  similarity.
  
  1. Install missing dependencies.
  
      # More details: https://github.com/facebookresearch/faiss/blob/master/INSTALL.md
      conda install faiss-cpu -c pytorch
  
  2. Prepare data. For example, the chunking dataset from CoNLL 2000.
  
      wget https://www.clips.uantwerpen.be/conll2000/chunking/train.txt.gz
      gunzip train.txt.gz
      python diora/misc/convert_conll_to_jsonl.py --path train.txt > conll-train.jsonl
  
  3. Run this script.
  
      python diora/scripts/phrase_embed.py \
          --batch_size 10 \
          --emb w2v \
          --embeddings_path ~/data/glove.6B/glove.6B.50d.txt \
          --hidden_dim 50 \
          --log_every_batch 100 \
          --save_after 1000 \
          --data_type conll_jsonl \
          --validation_path ./conll-train.jsonl \
          --validation_filter_length 10
  
  Can control the number of neighbors to show with the `--k_top` flag.
  
  Can control the number of candidates to consider with `--k_candidates` flag.
  
  """
  
  
  import json
  import types
  import itertools
  
  import torch
  
  import numpy as np
  
  from train import argument_parser, parse_args, configure
  from train import get_validation_dataset, get_validation_iterator
  from train import build_net
  
  from diora.logging.configuration import get_logger
  
  try:
      import faiss
      from faiss import normalize_L2
  except:
      print('Could not import `faiss`, which is used to find nearest neighbors.')
  
  
  def get_cell_index(entity_labels, i_label=0, i_pos=1, i_size=2):
      def helper():
          for i, lst in enumerate(entity_labels):
              for el in lst:
                  if el is None:
                      continue
                  pos = el[i_pos]
                  size = el[i_size]
                  label = el[i_label]
                  yield (i, pos, size, label)
      lst = list(helper())
      if len(lst) == 0:
          return None, []
      batch_index = [x[0] for x in lst]
      positions = [x[1] for x in lst]
      sizes = [x[2] for x in lst]
      labels = [x[3] for x in lst]
  
      return batch_index, positions, sizes, labels
  
  
  def get_many_cells(diora, chart, batch_index, positions, sizes):
      cells = []
      length = diora.length
  
      idx = []
      for bi, pos, size in zip(batch_index, positions, sizes):
          level = size - 1
          offset = diora.index.get_offset(length)[level]
          absolute_pos = offset + pos
          idx.append(absolute_pos)
  
      cells = chart[batch_index, idx]
  
      return cells
  
  
  def get_many_phrases(batch, batch_index, positions, sizes):
      batch = batch.tolist()
      lst = []
      for bi, pos, size in zip(batch_index, positions, sizes):
          phrase = tuple(batch[bi][pos:pos+size])
          lst.append(phrase)
      return lst
  
  
  class BatchRecorder(object):
      def __init__(self, dtype={}):
          super(BatchRecorder, self).__init__()
          self.cache = {}
          self.dtype = dtype
          self.dtype2flatten = {
              'list': self._flatten_list,
              'np': self._flatten_np,
              'torch': self._flatten_torch,
          }
  
      def _flatten_list(self, v):
          return list(itertools.chain(*v))
  
      def _flatten_np(self, v):
          return np.concatenate(v, axis=0)
  
      def _flatten_torch(self, v):
          return torch.cat(v, 0).cpu().data.numpy()
  
      def get_flattened_result(self):
          def helper():
              for k, v in self.cache.items():
                  flatten = self.dtype2flatten[self.dtype.get(k, 'list')]
                  yield k, flatten(v)
          return {k: v for k, v in helper()}
              
      def record(self, **kwargs):
          for k, v in kwargs.items():
              self.cache.setdefault(k, []).append(v)
  
  
  class Index(object):
      def __init__(self, dim=None):
          super(Index, self).__init__()
          self.D, self.I = None, None
          self.index = faiss.IndexFlatIP(dim)
  
      def add(self, vecs):
          self.index.add(vecs)
  
      def cache(self, vecs, k):
          self.D, self.I = self.index.search(vecs, k)
  
      def topk(self, q, k):
          for j in range(k):
              idx = self.I[q][j]
              dist = self.D[q][j]
              yield idx, dist
  
  
  class NearestNeighborsLookup(object):
      def __init__(self):
          super(NearestNeighborsLookup, self).__init__()
  
  
  def run(options):
      logger = get_logger()
  
      validation_dataset = get_validation_dataset(options)
      validation_iterator = get_validation_iterator(options, validation_dataset)
      word2idx = validation_dataset['word2idx']
      embeddings = validation_dataset['embeddings']
  
      idx2word = {v: k for k, v in word2idx.items()}
  
      logger.info('Initializing model.')
      trainer = build_net(options, embeddings, validation_iterator)
      diora = trainer.net.diora
  
      # 1. Get all relevant phrase vectors.
  
      dtype = {
          'example_ids': 'list',
          'labels': 'list',
          'positions': 'list',
          'sizes': 'list',
          'phrases': 'list',
          'inside': 'torch',
          'outside': 'torch',
      }
      batch_recorder = BatchRecorder(dtype=dtype)
  
      ## Eval mode.
      trainer.net.eval()
  
      batches = validation_iterator.get_iterator(random_seed=options.seed)
  
      logger.info('Beginning to embed phrases.')
  
      with torch.no_grad():
          for i, batch_map in enumerate(batches):
              sentences = batch_map['sentences']
              batch_size = sentences.shape[0]
              length = sentences.shape[1]
  
              # Skips very short examples.
              if length <= 2:
                  continue
  
              _ = trainer.step(batch_map, train=False, compute_loss=False)
  
              entity_labels = batch_map['entity_labels']
              batch_index, positions, sizes, labels = get_cell_index(entity_labels)
  
              # Skip short phrases.
              batch_index = [x for x, y in zip(batch_index, sizes) if y >= 2]
              positions = [x for x, y in zip(positions, sizes) if y >= 2]
              labels = [x for x, y in zip(labels, sizes) if y >= 2]
              sizes = [y for y in sizes if y >= 2]
  
              cell_index = (batch_index, positions, sizes)
  
              batch_result = {}
              batch_result['example_ids'] = [batch_map['example_ids'][idx] for idx in cell_index[0]]
              batch_result['labels'] = labels
              batch_result['positions'] = cell_index[1]
              batch_result['sizes'] = cell_index[2]
              batch_result['phrases'] = get_many_phrases(sentences, *cell_index)
              batch_result['inside'] = get_many_cells(diora, diora.inside_h, *cell_index)
              batch_result['outside'] = get_many_cells(diora, diora.outside_h, *cell_index)
  
              batch_recorder.record(**batch_result)
  
      result = batch_recorder.get_flattened_result()
  
      # 2. Build an index of nearest neighbors.
  
      vectors = np.concatenate([result['inside'], result['outside']], axis=1)
      normalize_L2(vectors)
  
      index = Index(dim=vectors.shape[1])
      index.add(vectors)
      index.cache(vectors, options.k_candidates)
  
      # 3. Print a summary.
  
      example_ids = result['example_ids']
      phrases = result['phrases']
  
      assert len(example_ids) == len(phrases)
      assert len(example_ids) == vectors.shape[0]
  
      def stringify(phrase):
          return ' '.join([idx2word[idx] for idx in phrase])
  
      for i in range(vectors.shape[0]):
          topk = []
  
          for j, score in index.topk(i, options.k_candidates):
              # Skip same example.
              if example_ids[i] == example_ids[j]:
                  continue
              # Skip string match.
              if phrases[i] == phrases[j]:
                  continue
              topk.append((j, score))
              if len(topk) == options.k_top:
                  break
          assert len(topk) == options.k_top, 'Did not find enough valid candidates.'
  
          # Print.
          print('[query] example_id={} phrase={}'.format(
              example_ids[i], stringify(phrases[i])))
          for rank, (j, score) in enumerate(topk):
              print('rank={} score={:.3f} example_id={} phrase={}'.format(
                  rank, score, example_ids[j], stringify(phrases[j])))
  
  
  if __name__ == '__main__':
      parser = argument_parser()
      parser.add_argument('--k_candidates', default=100, type=int)
      parser.add_argument('--k_top', default=3, type=int)
      options = parse_args(parser)
      configure(options)
  
      run(options)
  