From d2affb87780e03c7d40ee7ff969be2bee373229c Mon Sep 17 00:00:00 2001 From: saattrupdan Date: Mon, 4 Mar 2024 17:19:36 +0100 Subject: [PATCH] feat: Add whitespace_pattern to examples --- examples/transformers_integration.py | 2 +- examples/vllm_integration.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/transformers_integration.py b/examples/transformers_integration.py index b4019a479..05acbfa14 100644 --- a/examples/transformers_integration.py +++ b/examples/transformers_integration.py @@ -13,7 +13,7 @@ class Person(BaseModel): pipe = pipeline("text-generation", model="mistralai/Mistral-7B-v0.1") prefix_allowed_tokens_fn = JSONPrefixAllowedTokens( - schema=Person, tokenizer_or_pipe=pipe + schema=Person, tokenizer_or_pipe=pipe, whitespace_pattern=" ?" ) results = pipe( ["He is Tom Jones", "She saw Linda Smith"], diff --git a/examples/vllm_integration.py b/examples/vllm_integration.py index dc5579465..5a9c04f01 100644 --- a/examples/vllm_integration.py +++ b/examples/vllm_integration.py @@ -12,7 +12,7 @@ class Person(BaseModel): llm = vllm.LLM(model="mistralai/Mistral-7B-v0.1", max_model_len=512) -logits_processor = JSONLogitsProcessor(schema=Person, llm=llm) +logits_processor = JSONLogitsProcessor(schema=Person, llm=llm, whitespace_pattern=" ?") result = llm.generate( ["He is Tom Jones", "She saw Linda Smith"], sampling_params=vllm.SamplingParams(