Method beam_decode

lm_inference.py:146–189

Beam search.

(
      self,
      inputs: str,
      eos_tokens: np.ndarray = None,
      mask_tokens: np.ndarray = None,
      dstate: dict[str, np.ndarray] = None,
  )

    return metrics_np

  def beam_decode(
      self,
      inputs: str,
      eos_tokens: np.ndarray = None,
      mask_tokens: np.ndarray = None,
      dstate: dict[str, np.ndarray] = None,
  ) -> MetricsOutput:
    """Beam search."""
    inputs = jax.numpy.array([self.vocab.encode(inputs)] * self.batch_size)

    eos = self.eos
    if eos_tokens is not None:
      eos_ids = self.encode_list(eos_tokens)
      eos = np.array(
          [1 if idx in eos_ids else 0 for idx in range(1024)], dtype=np.bfloat16
      ).reshape((1, 1, 1024))

    mask = self.mask
    if mask_tokens is not None:
      mask_ids = self.encode_list(mask_tokens)
      mask = np.array(
          [0 if idx in mask_ids else 1 for idx in range(1024)],
          dtype=np.bfloat16,
      ).reshape((1, 1, 1024))

    metrics_np = self.call(inputs, dstate=dstate, eos=eos, mask=mask)

    finished_seqs = metrics_np['finished_seqs']
    finished_scores = metrics_np['finished_scores']

    seqs = []
    scores = []
    for seq, score in zip(finished_seqs, finished_scores):
      seq = self.decode(seq[1:])
      seqs.append(seq)
      scores.append(score)

    return {
        'finished_seqs': finished_seqs,
        'finished_scores': finished_scores,
        'seqs_str': seqs,
        'scores': scores,
        'dstate': metrics_np['dstate'],
    }