outlines

1
"""Example of integrating `outlines` with `transformers`."""
2

3
from pydantic import BaseModel
4
from transformers import pipeline
5

6
from outlines.integrations.transformers import JSONPrefixAllowedTokens
7

8

9
class Person(BaseModel):
10
    first_name: str
11
    surname: str
12

13

14
pipe = pipeline("text-generation", model="mistralai/Mistral-7B-v0.1")
15
prefix_allowed_tokens_fn = JSONPrefixAllowedTokens(
16
    schema=Person, tokenizer_or_pipe=pipe, whitespace_pattern=r" ?"
17
)
18
results = pipe(
19
    ["He is Tom Jones", "She saw Linda Smith"],
20
    return_full_text=False,
21
    do_sample=False,
22
    max_new_tokens=50,
23
    prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
24
)
25
print(results)
26