from pydantic import BaseModel

from distilabel.llms import TransformersLLM

# Example from outlines docs
class User(BaseModel):
    name: str
    last_name: str
    id: int

llm = TransformersLLM(
    model="argilla/notus-7b-v1",
    structured_output={"format": "json", "structure": User},
)

llm.load()
result = llm.generate(
    [
        [{"role": "user", "content": "Create a user profile with the fields name, last_name and id"}]
    ],
    max_new_tokens=30
)
print(result)
# [[{"name": "John", "last_name": "Doe", "id": 11}]]