Add task decorator

[gabrielmbmb]

Description

This PR adds a new task decorator that allows to easily create new Tasks defining an output formatting function and the system prompt and user message template using the docstring of the function.

from typing import Dict, Any, Union
from distilabel.steps.tasks import task
from distilabel.llms import InferenceEndpointsLLM


@task(inputs=["task", "instruction"], outputs=["response"])
def MyTask(
    output: Union[str, None], input: Union[Dict[str, Any], None] = None
) -> Dict[str, Any]:
    """
    `MyTask` is a simple `Task` for bla bla bla

    ---
    system_prompt: |
        You are an AI assistant designed to {task}

    user_message_template: |
        {instruction}
    ---
    """
    return {"response": output}


task = MyTask(llm=InferenceEndpointsLLM(model_id="Qwen/Qwen2.5-72B-Instruct"))

task.load()

result = next(
    task.process_applying_mappings(
        [
            {
                "task": "extract information",
                "instruction": "Who won the nobel? The Nobel Committee lauded David Baker, a US biochemist, for completing “the almost impossible feat of building entirely new kinds of proteins,” and Demis Hassabis and John Jumper, who work at Google DeepMind, for developing an AI model to predict proteins’ complex structures – a problem that had been unsolved for 50 years.",
            }
        ]
    )
)

print(result)

[github-actions]

Documentation for this PR has been built. You can view it at: https://distilabel.argilla.io/pr-1028/

[codspeed-hq]

CodSpeed Performance Report

Merging #1028 will not alter performance

_{Comparing task-decorator (e51f8b9) with develop (106e402)}

Summary

✅ 1 untouched benchmarks