[FEATURE] Add structured outputs using outlines

[plaguss]

Description

This PR adds an integration with outlines to generate structured outputs for LLMs.

LLMs available:

The following LLMs are implemented, just pass the StructuredOutputType:

• Transformers
• LlamaCpp
• vLLM

Example pipeline LlamaCPP:

from enum import Enum
from pathlib import Path

from distilabel.llms import vLLM, LlamaCppLLM
from distilabel.pipeline import Pipeline
from distilabel.steps import LoadDataFromDicts
from distilabel.steps.tasks import TextGeneration
from pydantic import BaseModel, StringConstraints, conint
from typing_extensions import Annotated


class Weapon(str, Enum):
    sword = "sword"
    axe = "axe"
    mace = "mace"
    spear = "spear"
    bow = "bow"
    crossbow = "crossbow"


class Armor(str, Enum):
    leather = "leather"
    chainmail = "chainmail"
    plate = "plate"
    mithril = "mithril"


class Character(BaseModel):
    name: Annotated[str, StringConstraints(max_length=30)]
    age: conint(gt=1, lt=3000)
    armor: Armor
    weapon: Weapon


# Download the model with
# curl -L -o ~/Downloads/openhermes-2.5-mistral-7b.Q4_K_M.gguf https://huggingface.co/TheBloke/OpenHermes-2.5-Mistral-7B-GGUF/resolve/main/openhermes-2.5-mistral-7b.Q4_K_M.gguf

model_path = "Downloads/openhermes-2.5-mistral-7b.Q4_K_M.gguf"

with Pipeline("RPG-characters") as pipeline:
    system_prompt = (
        "You are a leading role play gamer. You have seen thousands of different characters and their attributes."
        " Please return a JSON object with common attributes of an RPG character."
    )

    load_dataset = LoadDataFromDicts(
        name="load_instructions",
        data=[
            {
                "system_prompt": system_prompt,
                "instruction": f"Give me a character description for a {char}",
            }
            for char in ["dwarf", "elf", "human", "ork"]
        ],
    )
    llm=LlamaCppLLM(
        model_path=str(Path.home() / model_path),  # type: ignore
        n_gpu_layers=-1,
        n_ctx=1024,
        structured_output={"format": "json", "schema": Character},
    )
    # Change to vLLM as such:
    # llm = vLLM(
    #     model="teknium/OpenHermes-2.5-Mistral-7B",
    #     extra_kwargs={"tensor_parallel_size": 1},
    #     structured_output={"format": "json", "schema": Character},
    # )

    text_generation = TextGeneration(
        name="text_generation_rpg",
        llm=llm,
        input_batch_size=8,
        output_mappings={"model_name": "generation_model"},
    )
    load_dataset >> text_generation


if __name__ == "__main__":
    distiset = pipeline.run(
        parameters={
            text_generation.name: {
                "llm": {"generation_kwargs": {"max_new_tokens": 256}}
            }
        },
        use_cache=False,
    )
    print(distiset)
    df = distiset["default"]["train"].to_pandas()
    df.to_csv("rpg_characters.csv", index=False)

    for num, character in enumerate(distiset["default"]["train"]["generation"]):
        print(f"Character: {num}")
        print(character)

# Character: 0
# {
# "name": "Gimli",
# "age": 42,
# "armor": "plate",
# "weapon": "axe" }
# Character: 1
# {"name":"Gaelen","age":600,"armor":"leather","weapon":"bow"}
# Character: 2
# {"name": "John Smith","age": 35,"armor": "leather","weapon": "sword"}
# Character: 3
# { "name": "Grug", "age": 35, "armor": "leather", "weapon": "axe"}

• OpenAI
  Outlines offers integrations, but given that OpenAI already offers JSON support, and the other types of output are not available, this one should be done directly from the corresponding LLM:

OpenAILLM(..., response_format="json")  # To generate JSON: https://platform.openai.com/docs/api-reference/chat/create#chat-create-response_format

More information from OpenAI for function calling can be found here.

[alvarobartt]

If we're planning to use a real LLM shouldn't this be an integration test instead?

[plaguss]

Sure, I started with the simplest/fastest possible to iterate, but I wanted to discuss once the other pending details are done 😄

[alvarobartt]

Just chiming in as mentioned this morning that we were using real LLMs in some cases, which shouldn't be the case for unit tests if we can avoid that (talking about the llama-cpp case now)

[gabrielmbmb]

LGTM! This feat is awesome 🙌 just some small comments

[gabrielmbmb]

Can we be more specific on the type hinting of this attribute? or we can't because circular imports, etc?

[plaguss]

I couldn't due to circular imports, and/or validation from pydantic, would love to improve these if we can find a way