Fix outlines compatibility with speculative decoding (#578)

predibase · Aug 16, 2024 · ca2e643 · ca2e643
1 parent 12dc740
commit ca2e643
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 9 deletions.
diff --git a/server/lorax_server/models/flash_causal_lm.py b/server/lorax_server/models/flash_causal_lm.py
@@ -268,7 +268,9 @@ def from_pb(
 
         adapter_indices = torch.cat(adapter_indices_list).to(dtype=torch.int64, device=device)
 
-        request_tokenizers = [tokenizers.get_tokenizer(r.adapter_index, tokenizer) for r in pb.requests]
+        # always use the base model tokenizer for the next token chooser until we revisit adding back support
+        # for per-request tokenizers
+        request_tokenizers = [tokenizer for _ in pb.requests]
         next_token_chooser = HeterogeneousNextTokenChooser.from_pb(
             next_token_chooser_parameters, request_tokenizers, dtype, device
         )

diff --git a/server/poetry.lock b/server/poetry.lock
diff --git a/server/pyproject.toml b/server/pyproject.toml
@@ -38,7 +38,7 @@ boto3 = "^1.28.34"
 urllib3 = "<=1.26.18"
 hqq = { version = "^0.1.7", optional = true }
 stanford-stk = { version = "^0.7.0", markers = "sys_platform == 'linux'" }
-outlines = { version = "^0.0.40", optional = true }
+outlines = { version = "^0.0.46", optional = true }
 prometheus-client = "^0.20.0"
 py-cpuinfo = "^9.0.0"
 nvidia-ml-py = "^12.555.43"