Add StructuredGeneration task and support for grammar in InferenceEndpointsLLM

[alvarobartt]

Description

This PR adds the StructuredGeneration task, similarly to the TextGeneration one, but also expecting the input grammar and producing both the chat-like input and the grammar within the format_input. In order to achieve that, the typing has been updated / modified in Task.process and also LLM.generate so that the received input/s contain either only the chat for most of the cases, or the chat and the grammar for the StructuredGeneration case.

Note

This is still a work in progress and subject to changes, but at the moment this seems the most straight forward / intuitive way to do so.

Additionally, this PR adds the grammar arg within InferenceEndpointsLLM so that it can be provided via the init or via runtime parameter.

Note

The main difference between using the grammar arg compared to using the StructuredGeneration task relies on the fact that the grammar arg is intended to be used with any task, whenever we want the output to match a certain format e.g. in UltraFeedback we may want the output to match a certain regex to avoid output parsing issues; while on the other hand, the task StructuredGeneration is intended when we have different grammars per row and we want to generate an output for the given instruction based on a different grammar per row, e.g. a function calling scenario where we want each generation for a given instruction to match a certain function schema.

Examples

grammar at LLM-level (same grammar for every generation)

from distilabel.llms import InferenceEndpointsLLM
from distilabel.pipeline import Pipeline
from distilabel.steps import LoadDataFromDicts
from distilabel.steps.tasks import TextGeneration
from pydantic import BaseModel


class Character(BaseModel):
    name: str
    description: str
    role: str
    weapon: str


with Pipeline(name="inference-endpoints-structured-generation") as pipeline:
    load_data = LoadDataFromDicts(
        name="load_data",
        data=[{"instruction": "Generate a character from a RPG game."}],
    )

    text_generation_cohere = TextGeneration(
        name="text_generation_cohere",
        llm=InferenceEndpointsLLM(
            model_id="CohereForAI/c4ai-command-r-plus",
            tokenizer_id="CohereForAI/c4ai-command-r-plus",
            api_key="***",  # type: ignore
            grammar={
                "type": "json",
                "value": Character.model_json_schema(),
            },
        ),
        use_system_prompt=False,
        input_batch_size=10,
        output_mappings={"model_name": "generation_model"},
    )

    load_data >> text_generation_cohere  # type: ignore


if __name__ == "__main__":
    distiset = pipeline.run(
        parameters={  # type: ignore
            text_generation_cohere.name: {
                "llm": {
                    "generation_kwargs": {
                        "temperature": 0.7,
                        "max_new_tokens": 4096,
                        "stop_sequences": ["<EOS_TOKEN>", "<|END_OF_TURN_TOKEN|>"],
                    }
                }
            },
        },
    )
    if distiset is not None:
        distiset.push_to_hub(
            "distilabel-internal-testing/inference-endpoints-structured-generation",
            token="***",
        )

grammar via StructuredGeneration (one grammar` per row)

from distilabel.llms import InferenceEndpointsLLM
from distilabel.pipeline import Pipeline
from distilabel.steps import LoadDataFromDicts
from distilabel.steps.tasks.structured_generation import StructuredGeneration
from pydantic import BaseModel


class Character(BaseModel):
    name: str
    description: str
    role: str
    weapon: str


class Animal(BaseModel):
    name: str
    species: str
    habitat: str
    diet: str


with Pipeline(name="inference-endpoints-structured-generation") as pipeline:
    load_data = LoadDataFromDicts(
        name="load_data",
        data=[
            {
                "instruction": "Generate a character from a RPG game.",
                "grammar": {
                    "type": "json",
                    "value": Character.model_json_schema(),
                },
            },
            {
                "instruction": "Generate an animal from a zoo.",
                "grammar": {
                    "type": "json",
                    "value": Animal.model_json_schema(),
                },
            },
            {
                "instruction": "What's the weather like today in Seattle in Celsius degrees?",
                "grammar": {
                    "type": "regex",
                    "value": "(\\d{1,2})°C",
                },
            },
        ],
    )

    task = StructuredGeneration(
        name="task",
        llm=InferenceEndpointsLLM(
            model_id="CohereForAI/c4ai-command-r-plus",
            tokenizer_id="CohereForAI/c4ai-command-r-plus",
            api_key="***",  # type: ignore
        ),
        use_system_prompt=False,
        output_mappings={"model_name": "generation_model"},
    )

    load_data >> task  # type: ignore


if __name__ == "__main__":
    distiset = pipeline.run(
        parameters={  # type: ignore
            task.name: {
                "llm": {
                    "generation_kwargs": {
                        "temperature": 0.7,
                        "max_new_tokens": 4096,
                        "stop_sequences": ["<EOS_TOKEN>", "<|END_OF_TURN_TOKEN|>"],
                    }
                }
            },
        },
    )
    if distiset is not None:
        distiset.push_to_hub(
            "distilabel-internal-testing/inference-endpoints-structured-generation-multiple",
            token="***",
        )

[codspeed-hq]

CodSpeed Performance Report

Merging #680 will not alter performance

_{Comparing inference-endpoints-structured-gen (e7399d1) with develop (1624b1e)}

Summary

✅ 1 untouched benchmarks