Process a new input.
(
self,
model_params: Params,
decode_state: DecodeState,
decode_slot: int,
input_tokens_padded: jax.Array,
input_true_length: int,
rng: PRNGKeyType,
return_prompt_logp: bool = False,
)
| 108 | return self._process_compiled(params, input_padding) |
| 109 | |
| 110 | def process( |
| 111 | self, |
| 112 | model_params: Params, |
| 113 | decode_state: DecodeState, |
| 114 | decode_slot: int, |
| 115 | input_tokens_padded: jax.Array, |
| 116 | input_true_length: int, |
| 117 | rng: PRNGKeyType, |
| 118 | return_prompt_logp: bool = False, |
| 119 | ) -> tuple[engine_api.ResultTokens, DecodeState]: |
| 120 | """Process a new input.""" |
| 121 | |
| 122 | process_fn = self._process_compiled(model_params, len(input_tokens_padded), return_prompt_logp) |
| 123 | return process_fn( |
| 124 | model_params, input_tokens_padded, decode_slot, input_true_length, decode_state, rng, return_prompt_logp |
| 125 | ) |
| 126 | |
| 127 | def _process_compiled(self, params: Params, padded_length: int, return_prompt_logp: bool = False): |
| 128 | """Ahead-of-time compilation wrapper of _process().""" |
nothing calls this directly
no test coverage detected