fix

octoml · Jan 9, 2024 · 2149d5d · 2149d5d
1 parent 58a1720
commit 2149d5d
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 7 deletions.
diff --git a/serve/benchmarks/benchmark_latency.py b/serve/benchmarks/benchmark_latency.py
@@ -1,4 +1,4 @@
-"""Benchmark offline user metric."""
+"""Benchmark latency offline."""
 import argparse
 import time, numpy as np
 from mlc_serve.engine import (

diff --git a/serve/mlc_serve/utils.py b/serve/mlc_serve/utils.py
@@ -56,24 +56,23 @@ def create_mlc_engine(args: argparse.Namespace):
             "max_decode_steps": args.max_decode_steps,
         }
     )
-    # type: off
+
     if args.use_staging_engine:
-        engine = StagingInferenceEngine(
+        engine = StagingInferenceEngine( # type: ignore
             tokenizer_module=HfTokenizerModule(args.model_artifact_path),
-            model_module_loader=PagedCacheModelModule, 
+            model_module_loader=PagedCacheModelModule, # type: ignore
             model_module_loader_kwargs={
                 "model_artifact_path": args.model_artifact_path,
                 "engine_config": engine_config,
             },
         )
         engine.start()
     else:
-        engine = SynchronousInferenceEngine( 
-            PagedCacheModelModule( 
+        engine = SynchronousInferenceEngine( # type: ignore
+            PagedCacheModelModule( # type: ignore
                 model_artifact_path=args.model_artifact_path,
                 engine_config=engine_config,
             )
         )
-    # type: on
     return engine