refactor(Jetstream PT): drop HfEngine

Now that sampling can be passed as parameter to prefill and generate,
the custom engine is not required anymore.

text-generation-inference/server/text_generation_server/jetstream_pt_support/engine_loader.py

@@ -5,6 +5,7 @@
 
 import jax
 from jetstream_pt import fetch_models, torchjax
+from jetstream_pt.engine import PyTorchEngine
 from jetstream_pt.environment import (
     JetEngineEnvironment,
     JetEngineEnvironmentData,
@@ -13,7 +14,6 @@
 from loguru import logger
 from transformers import AutoConfig, PretrainedConfig
 
-from .engine import HfEngine
 from .llama_model_exportable_hf import TransformerHf
 
 
@@ -131,7 +131,7 @@ def create_engine(
     sequence_length: int,
     max_input_tokens: int,
     max_output_tokens: int,
-) -> HfEngine:
+) -> PyTorchEngine:
     # NOTE: for now no quantization is done
     env_data = create_engine_env_data(model_path, batch_size, sequence_length, max_input_tokens, max_output_tokens)
     if env_data is None:
@@ -142,7 +142,7 @@ def create_engine(
     weight_shardings = model.get_sharding_annotations()
     sharded_weights = shard_weights(env, model.state_dict(), weight_shardings)
 
-    return HfEngine(
+    return PyTorchEngine(
         pt_model=model,
         env=env,
         weights=torchjax.from_torch_with_copy(sharded_weights),

text-generation-inference/server/text_generation_server/jetstream_pt_support/generator.py

@@ -10,6 +10,7 @@
 import torch
 import torch_xla2
 from jetstream.engine.token_utils import pad_tokens, take_nearest_length
+from jetstream_pt.engine import PyTorchEngine
 from loguru import logger
 from transformers import AutoTokenizer, PreTrainedTokenizerBase
 from transformers.generation import GenerationConfig
@@ -27,7 +28,6 @@
     StoppingCriteriaParameters,
     Tokens,
 )
-from .engine import HfEngine
 from .engine_loader import create_engine
 from .token_selector import TokenSelector
 
@@ -243,7 +243,7 @@ class TpuGeneratorJetStream(Generator):
 
     def __init__(
         self,
-        engine: HfEngine,
+        engine: PyTorchEngine,
         tokenizer: PreTrainedTokenizerBase,
     ):
         self.engine = engine