feat: move generation outside of pipeline

src/distilabel_dataset_generator/apps/sft.py

@@ -1,23 +1,24 @@
 import io
-import multiprocessing
-import time
 from typing import Union
 
 import gradio as gr
 import pandas as pd
 from datasets import Dataset
 from distilabel.distiset import Distiset
+from distilabel.steps.tasks.text_generation import TextGeneration
 from gradio.oauth import OAuthToken
 from huggingface_hub import upload_file
 
 from src.distilabel_dataset_generator.pipelines.sft import (
+    DEFAULT_BATCH_SIZE,
     DEFAULT_DATASET_DESCRIPTIONS,
     DEFAULT_DATASETS,
     DEFAULT_SYSTEM_PROMPTS,
     PROMPT_CREATION_PROMPT,
     generate_pipeline_code,
-    get_pipeline,
-    get_prompt_generation_step,
+    get_magpie_generator,
+    get_prompt_generator,
+    get_response_generator,
 )
 from src.distilabel_dataset_generator.utils import (
     get_login_button,
@@ -26,22 +27,15 @@ from src.distilabel_dataset_generator.utils import (
 )
 
 
-def _run_pipeline(result_queue, num_turns, num_rows, system_prompt, is_sample):
-    pipeline = get_pipeline(num_turns, num_rows, system_prompt, is_sample)
-    distiset: Distiset = pipeline.run(use_cache=False)
-    result_queue.put(distiset)
-
-
 def generate_system_prompt(dataset_description, progress=gr.Progress()):
+    progress(0.0, desc="Generating system prompt")
     if dataset_description in DEFAULT_DATASET_DESCRIPTIONS:
         index = DEFAULT_DATASET_DESCRIPTIONS.index(dataset_description)
         if index < len(DEFAULT_SYSTEM_PROMPTS):
             return DEFAULT_SYSTEM_PROMPTS[index]
 
-    progress(0.1, desc="Initializing text generation")
-    generate_description = get_prompt_generation_step()
-    progress(0.4, desc="Loading model")
-    generate_description.load()
+    progress(0.3, desc="Initializing text generation")
+    generate_description: TextGeneration = get_prompt_generator()
     progress(0.7, desc="Generating system prompt")
     result = next(
         generate_description.process(
@@ -62,12 +56,9 @@ def generate_sample_dataset(system_prompt, progress=gr.Progress()):
         index = DEFAULT_SYSTEM_PROMPTS.index(system_prompt)
         if index < len(DEFAULT_DATASETS):
             return DEFAULT_DATASETS[index]
-
-    progress(0.1, desc="Initializing sample dataset generation")
     result = generate_dataset(
         system_prompt, num_turns=1, num_rows=1, progress=progress, is_sample=True
     )
-    progress(1.0, desc="Sample dataset generated")
     return result
 
 
@@ -92,52 +83,98 @@ def generate_dataset(
     is_sample: bool = False,
     progress=gr.Progress(),
 ):
-    if num_rows < 5:
-        duration = 25
-    elif num_rows < 10:
-        duration = 60
-    elif num_rows < 30:
-        duration = 120
-    elif num_rows < 100:
-        duration = 240
-    elif num_rows < 300:
-        duration = 600
-    elif num_rows < 1000:
-        duration = 1200
-    else:
-        duration = 2400
-
-    result_queue = multiprocessing.Queue()
-    p = multiprocessing.Process(
-        target=_run_pipeline,
-        args=(result_queue, num_turns, num_rows, system_prompt, is_sample),
+    progress(0.0, desc="(1/2) Generating instructions")
+    magpie_generator = get_magpie_generator(
+        num_turns, num_rows, system_prompt, is_sample
     )
+    response_generator = get_response_generator(num_turns, system_prompt, is_sample)
+    total_steps: int = num_rows * 2
+    batch_size = DEFAULT_BATCH_SIZE
+
+    # create instructions
+    magpie_results = []
+    for i in range(0, num_rows, batch_size):
+        progress(
+            0.5 * min(i + batch_size, num_rows) / num_rows,
+            total=total_steps,
+            desc="(1/2) Generating instructions",
+        )
+        batch = list(magpie_generator.process())[:batch_size]
+        magpie_results.extend([item[0] for item in batch])
+    progress(0.5, desc="(1/2) Generating instructions")
 
-    try:
-        p.start()
-        total_steps = 100
-        for step in range(total_steps):
-            if not p.is_alive() or p._popen.poll() is not None:
-                break
+    # generate responses
+    response_results = []
+    if num_turns == 1:
+        for i in range(0, num_rows, batch_size):
             progress(
-                (step + 1) / total_steps,
-                desc=f"Generating dataset with {num_rows} rows. Don't close this window.",
+                0.5 + 0.5 * min(i + batch_size, num_rows) / num_rows,
+                total=total_steps,
+                desc="(2/2) Generating responses",
             )
-            time.sleep(duration / total_steps)  # Adjust this value based on your needs
-        p.join()
-    except Exception as e:
-        raise gr.Error(f"An error occurred during dataset generation: {str(e)}")
+            batch = magpie_results[i : i + batch_size]
+            batch = [entry[0] for entry in batch]
+            responses = list(response_generator.process(inputs=batch))
+            response_results.extend(responses)
+        for result in response_results[0]:
+            result["prompt"] = result["instruction"]
+            result["completion"] = result["generation"]
+            result["system_prompt"] = system_prompt
+    else:
+        for result in magpie_results:
+            result[0]["conversation"].insert(
+                0, {"role": "system", "content": system_prompt}
+            )
+            result[0]["messages"] = result[0]["conversation"]
+        for i in range(0, num_rows, batch_size):
+            progress(
+                0.5 + 0.5 * min(i + batch_size, num_rows) / num_rows,
+                total=total_steps,
+                desc="(2/2) Generating responses",
+            )
+            batch = magpie_results[i : i + batch_size]
+            batch = [entry[0] for entry in batch]
+            responses = list(response_generator.process(inputs=batch))
+            response_results.extend(responses)
+
+        for result in response_results[0]:
+            result["messages"].append(
+                {"role": "assistant", "content": result["generation"]}
+            )
+    progress(
+        1,
+        total=total_steps,
+        desc="(2/2) Generating responses",
+    )
 
-    distiset = result_queue.get()
+    # create distiset
+    distiset_results = []
+    for result in response_results[0]:
+        record = {}
+        for relevant_keys in [
+            "messages",
+            "prompt",
+            "completion",
+            "model_name",
+            "system_prompt",
+        ]:
+            if relevant_keys in result:
+                record[relevant_keys] = result[relevant_keys]
+        distiset_results.append(record)
+
+    distiset = Distiset(
+        {
+            "default": Dataset.from_list(distiset_results),
+        }
+    )
 
     # If not pushing to hub generate the dataset directly
-    distiset = distiset["default"]["train"]
+    distiset = distiset["default"]
     if num_turns == 1:
-        outputs = distiset.to_pandas()[["prompt", "completion"]]
+        outputs = distiset.to_pandas()[["system_prompt", "prompt", "completion"]]
     else:
         outputs = distiset.to_pandas()[["messages"]]
     dataframe = pd.DataFrame(outputs)
-
     progress(1.0, desc="Dataset generation completed")
     return dataframe
 
@@ -233,7 +270,7 @@ with gr.Blocks(
         )
 
         with gr.Row():
-            sample_dataset = gr.DataFrame(
+            sample_dataset = gr.Dataframe(
                 value=DEFAULT_DATASETS[0],
                 label="Sample dataset. Prompts and completions truncated to 256 tokens.",
                 interactive=False,
@@ -311,7 +348,7 @@ with gr.Blocks(
                     value="Push to Hub", variant="primary", scale=2
                 )
             with gr.Row():
-                final_dataset = gr.DataFrame(
+                final_dataset = gr.Dataframe(
                     value=DEFAULT_DATASETS[0],
                     label="Generated dataset",
                     interactive=False,

src/distilabel_dataset_generator/pipelines/sft.py

@@ -1,10 +1,12 @@
 import pandas as pd
+from datasets import Dataset
+from distilabel.distiset import Distiset
 from distilabel.llms import InferenceEndpointsLLM
 from distilabel.pipeline import Pipeline
 from distilabel.steps import KeepColumns
-from distilabel.steps.tasks import MagpieGenerator, TextGeneration
+from distilabel.steps.tasks import ChatGeneration, MagpieGenerator, TextGeneration
 
-from src.distilabel_dataset_generator.utils import HF_TOKENS
+from distilabel_dataset_generator.utils import HF_TOKENS
 
 INFORMATION_SEEKING_PROMPT = (
     "You are an AI assistant designed to provide accurate and concise information on a wide"
@@ -118,7 +120,7 @@ The prompt you write should follow the same style and structure as the following
 User dataset description:
 """
 
-MODEL = "meta-llama/Meta-Llama-3.1-70B-Instruct"
+MODEL = "meta-llama/Meta-Llama-3.1-8B-Instruct"
 DEFAULT_DATASET_DESCRIPTIONS = (
     "rude customer assistant for a phone company",
     "assistant that solves math puzzles using python",
@@ -155,7 +157,7 @@ _STOP_SEQUENCES = [
     "assistant",
     " \n\n",
 ]
-DEFAULT_BATCH_SIZE = 50
+DEFAULT_BATCH_SIZE = 5
 TOKEN_INDEX = 0
 
 
@@ -198,7 +200,7 @@ with Pipeline(name="sft") as pipeline:
         output_mappings={input_mappings},
     )
     keep_columns = KeepColumns(
-        columns={list(input_mappings.values())} + ["model_name"],
+        columns={list(input_mappings.values())} + ["model_name", "system_prompt"],
     )
     magpie.connect(keep_columns)
 
@@ -208,92 +210,101 @@ if __name__ == "__main__":
     return code
 
 
-def get_pipeline(num_turns, num_rows, system_prompt, is_sample):
+def _get_next_api_key():
     global TOKEN_INDEX
-    input_mappings = _get_output_mappings(num_turns)
-    output_mappings = input_mappings
     api_key = HF_TOKENS[TOKEN_INDEX % len(HF_TOKENS)]
     TOKEN_INDEX += 1
-    MODEL = "meta-llama/Meta-Llama-3.1-8B-Instruct"
-    print("is sample?", is_sample)
-    if num_turns == 1:
-        with Pipeline(name="sft") as pipeline:
-            magpie = MagpieGenerator(
-                llm=InferenceEndpointsLLM(
-                    model_id=MODEL,
-                    tokenizer_id=MODEL,
-                    api_key=api_key,
-                    magpie_pre_query_template="llama3",
-                    generation_kwargs={
-                        "temperature": 0.8,  # it's the best value for Llama 3.1 70B Instruct
-                        "do_sample": True,
-                        "max_new_tokens": 256 if is_sample else 512,
-                        "stop_sequences": _STOP_SEQUENCES,
-                    },
-                ),
-                batch_size=DEFAULT_BATCH_SIZE,
-                n_turns=num_turns,
-                num_rows=num_rows,
-                system_prompt=system_prompt,
-                output_mappings={"instruction": "prompt"},
-                only_instruction=True,
-            )
+    return api_key
 
-            generate_response = TextGeneration(
-                llm=InferenceEndpointsLLM(
-                    model_id=MODEL,
-                    tokenizer_id=MODEL,
-                    api_key=api_key,
-                    generation_kwargs={
-                        "temperature": 0.8,
-                        "max_new_tokens": 256 if is_sample else 1024,
-                    },
-                ),
-                system_prompt=system_prompt,
-                output_mappings={"generation": "completion"},
-                input_mappings={"instruction": "prompt"},
-            )
 
-            keep_columns = KeepColumns(
-                columns=list(output_mappings.values()) + ["model_name"],
-            )
+def get_magpie_generator(num_turns, num_rows, system_prompt, is_sample):
+    input_mappings = _get_output_mappings(num_turns)
+    output_mappings = input_mappings.copy()
+    if num_turns == 1:
+        magpie_generator = MagpieGenerator(
+            llm=InferenceEndpointsLLM(
+                model_id=MODEL,
+                tokenizer_id=MODEL,
+                api_key=_get_next_api_key(),
+                magpie_pre_query_template="llama3",
+                generation_kwargs={
+                    "temperature": 0.8,
+                    "do_sample": True,
+                    "max_new_tokens": 256 if is_sample else 512,
+                    "stop_sequences": _STOP_SEQUENCES,
+                },
+            ),
+            batch_size=DEFAULT_BATCH_SIZE,
+            n_turns=num_turns,
+            num_rows=num_rows,
+            system_prompt=system_prompt,
+            output_mappings=output_mappings,
+            only_instruction=True,
+        )
+    else:
+        magpie_generator = MagpieGenerator(
+            llm=InferenceEndpointsLLM(
+                model_id=MODEL,
+                tokenizer_id=MODEL,
+                api_key=_get_next_api_key(),
+                magpie_pre_query_template="llama3",
+                generation_kwargs={
+                    "temperature": 0.8,
+                    "do_sample": True,
+                    "max_new_tokens": 256 if is_sample else 1024,
+                    "stop_sequences": _STOP_SEQUENCES,
+                },
+            ),
+            batch_size=DEFAULT_BATCH_SIZE,
+            end_with_user=True,
+            n_turns=num_turns,
+            num_rows=num_rows,
+            system_prompt=system_prompt,
+            output_mappings=output_mappings,
+        )
+    magpie_generator.load()
+    return magpie_generator
 
-            magpie.connect(generate_response)
-            generate_response.connect(keep_columns)
-        return pipeline
+
+def get_response_generator(num_turns, system_prompt, is_sample):
+    if num_turns == 1:
+        response_generator = TextGeneration(
+            llm=InferenceEndpointsLLM(
+                model_id=MODEL,
+                tokenizer_id=MODEL,
+                api_key=_get_next_api_key(),
+                generation_kwargs={
+                    "temperature": 0.8,
+                    "max_new_tokens": 256 if is_sample else 1024,
+                },
+            ),
+            system_prompt=system_prompt,
+            output_mappings={"generation": "completion"},
+            input_mappings={"instruction": "prompt"},
+        )
     else:
-        with Pipeline(name="sft") as pipeline:
-            magpie = MagpieGenerator(
-                llm=InferenceEndpointsLLM(
-                    model_id=MODEL,
-                    tokenizer_id=MODEL,
-                    api_key=api_key,
-                    magpie_pre_query_template="llama3",
-                    generation_kwargs={
-                        "temperature": 0.8,  # it's the best value for Llama 3.1 70B Instruct
-                        "do_sample": True,
-                        "max_new_tokens": 2048,
-                        "stop_sequences": _STOP_SEQUENCES,
-                    },
-                ),
-                batch_size=DEFAULT_BATCH_SIZE,
-                n_turns=num_turns,
-                num_rows=num_rows,
-                system_prompt=system_prompt,
-                output_mappings=output_mappings,
-            )
-            keep_columns = KeepColumns(
-                columns=list(output_mappings.values()) + ["model_name"],
-            )
-            magpie.connect(keep_columns)
-        return pipeline
+        response_generator = ChatGeneration(
+            llm=InferenceEndpointsLLM(
+                model_id=MODEL,
+                tokenizer_id=MODEL,
+                api_key=_get_next_api_key(),
+                generation_kwargs={
+                    "temperature": 0.8,
+                    "max_new_tokens": 2048,
+                },
+            ),
+            output_mappings={"generation": "completion"},
+            input_mappings={"conversation": "messages"},
+        )
+    response_generator.load()
+    return response_generator
 
 
-def get_prompt_generation_step():
+def get_prompt_generator():
     global TOKEN_INDEX
     api_key = HF_TOKENS[TOKEN_INDEX % len(HF_TOKENS)]
     TOKEN_INDEX += 1
-    generate_description = TextGeneration(
+    prompt_generator = TextGeneration(
         llm=InferenceEndpointsLLM(
             api_key=api_key,
             model_id=MODEL,
@@ -306,13 +317,30 @@ def get_prompt_generation_step():
         ),
         use_system_prompt=True,
     )
-    return generate_description
+    prompt_generator.load()
+    return prompt_generator
+
+
+def get_pipeline(num_turns, num_rows, system_prompt, is_sample):
+    input_mappings = _get_output_mappings(num_turns)
+    output_mappings = input_mappings
+
+    with Pipeline(name="sft") as pipeline:
+        magpie = get_magpie_generator(num_turns, num_rows, system_prompt, is_sample)
+        generate_response = get_response_generator(system_prompt, is_sample)
+
+        keep_columns = KeepColumns(
+            columns=list(output_mappings.values()) + ["model_name"],
+        )
+
+        magpie.connect(generate_response)
+        generate_response.connect(keep_columns)
+        return pipeline
 
 
 if __name__ == "__main__":
-    prompt_generation_step = get_prompt_generation_step()
-    prompt_generation_step.load()
-    result = next(
+    prompt_generation_step = get_prompt_generator()
+    system_prompt = next(
         prompt_generation_step.process(
             [
                 {
@@ -322,5 +350,64 @@ if __name__ == "__main__":
             ]
         )
     )[0]["generation"]
-    pipeline = get_pipeline(num_rows=100, num_turns=1, system_prompt=result)
-    pipeline.run()
+    num_rows = 2
+    num_turns = 1
+    magpie_generator = get_magpie_generator(num_turns, num_rows, system_prompt, False)
+    response_generator = get_response_generator(num_turns, system_prompt, False)
+    total_steps = num_rows * 2
+    batch_size = 5  # Adjust this value as needed
+
+    # create instructions
+    magpie_results = []
+    for i in range(0, num_rows, batch_size):
+        batch = list(magpie_generator.process())[:batch_size]
+        magpie_results.extend([item[0] for item in batch])
+
+    # generate responses
+    response_results = []
+    if num_turns == 1:
+        for i in range(0, len(magpie_results), batch_size):
+            batch = magpie_results[i : i + batch_size]
+            batch = [entry[0] for entry in batch]
+            responses = list(response_generator.process(inputs=batch))
+            response_results.extend(responses)
+        for result in response_results:
+            result[0]["prompt"] = result[0]["instruction"]
+            result[0]["completion"] = result[0]["generation"]
+            result[0]["system_prompt"] = system_prompt
+    else:
+        for result in magpie_results:
+            result[0]["conversation"].insert(
+                0, {"role": "system", "content": system_prompt}
+            )
+            result[0]["messages"] = result[0]["conversation"]
+        for i in range(0, len(magpie_results), batch_size):
+            batch = magpie_results[i : i + batch_size]
+            batch = [entry[0] for entry in batch]
+            responses = list(response_generator.process(inputs=batch))
+            response_results.extend(responses)
+
+        for result in response_results:
+            result[0]["messages"].append(
+                {"role": "assistant", "content": result[0]["generation"]}
+            )
+
+    distiset_results = []
+    for result in response_results[0]:
+        record = {}
+        for relevant_keys in [
+            "messages",
+            "prompt",
+            "completion",
+            "model_name",
+            "system_prompt",
+        ]:
+            if relevant_keys in result:
+                record[relevant_keys] = result[relevant_keys]
+        distiset_results.append(record)
+
+    distiset = Distiset(
+        {
+            "default": Dataset.from_list(distiset_results),
+        }
+    )