Add DeitaFiltering step

[gabrielmbmb]

Description

This PR adds the DeitaFiltering steps and introduces some parameters in EvolQuality and EvolInstruct that allows including the original instruction in the list of evolved_{response,instruction}. In addition, it adds the ExpandColumns and GenerateConversation steps auxiliary steps that were needed to build the whole DEITA pipeline.

Pipeline

from distilabel.llm.huggingface.transformers import TransformersLLM
from distilabel.llm.openai import OpenAILLM
from distilabel.pipeline.local import Pipeline
from distilabel.steps.conversation import ConversationTemplate
from distilabel.steps.deita import DeitaFiltering
from distilabel.steps.expand import ExpandColumns
from distilabel.steps.generators.huggingface import LoadHubDataset
from distilabel.steps.task.complexity_scorer import ComplexityScorer
from distilabel.steps.task.evol_instruct.base import EvolInstruct
from distilabel.steps.task.evol_quality.base import EvolQuality
from distilabel.steps.task.generate_embeddings import GenerateEmbeddings
from distilabel.steps.task.quality_scorer import QualityScorer
from distilabel.steps.globals.huggingface import PushToHub


with Pipeline(name="DEITA", description="") as pipeline:
    load_data = LoadHubDataset(
        name="load_data", batch_size=100, output_mappings={"prompt": "instruction"}
    )

    evol_instruction_complexity = EvolInstruct(
        name="evol_instruction_complexity",
        llm=OpenAILLM(model="gpt-3.5-turbo"),
        num_evolutions=5,
        store_evolutions=True,
        generate_answers=True,
        include_original_instruction=True,
    )

    instruction_complexity_scorer = ComplexityScorer(
        name="instruction_complexity_scorer",
        llm=OpenAILLM(model="gpt-3.5-turbo"),
        input_mappings={"instructions": "evolved_instructions"},
    )

    expand_evolved_instructions = ExpandColumns(
        name="expand_evolved_instructions",
        columns=["evolved_instructions", "answers", "scores"],
        output_mappings={
            "evolved_instructions": "evolved_instruction",
            "answers": "answer",
            "scores": "evol_instruction_score",
        },
    )

    evol_response_quality = EvolQuality(
        name="evol_response_quality",
        llm=OpenAILLM(model="gpt-3.5-turbo"),
        num_evolutions=5,
        store_evolutions=True,
        include_original_response=True,
        input_mappings={
            "instruction": "evolved_instruction",
            "response": "answer",
        },
    )

    response_quality_scorer = QualityScorer(
        name="response_quality_scorer",
        llm=OpenAILLM(model="gpt-3.5-turbo"),
        input_mappings={
            "instruction": "evolved_instruction",
            "responses": "evolved_responses",
        },
    )

    expand_evolved_responses = ExpandColumns(
        name="expand_evolved_responses",
        columns=["evolved_responses", "scores"],
        output_mappings={
            "evolved_responses": "evolved_response",
            "scores": "evol_response_score",
        },
    )

    generate_conversation = ConversationTemplate(
        name="generate_conversation",
        input_mappings={
            "instruction": "evolved_instruction",
            "response": "evolved_response",
        },
    )

    push_to_hub_after_conversation = PushToHub(
        name="push_to_hub_after_conversation",
    )

    generate_embeddings = GenerateEmbeddings(
        name="generate_embeddings",
        llm=TransformersLLM(
            model="argilla/notus-7b-v1",
            device="cuda",
            torch_dtype="float16",
        ),
        input_mappings={"text": "conversation"},
        input_batch_size=5,
    )

    push_to_hub_after_embeddings = PushToHub(
        name="push_to_hub_after_embeddings",
    )

    deita_filtering = DeitaFiltering(name="deita_filtering")

    load_data.connect(evol_instruction_complexity)
    evol_instruction_complexity.connect(instruction_complexity_scorer)
    instruction_complexity_scorer.connect(expand_evolved_instructions)
    expand_evolved_instructions.connect(evol_response_quality)
    evol_response_quality.connect(response_quality_scorer)
    response_quality_scorer.connect(expand_evolved_responses)
    expand_evolved_responses.connect(generate_conversation)
    generate_conversation.connect(generate_embeddings)
    generate_conversation.connect(push_to_hub_after_conversation)
    generate_embeddings.connect(deita_filtering)
    generate_embeddings.connect(push_to_hub_after_embeddings)

if __name__ == "__main__":
    distiset = pipeline.run(
        parameters={
            "load_data": {
                "repo_id": "HuggingFaceH4/instruction-dataset",
                "split": "test",
            },
            "evol_instruction_complexity": {
                "llm": {
                    "generation_kwargs": {"max_new_tokens": 512, "temperature": 0.7}
                }
            },
            "instruction_complexity_scorer": {
                "llm": {"generation_kwargs": {"temperature": 0.0}}
            },
            "evol_response_quality": {
                "llm": {
                    "generation_kwargs": {"max_new_tokens": 512, "temperature": 0.7}
                }
            },
            "response_quality_scorer": {
                "llm": {"generation_kwargs": {"temperature": 0.0}}
            },
            "push_to_hub_after_conversation": {
                "repo_id": "distilabel-internal-testing/deita-after-conversation",
            },
            "push_to_hub_after_embeddings": {
                "repo_id": "distilabel-internal-testing/deita-after-embeddings",
            },
            "deita_filtering": {"data_budget": 3000, "diversity_threshold": 0.04},
        },
        use_cache=True,
    )
    distiset.push_to_hub("distilabel-internal-testing/deita")