Create a GeneratorStep from a dataset using a helper function

[plaguss]

Description

This PR simplifies the process to create a generator step from a dataset:

1. Helper function to create the GeneratorStep from the dataset already processed:

From the example by @dvsrepo, we can update the code like so:

dataset = load_dataset(
  "DIBT/10k_prompts_ranked",
  split="train"
).filter(
  lambda r: r['avg_rating']>=4 and r['num_responses']>=2
)

dataset = dataset.to_list()
load_dataset = LoadDataFromDicts(
  name="load_dataset",
  data=dataset[0:500], # during development I used [0:1] to iterate quickly
  output_mappings={"prompt": "instruction"}
)

can be integrated in a pipeline as:

from distilabel.pipeline import Pipeline
from distilabel.steps import make_generator_step

dataset = load_dataset(
  "DIBT/10k_prompts_ranked",
  split="train"
).filter(
  lambda r: r['avg_rating']>=4 and r['num_responses']>=2
)

with Pipeline() as pipeline:
    load_dataset = make_generator_step(
        dataset,
        output_mappings={"prompt": "instruction"}
    )
    dummy = DummyStep()
    load_dataset >> dummy

distiset = pipeline.run()

New entry in the docs:

2. Pass the dataset via pipeline.run(dataset=....)

Internally we will create the step (has less flexibility, but it's more direct, and easier if you don't need the flexibility):

from distilabel.pipeline import Pipeline

dataset = load_dataset(
  "DIBT/10k_prompts_ranked",
  split="train"
).filter(
  lambda r: r['avg_rating']>=4 and r['num_responses']>=2
)

with Pipeline() as pipeline:
    dummy = DummyStep()

distiset = pipeline.run(dataset=dataset)

Example in the docs with the new functionality:

A new example that will appear in the quick start with the new simplifications (the original pipeline can be compared here):

---

The pipeline from this blogpost you did @dvsrepo could now do with a less lines:

from distilabel.llms import InferenceEndpointsLLM
from distilabel.pipeline import Pipeline
from distilabel.steps import (
    GroupColumns,
    KeepColumns,
-    LoadDataFromDicts
)
from distilabel.steps.tasks import TextGeneration, UltraFeedback

from datasets import load_dataset

dataset = load_dataset(
    "DIBT/10k_prompts_ranked",
    split="train"
).filter(
    lambda r: r['avg_rating']>=4 and r['num_responses']>=2
+).select(range(100))
- dataset = dataset.to_list()

with Pipeline(
-    name="prefs-with-llama-3",
    description="Pipeline for building preference datasets using Llama 3",
) as pipeline:
-    load_dataset = LoadDataFromDicts(
-        name="load_dataset",
-        data=dataset[0:100],
-        output_mappings={"prompt": "instruction"}
-    )
    generate_with_llama3_70B = TextGeneration(
-        name="generate_with_llama3",
        llm=InferenceEndpointsLLM(
            model_id="meta-llama/Meta-Llama-3-70B-Instruct",
            tokenizer_id="meta-llama/Meta-Llama-3-8B-Instruct",
        ),
    )

    generate_with_llama3_8B = TextGeneration(
-        name="generate_with_llama3_8B",
        llm=InferenceEndpointsLLM(
            model_id="meta-llama/Meta-Llama-3-8B-Instruct",
            tokenizer_id="meta-llama/Meta-Llama-3-8B-Instruct",
        ),
    )

    combine_columns = GroupColumns(
-      name="combine_columns",
      columns=["generation", "model_name"],
      output_columns=["generations", "generation_models"],
    )

    ultrafeedback = UltraFeedback(
-      name="ultrafeedback",
      aspect='overall-rating',
      llm=InferenceEndpointsLLM(
            model_id="meta-llama/Meta-Llama-3-70B-Instruct",
            tokenizer_id="meta-llama/Meta-Llama-3-70B-Instruct",
        ),
    )
    keep_columns = KeepColumns(
-        name="keep_columns",
        columns=[
            "instruction",
            "generations",
            "generation_models",
            "ratings",
            "rationales",
        ],
    )

-    generate_with_llama3_70B.connect(combine_columns)
-    generate_with_llama3_8B.connect(combine_columns)
-    load_dataset.connect(generate_with_llama3_70B)
-    load_dataset.connect(generate_with_llama3_8B)
-    combine_columns.connect(ultrafeedback)
-    ultrafeedback.connect(keep_columns)
+    (
+        [generate_with_llama3_70B, generate_with_llama3_8B]
+        >> combine_columns
+        >> ultrafeedback
+        >> keep_columns
+    )


if __name__ == "__main__":
    distiset = pipeline.run(
+       dataset=dataset,
        parameters={
-            "load_dataset": {
-                "repo_id": "distilabel-internal-testing/instruction-dataset-mini",
-                "split": "test",
-            },
            "generate_with_llama3": {
                "llm": {
                    "generation_kwargs": {"max_new_tokens": 1024, "temperature": 0.7, "stop_sequences": ["<|eot_id|>", "<|end_of_text|>"]}
                }
            },
            "generate_with_llama3_8B": {
                "llm": {
                    "generation_kwargs": {"max_new_tokens": 1024, "temperature": 0.7, "stop_sequences": ["<|eot_id|>", "<|end_of_text|>"]}
                }
            },
            "ultrafeedback": {
                "llm": {
                    "generation_kwargs": {"max_new_tokens": 1024, "temperature": 0.1, "stop_sequences": ["<|eot_id|>", "<|end_of_text|>"]}
                }
            },
        }
    )
    distiset.push_to_hub(
        "dvilasuero/distillama3-prompts10k",
+        include_script=True
    )

[github-actions]

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

[codspeed-hq]

CodSpeed Performance Report

Merging #812 will not alter performance

_{Comparing pipeline-with-dataset (64e4ff2) with develop (25601bb)}

Summary

✅ 1 untouched benchmarks

[gabrielmbmb]

LGTM! Just some comments regarding the quickstart on the docs, docstrings and some suggestion for the code.

[gabrielmbmb]

I think that we could leave here the new example that you've added @plaguss as it's the quicker way to get started. WDYT @dvsrepo @davidberenstein1957 ?