Class DecoderOnlyLanguageModelGenerate

models.py:28–178

Decoder only language modeling in inference mode.

@gin.configurable
class DecoderOnlyLanguageModelGenerate(models.DecoderOnlyLanguageModel):
  """Decoder only language modeling in inference mode."""

  decoder_factory = decoder_stack.DecoderStackGenerate

  num_heads: int = gin.REQUIRED
  head_size: int = gin.REQUIRED

  def get_fake_input(self) -> dict[str, Any]:
    fake_input_dict = super().get_fake_input()
    b = self.task_config.batch_size
    n = self.num_heads
    h = self.head_size
    fake_input_dict.update({
        'dstate': tuple(
            [{
                'current_index': jnp.array([0] * b, dtype=jnp.int32),
                'keys': jnp.zeros((b, 2048, n, h), dtype=jnp.bfloat16),
                'values': jnp.zeros((b, 2048, n, h), dtype=jnp.bfloat16),
                'recurrent_kvq': None,
                'relative_position_bias': jnp.zeros(
                    (b, n, 1, 1024), dtype=jnp.bfloat16
                ),
            }]
            * 12
        ),
        'eos': jnp.zeros([1024], dtype=jnp.bfloat16),
        'mask': jnp.ones([1024], dtype=jnp.bfloat16),
        'length': 1,
        'temperature': 1.0,
    })
    return fake_input_dict

  def __call__(self, inputs: ...) -> tuple[Any, dict[str, Any]]:
    # Make sure this code is not used on untested cases.
    if self.mode not in ['init', 'beam_search']:
      raise ValueError(f'{type(self)} cannot do mode {self.mode}')
    if self.decoder.supports_generate():
      raise ValueError(f'{type(self)}.decoder cannot supports_generate()')

    self.decoder(
        input_tokens=inputs['targets'][:, 0:1],
        target_tokens=None,
        start_of_sequence=inputs['start_of_sequence'],
    )

    b = inputs['targets'].shape[0]
    no_start_of_seq = jnp.array([False] * b, dtype=jnp.bool_)

    # This fn is used in both beam_search or topk_sampling.
    def tokens_to_logits_fn(
        input_token: jnp.ndarray, dstate: tuple[dict[str, jnp.ndarray], ...]
    ) -> tuple[jnp.ndarray, tuple[dict[str, jnp.ndarray], ...]]:
      (logits, dstate, _) = self.decoder(
          input_tokens=input_token,
          target_tokens=None,
          start_of_sequence=no_start_of_seq,
          decoder_state=dstate,