Add tasks to replicate Math-shepherd

[plaguss]

Description

This task Integrates the tasks to replicate:
Math-Shepherd: Verify and Reinforce LLMs Step-by-step without Human Annotations

Examples:

• Dataset generated with ray: https://huggingface.co/datasets/plaguss/test_math_shepherd_prm_ray
• Generated with inference endpoints: https://huggingface.co/datasets/plaguss/test_math_shepherd_all_structured

It integrates structured outputs as I found while testing that it highly improves the success rate of the generations.

Example pipeline:

from datasets import load_dataset

from distilabel.steps.tasks.math_shepherd.generator import MathShepherdGenerator
from distilabel.steps.tasks.math_shepherd.completer import MathShepherdCompleter
from distilabel.steps.tasks.math_shepherd.utils import FormatPRM
from distilabel.models import InferenceEndpointsLLM
from distilabel.pipeline import Pipeline
from distilabel.steps import CombineOutputs, ExpandColumns

ds_name = "openai/gsm8k"

ds = load_dataset(ds_name, "main", split="test").rename_column("question", "instruction").select(range(3))


with Pipeline(name="Math-Shepherd") as pipe:
    model_id_70B = "meta-llama/Meta-Llama-3.1-70B-Instruct"
    model_id_8B = "meta-llama/Meta-Llama-3.1-8B-Instruct"

    llm_70B = InferenceEndpointsLLM(
        model_id=model_id_8B,
        tokenizer_id=model_id_8B,
        generation_kwargs={"max_new_tokens": 1024, "temperature": 0.6},
    )
    llm_8B = InferenceEndpointsLLM(
        model_id=model_id_8B,
        tokenizer_id=model_id_8B,
        generation_kwargs={"max_new_tokens": 2048, "temperature": 0.6},
    )

    generator_golden = MathShepherdGenerator(
        name="golden_generator",
        llm=llm_70B,
    )
    generator = MathShepherdGenerator(
        name="generator",
        llm=llm_8B,
        M=5  # Generate 5 sample solutions
    )
    completer = MathShepherdCompleter(
        name="completer",
        llm=llm_8B,
        N=4  # Each solution will be tested with 4 completions during labelling
    )

    combine = CombineOutputs()
    expand = ExpandColumns(
        name="expand_columns",
        columns=["solutions"],
        split_statistics=True,
    )
    formatter = FormatPRM(name="format_prm", format="trl")
    [generator_golden, generator] >> combine >> completer >> expand >> formatter


if __name__ == "__main__":
    distiset = pipe.run(use_cache=False, dataset=ds)
    distiset.push_to_hub("plaguss/test_math_shepherd_prm")

A sample dataset can be seen at plaguss/test_math_shepherd_prm

[github-actions]

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

[codspeed-hq]

CodSpeed Performance Report

Merging #1052 will not alter performance

_{Comparing math-shepherd (750f8d6) with develop (f8e41cd)}

Summary

✅ 1 untouched benchmarks