Skip to content

Commit

Permalink
Fix llamacpp caching by making LlamaCppTokenizer pickleable
Browse files Browse the repository at this point in the history
  • Loading branch information
lapp0 committed May 31, 2024
1 parent 7723ce8 commit 62e1ef6
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 0 deletions.
10 changes: 10 additions & 0 deletions outlines/integrations/llamacpp.py
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,16 @@ def __init__(self, model: "Llama"):
def convert_token_to_string(self, token: str) -> str:
return token

def __getstate__(self):
"""Allow tokenizer to be used as hash key by excluding self.decode"""
return (
self.vocabulary.items(),
self.eos_token_id,
self.eos_token,
self.pad_token_id,
sorted(self.special_tokens),
)


class LogitsProcessor:
"""Bias LlamaCpp generation using a finite state machine.
Expand Down
8 changes: 8 additions & 0 deletions tests/generate/test_integration_llamacpp.py
Original file line number Diff line number Diff line change
Expand Up @@ -279,3 +279,11 @@ def test_llama_cpp_pre_tokenizer_remains_broken():
model = models.llamacpp(repo, model_path)
with pytest.raises(RuntimeError):
generate.choice(model, ["skirt", "dress", "pen", "jacket"])


def test_create_states_mapping_llamacpp_tokenizer_regression(model):
"""Minimal reproducer for #922, error passing llamacpp tokenizer to create_states_mapping"""
from outlines.fsm.guide import create_states_mapping
from outlines.integrations.llamacpp import LlamaCppTokenizer

create_states_mapping("a", LlamaCppTokenizer(model.model))

0 comments on commit 62e1ef6

Please sign in to comment.