feat: Add whitespace_pattern to examples

dottxt-ai · Mar 4, 2024 · d2affb8 · d2affb8
1 parent f187085
commit d2affb8
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git 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
@@ -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(