Move create_pytorch_engine to init. (#7)

Makes the UX a little bit cleaner

Author: qihqi

jetstream_pt/__init__.py

@@ -0,0 +1 @@
+from jetstream_pt.engine import create_pytorch_engine

run_server.py

@@ -6,8 +6,9 @@
 from absl import flags
 
 from jetstream.core import server_lib
+import jetstream_pt
 from jetstream_pt import config
-from jetstream_pt import engine as je
+from jetstream.core.config_lib import ServerConfig
 
 
 _PORT = flags.DEFINE_integer('port', 9000, 'port to listen on')
@@ -61,15 +62,14 @@
 _QUANTIZE_KV_CACHE = flags.DEFINE_bool('quantize_kv_cache', False, 'kv_cache_quantize')
 _MAX_CACHE_LENGTH = flags.DEFINE_integer('max_cache_length', 1024, 'kv_cache_quantize')
 
-from jetstream.core.config_lib import ServerConfig
 
 def main(argv: Sequence[str]):
   del argv
   os.environ['XLA_FLAGS'] = '--xla_dump_to=/tmp/xla_logs --xla_dump_hlo_as_text'
   # No devices for local cpu test. A None for prefill and a None for generate.
   devices = server_lib.get_devices()
   print(f"devices: {devices}")
-  engine = je.create_pytorch_engine(
+  engine = jetstream_pt.create_pytorch_engine(
         devices=devices,
         tokenizer_path=_TOKENIZER_PATH.value,
         ckpt_path=_CKPT_PATH.value,