Class LanguageModelInference

lm_inference.py:36–189

Meliad wrapper for LM inference.

class LanguageModelInference:
  """Meliad wrapper for LM inference."""

  def __init__(self, vocab_path: str, load_dir: str, mode='beam_search'):
    self.vocab = t5.data.SentencePieceVocabulary(vocab_path)

    # This task won't be pulling from a dataset.
    def null_iter_fn() -> None:
      return None

    process_summaries_f = inference_utils.models.process_summaries_function(
        self.vocab
    )

    trainer = inference_utils.training_loop.Trainer(
        get_training_dataset_iterator=null_iter_fn,
        get_test_dataset_iterator=None,
        pretty_print_input_function=None,
        process_summaries_function=process_summaries_f,
        load_dir=load_dir,
        workdir='',  # Don't log or save checkpoints.
        replicate_mode=False,
    )  # Run on a single device at batch size 1.
    self.trainer = trainer

    # Create and initialize the model.
    (tstate, _, imodel, prngs) = trainer.initialize_model()
    self.imodel = imodel
    self.batch_size = imodel.task_config.batch_size

    self.n = imodel.num_heads
    self.h = imodel.head_size

    # Create an inference task.
    writers = {}
    self.task = trainer.create_training_task(mode, imodel, prngs, writers)  # pylint: disable=too-many-function-args

    # Register any additional actions.
    # Actions are cleared first for use with colab.
    inference_utils.training_loop.clear_interstep_callbacks()
    inference_utils.training_loop.register_interstep_callbacks()
    self.tstate = tstate

    # some default parameters.
    eos = [0] * 1024
    for idx in self.encode_list(['.', ';']):
      eos[idx] = 1

    self.eos = np.array(eos, dtype=np.bfloat16)
    self.mask = jax.numpy.ones([1024], dtype=np.bfloat16)

  def decode(self, ids: list[int]) -> str:
    return self.vocab.decode(ids)

  def decode_list(self, tokens: list[int]) -> list[str]:
    return [self.decode([tok]) for tok in tokens]

  def encode(self, inputs_str: str) -> list[int]: