Add an optional parameter for sampling in prefill / sample. (#133)

* Add an optional parameter for sampling in prefill / sample.

This is needed because we want to enable per-request
sampling parameters. This allows jetstream to be
used as backend for HuggingFace TGI.

* lint

Author: qihqi

jetstream/engine/engine_api.py

@@ -19,7 +19,7 @@
 """
 
 import abc
-from typing import Any, Optional, Tuple, Union
+from typing import Any, Optional, Tuple, Union, Callable
 
 from flax import struct
 import jax
@@ -142,18 +142,24 @@ def prefill(
       existing_prefix: Optional[Prefix] = None,
       padded_tokens: jax.Array,
       true_length: int,
+      sampler: Optional[Callable[[Any], Any]] = None,
   ) -> Tuple[Prefix, ResultTokens]:
     """Computes a kv-cache for a set of tokens conditional on existing cache.
 
     existing_prefix (if provided) represents a prefix that has already been
     processed by the underlying model. tokens is logically appended
     to the text represented by `existing_prefix`. This method returns a new
     kv_cache (typically) for the resulting text.
+
+    If sampler is passed, then the engine should use it do sample next token.
     """
 
   @abc.abstractmethod
   def generate(
-      self, params: Params, decode_state: DecodeState
+      self,
+      params: Params,
+      decode_state: DecodeState,
+      sampler: Optional[Callable[[Any], Any]] = None,
   ) -> Tuple[DecodeState, ResultTokens]:
     """Generates tokens for each sequence being decoded in parallel.
 
@@ -165,6 +171,8 @@ def generate(
     consists of each microbatch progressing through every stage), in
     non-pipelined code this is a full forward pass. In both cases, this accounts
     for a full embed-layerstack-unembed-sample operation.
+
+    If sampler is passed, then the engine should use it do sample next token.
     """
 
   @abc.abstractmethod