Update vllm_integration.py to the latest version of outlines (#1352)

dottxt-ai · Jan 15, 2025 · 7b9012b · 7b9012b
1 parent a03c254
commit 7b9012b
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/examples/vllm_integration.py b/examples/vllm_integration.py
@@ -2,17 +2,23 @@
 
 import vllm
 from pydantic import BaseModel
+from transformers import AutoTokenizer
 
-from outlines.processors.structured import JSONLogitsProcessor
+from outlines.models.vllm import adapt_tokenizer
+from outlines.processors import JSONLogitsProcessor
 
 
 class Person(BaseModel):
     first_name: str
     surname: str
 
 
-llm = vllm.LLM(model="mistralai/Mistral-7B-v0.1", max_model_len=512)
-logits_processor = JSONLogitsProcessor(schema=Person, llm=llm, whitespace_pattern=r" ?")
+MODEL_ID = "mistralai/Mistral-7B-v0.1"
+llm = vllm.LLM(model=MODEL_ID, max_model_len=512)
+tokenizer = adapt_tokenizer(AutoTokenizer.from_pretrained(MODEL_ID))
+logits_processor = JSONLogitsProcessor(
+    schema=Person, tokenizer=tokenizer, whitespace_pattern=r" ?"
+)
 result = llm.generate(
     ["He is Tom Jones", "She saw Linda Smith"],
     sampling_params=vllm.SamplingParams(