Merge branch 'argilla-io:main' into main

README.md

@@ -86,12 +86,14 @@ You can set the following environment variables to customize the generation proc
 Optionally, you can use different API providers and models.
 
 - `MODEL`: The model to use for generating the dataset, e.g. `meta-llama/Meta-Llama-3.1-8B-Instruct`, `gpt-4o`, `llama3.1`.
-- `API_KEY`: The API key to use for the generation API, e.g. `hf_...`, `sk-...`. If not provided, it will default to the provided `HF_TOKEN` environment variable.
+- `API_KEY`: The API key to use for the generation API, e.g. `hf_...`, `sk-...`. If not provided, it will default to the `HF_TOKEN` environment variable.
 - `OPENAI_BASE_URL`: The base URL for any OpenAI compatible API, e.g. `https://api.openai.com/v1/`.
 - `OLLAMA_BASE_URL`: The base URL for any Ollama compatible API, e.g. `http://127.0.0.1:11434/`.
 - `HUGGINGFACE_BASE_URL`: The base URL for any Hugging Face compatible API, e.g. TGI server or Dedicated Inference Endpoints. If you want to use serverless inference, only set the `MODEL`.
 - `VLLM_BASE_URL`: The base URL for any VLLM compatible API, e.g. `http://localhost:8000/`.
 
+To use a specific model exclusively for generating completions, set the corresponding environment variables by appending `_COMPLETION` to the ones mentioned earlier. For example, you can use `MODEL_COMPLETION` and `OPENAI_BASE_URL_COMPLETION`.
+
 SFT and Chat Data generation is not supported with OpenAI Endpoints. Additionally, you need to configure it per model family based on their prompt templates using the right `TOKENIZER_ID` and `MAGPIE_PRE_QUERY_TEMPLATE` environment variables.
 
 - `TOKENIZER_ID`: The tokenizer ID to use for the magpie pipeline, e.g. `meta-llama/Meta-Llama-3.1-8B-Instruct`.

examples/hf-serverless-deployment-deepseek.py

@@ -0,0 +1,16 @@
+# /// script
+# requires-python = ">=3.11,<3.12"
+# dependencies = [
+#     "synthetic-dataset-generator",
+# ]
+# ///
+import os
+
+from synthetic_dataset_generator import launch
+
+os.environ["HF_TOKEN"] = "hf_..."  # push the data to huggingface
+os.environ["MODEL"] = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"  # use model for instructions
+os.environ["MAGPIE_PRE_QUERY_TEMPLATE"] = "<｜begin▁of▁sentence｜>User: "  # use the custom template for the model
+
+
+launch()

examples/hf-serverless-deployment.py

@@ -9,7 +9,7 @@ import os
 from synthetic_dataset_generator import launch
 
 os.environ["HF_TOKEN"] = "hf_..."  # push the data to huggingface
-os.environ["MODEL"] = "meta-llama/Llama-3.1-8B-Instruct"  # use instruct model
+os.environ["MODEL"] = "meta-llama/Llama-3.1-8B-Instruct"  # use model for generation
 os.environ["MAGPIE_PRE_QUERY_TEMPLATE"] = "llama3"  # use the template for the model
 
 launch()

examples/hf-serverless-different-model-for-completion.py

@@ -0,0 +1,16 @@
+# /// script
+# requires-python = ">=3.11,<3.12"
+# dependencies = [
+#     "synthetic-dataset-generator",
+# ]
+# ///
+import os
+
+from synthetic_dataset_generator import launch
+
+os.environ["HF_TOKEN"] = "hf_..."  # push the data to huggingface
+os.environ["MODEL"] = "meta-llama/Llama-3.1-8B-Instruct"  # use model for instruction generation
+os.environ["MODEL_COMPLETION"] = "meta-llama/Llama-3.1-70B-Instruct"  # use model for completion generation
+os.environ["MAGPIE_PRE_QUERY_TEMPLATE"] = "llama3"  # use the template for the model
+
+launch()

examples/ollama-different-model-for-completion.py

@@ -0,0 +1,26 @@
+# /// script
+# requires-python = ">=3.11,<3.12"
+# dependencies = [
+#     "synthetic-dataset-generator",
+# ]
+# ///
+# ollama serve
+# ollama run llama3.2
+# ollama run llama3.2:1b
+import os
+
+from synthetic_dataset_generator import launch
+
+os.environ["OLLAMA_BASE_URL"] = (
+    "http://127.0.0.1:11434/"  # in this case, the same base url for both models
+)
+
+os.environ["MODEL"] = "llama3.2" # model for instruction generation
+os.environ["MODEL_COMPLETION"] = "llama3.2:1b" # model for completion generation
+
+os.environ["TOKENIZER_ID"] = "meta-llama/Llama-3.2-1B-Instruct" # tokenizer for instruction generation
+os.environ["TOKENIZER_ID_COMPLETION"] = "meta-llama/Llama-3.2-3B-Instruct" # tokenizer for completion generation
+
+os.environ["MAGPIE_PRE_QUERY_TEMPLATE"] = "llama3" # magpie template required for instruction generation
+
+launch()

src/synthetic_dataset_generator/apps/base.py

@@ -1,12 +1,16 @@
 import io
 import uuid
+from tqdm import tqdm
 from typing import Union
 
 import argilla as rg
 import gradio as gr
-from datasets import Dataset, concatenate_datasets, load_dataset
+import pandas as pd
+from datasets import Dataset, concatenate_datasets, get_dataset_config_names, get_dataset_split_names, load_dataset
 from gradio import OAuthToken
 from huggingface_hub import HfApi, upload_file, repo_exists
+from unstructured.chunking.title import chunk_by_title
+from unstructured.partition.auto import partition
 
 from synthetic_dataset_generator.constants import MAX_NUM_ROWS
 from synthetic_dataset_generator.utils import get_argilla_client
@@ -64,7 +68,7 @@ def push_pipeline_code_to_hub(
     progress(1.0, desc="Pipeline code uploaded")
 
 
-def validate_push_to_hub(org_name, repo_name):
+def validate_push_to_hub(org_name: str, repo_name: str):
     repo_id = (
         f"{org_name}/{repo_name}"
         if repo_name is not None and org_name is not None
@@ -93,7 +97,7 @@ def combine_datasets(
         return dataset
 
 
-def show_success_message(org_name, repo_name) -> gr.Markdown:
+def show_success_message(org_name: str, repo_name: str) -> gr.Markdown:
     client = get_argilla_client()
     if client is None:
         return gr.Markdown(
@@ -179,3 +183,81 @@ def get_iframe(hub_repo_id: str) -> str:
     ></iframe>
     """
     return iframe
+
+
+def _get_valid_columns(dataframe: pd.DataFrame):
+    doc_valid_columns = []
+
+    for col in dataframe.columns:
+        sample_val = dataframe[col].iloc[0]
+        if isinstance(sample_val, str):
+            doc_valid_columns.append(col)
+
+    return doc_valid_columns
+
+
+def load_dataset_from_hub(
+    repo_id: str,
+    num_rows: int = 10,
+    token: Union[OAuthToken, None] = None,
+    progress=gr.Progress(track_tqdm=True),
+):
+    if not repo_id:
+        raise gr.Error("Please provide a Hub repo ID")
+    subsets = get_dataset_config_names(repo_id, token=token)
+    splits = get_dataset_split_names(repo_id, subsets[0], token=token)
+    ds = load_dataset(repo_id, subsets[0], split=splits[0], token=token, streaming=True)
+    rows = []
+    for idx, row in enumerate(tqdm(ds, desc="Loading the dataset", total=num_rows)):
+        rows.append(row)
+        if idx == num_rows:
+            break
+    ds = Dataset.from_list(rows)
+    dataframe = ds.to_pandas()
+    doc_valid_columns = _get_valid_columns(dataframe)
+    col_doc = doc_valid_columns[0] if doc_valid_columns else ""
+    return (
+        dataframe,
+        gr.Dropdown(
+            choices=doc_valid_columns,
+            label="Documents column",
+            value=col_doc,
+            interactive=(False if col_doc == "" else True),
+            multiselect=False,
+        ),
+    )
+
+
+def preprocess_input_data(
+    file_paths: list[str], num_rows: int, progress=gr.Progress(track_tqdm=True)
+):
+    if not file_paths:
+        raise gr.Error("Please provide an input file")
+
+    data = {}
+    total_chunks = 0
+
+    for file_path in tqdm(file_paths, desc="Processing files", total=len(file_paths)):
+        partitioned_file = partition(filename=file_path)
+        chunks = [str(chunk) for chunk in chunk_by_title(partitioned_file)]
+        data[file_path] = chunks
+        total_chunks += len(chunks)
+        if total_chunks >= num_rows:
+            break
+
+    dataframe = pd.DataFrame.from_records(
+        [(k, v) for k, values in data.items() for v in values],
+        columns=["filename", "chunks"],
+    )
+    col_doc = "chunks"
+
+    return (
+        dataframe,
+        gr.Dropdown(
+            choices=["chunks"],
+            label="Documents column",
+            value=col_doc,
+            interactive=(False if col_doc == "" else True),
+            multiselect=False,
+        ),
+    )

src/synthetic_dataset_generator/apps/chat.py

@@ -1,4 +1,5 @@
 import ast
+import json
 import random
 import uuid
 from typing import Dict, List, Union
@@ -8,11 +9,15 @@ import gradio as gr
 import pandas as pd
 from datasets import Dataset
 from distilabel.distiset import Distiset
+from gradio.oauth import OAuthToken
+from gradio_huggingfacehub_search import HuggingfaceHubSearch
 from huggingface_hub import HfApi
 
 from synthetic_dataset_generator.apps.base import (
     combine_datasets,
     hide_success_message,
+    load_dataset_from_hub,
+    preprocess_input_data,
     push_pipeline_code_to_hub,
     show_success_message,
     test_max_num_rows,
@@ -23,21 +28,25 @@ from synthetic_dataset_generator.constants import (
     BASE_URL,
     DEFAULT_BATCH_SIZE,
     MODEL,
+    MODEL_COMPLETION,
     SFT_AVAILABLE,
 )
 from synthetic_dataset_generator.pipelines.base import get_rewritten_prompts
 from synthetic_dataset_generator.pipelines.chat import (
     DEFAULT_DATASET_DESCRIPTIONS,
     generate_pipeline_code,
+    get_follow_up_generator,
     get_magpie_generator,
     get_prompt_generator,
     get_response_generator,
+    get_sentence_pair_generator,
 )
 from synthetic_dataset_generator.pipelines.embeddings import (
     get_embeddings,
     get_sentence_embedding_dimensions,
 )
 from synthetic_dataset_generator.utils import (
+    column_to_list,
     get_argilla_client,
     get_org_dropdown,
     get_random_repo_name,
@@ -45,6 +54,14 @@ from synthetic_dataset_generator.utils import (
 )
 
 
+def _get_dataframe():
+    return gr.Dataframe(
+        headers=["prompt", "completion"],
+        wrap=True,
+        interactive=False,
+    )
+
+
 def convert_dataframe_messages(dataframe: pd.DataFrame) -> pd.DataFrame:
     def convert_to_list_of_dicts(messages: str) -> List[Dict[str, str]]:
         return ast.literal_eval(
@@ -60,7 +77,7 @@ def convert_dataframe_messages(dataframe: pd.DataFrame) -> pd.DataFrame:
     return dataframe
 
 
-def generate_system_prompt(dataset_description, progress=gr.Progress()):
+def generate_system_prompt(dataset_description: str, progress=gr.Progress()):
     progress(0.1, desc="Initializing")
     generate_description = get_prompt_generator()
     progress(0.5, desc="Generating")
@@ -77,42 +94,73 @@ def generate_system_prompt(dataset_description, progress=gr.Progress()):
     return result
 
 
-def generate_sample_dataset(system_prompt, num_turns, progress=gr.Progress()):
-    progress(0.1, desc="Generating sample dataset")
+def load_dataset_file(
+    repo_id: str,
+    file_paths: list[str],
+    input_type: str,
+    num_rows: int = 10,
+    token: Union[OAuthToken, None] = None,
+    progress=gr.Progress(),
+):
+    progress(0.1, desc="Loading the source data")
+    if input_type == "dataset-input":
+        return load_dataset_from_hub(repo_id=repo_id, num_rows=num_rows, token=token)
+    else:
+        return preprocess_input_data(file_paths=file_paths, num_rows=num_rows)
+
+
+def generate_sample_dataset(
+    repo_id: str,
+    file_paths: list[str],
+    input_type: str,
+    system_prompt: str,
+    document_column: str,
+    num_turns: int,
+    num_rows: int,
+    oauth_token: Union[OAuthToken, None],
+    progress=gr.Progress(),
+):
+    if input_type == "prompt-input":
+        dataframe = pd.DataFrame(columns=["prompt", "completion"])
+    else:
+        dataframe, _ = load_dataset_file(
+            repo_id=repo_id,
+            file_paths=file_paths,
+            input_type=input_type,
+            num_rows=num_rows,
+            token=oauth_token,
+        )
+    progress(0.5, desc="Generating sample dataset")
     dataframe = generate_dataset(
+        input_type=input_type,
+        dataframe=dataframe,
         system_prompt=system_prompt,
+        document_column=document_column,
         num_turns=num_turns,
-        num_rows=10,
-        progress=progress,
+        num_rows=num_rows,
         is_sample=True,
     )
     progress(1.0, desc="Sample dataset generated")
     return dataframe
 
 
-def _get_dataframe():
-    return gr.Dataframe(
-        headers=["prompt", "completion"],
-        wrap=True,
-        interactive=False,
-    )
-
-
-def generate_dataset(
+def generate_dataset_from_prompt(
     system_prompt: str,
     num_turns: int = 1,
     num_rows: int = 10,
     temperature: float = 0.9,
+    temperature_completion: Union[float, None] = None,
     is_sample: bool = False,
     progress=gr.Progress(),
 ) -> pd.DataFrame:
     num_rows = test_max_num_rows(num_rows)
     progress(0.0, desc="(1/2) Generating instructions")
-    magpie_generator = get_magpie_generator(
-        system_prompt, num_turns, temperature, is_sample
-    )
+    magpie_generator = get_magpie_generator(num_turns, temperature, is_sample)
     response_generator = get_response_generator(
-        system_prompt, num_turns, temperature, is_sample
+        system_prompt=system_prompt,
+        num_turns=num_turns,
+        temperature=temperature or temperature_completion,
+        is_sample=is_sample,
     )
     total_steps: int = num_rows * 2
     batch_size = DEFAULT_BATCH_SIZE
@@ -217,6 +265,180 @@ def generate_dataset(
     return dataframe
 
 
+def generate_dataset_from_seed(
+    dataframe: pd.DataFrame,
+    document_column: str,
+    num_turns: int = 1,
+    num_rows: int = 10,
+    temperature: float = 0.9,
+    temperature_completion: Union[float, None] = None,
+    is_sample: bool = False,
+    progress=gr.Progress(),
+) -> pd.DataFrame:
+    num_rows = test_max_num_rows(num_rows)
+    progress(0.0, desc="Initializing dataset generation")
+    document_data = column_to_list(dataframe, document_column)
+    if len(document_data) < num_rows:
+        document_data += random.choices(document_data, k=num_rows - len(document_data))
+    instruction_generator = get_sentence_pair_generator(
+        temperature=temperature, is_sample=is_sample
+    )
+    response_generator = get_response_generator(
+        system_prompt=None,
+        num_turns=1,
+        temperature=temperature or temperature_completion,
+        is_sample=is_sample,
+    )
+    follow_up_generator_instruction = get_follow_up_generator(
+        type="instruction", temperature=temperature, is_sample=is_sample
+    )
+    follow_up_generator_response = get_follow_up_generator(
+        type="response",
+        temperature=temperature or temperature_completion,
+        is_sample=is_sample,
+    )
+    steps = 2 * num_turns
+    total_steps: int = num_rows * steps
+    step_progress = round(1 / steps, 2)
+    batch_size = DEFAULT_BATCH_SIZE
+
+    # create instructions
+    n_processed = 0
+    instruction_results = []
+    while n_processed < num_rows:
+        progress(
+            step_progress * n_processed / num_rows,
+            total=total_steps,
+            desc="Generating questions",
+        )
+        remaining_rows = num_rows - n_processed
+        batch_size = min(batch_size, remaining_rows)
+        batch = [
+            {"anchor": document}
+            for document in document_data[n_processed : n_processed + batch_size]
+        ]
+        questions = list(instruction_generator.process(inputs=batch))
+        instruction_results.extend(questions[0])
+        n_processed += batch_size
+    for result in instruction_results:
+        result["instruction"] = result["positive"]
+        result["prompt"] = result.pop("positive")
+
+    progress(step_progress, desc="Generating instructions")
+
+    # generate responses
+    n_processed = 0
+    response_results = []
+    while n_processed < num_rows:
+        progress(
+            step_progress + step_progress * n_processed / num_rows,
+            total=total_steps,
+            desc="Generating responses",
+        )
+        batch = instruction_results[n_processed : n_processed + batch_size]
+        responses = list(response_generator.process(inputs=batch))
+        response_results.extend(responses[0])
+        n_processed += batch_size
+    for result in response_results:
+        result["completion"] = result.pop("generation")
+
+    # generate follow-ups
+    if num_turns > 1:
+        n_processed = 0
+        final_conversations = []
+
+        while n_processed < num_rows:
+            progress(
+                step_progress + step_progress * n_processed / num_rows,
+                total=total_steps,
+                desc="Generating follow-ups",
+            )
+            batch = response_results[n_processed : n_processed + batch_size]
+            conversations_batch = [
+                {
+                    "messages": [
+                        {"role": "user", "content": result["prompt"]},
+                        {"role": "assistant", "content": result["completion"]},
+                    ]
+                }
+                for result in batch
+            ]
+
+            for _ in range(num_turns - 1):
+                follow_up_instructions = list(
+                    follow_up_generator_instruction.process(inputs=conversations_batch)
+                )
+                for conv, follow_up in zip(conversations_batch, follow_up_instructions[0]):
+                    conv["messages"].append(
+                        {"role": "user", "content": follow_up["generation"]}
+                    )
+
+                follow_up_responses = list(
+                    follow_up_generator_response.process(inputs=conversations_batch)
+                )
+                for conv, follow_up in zip(conversations_batch, follow_up_responses[0]):
+                    conv["messages"].append(
+                        {"role": "assistant", "content": follow_up["generation"]}
+                    )
+
+            final_conversations.extend(
+                [{"messages": conv["messages"]} for conv in conversations_batch]
+            )
+            n_processed += batch_size
+
+    # create distiset
+    distiset_results = []
+    if num_turns == 1:
+        for result in response_results:
+            record = {}
+            for relevant_keys in ["prompt", "completion"]:
+                if relevant_keys in result:
+                    record[relevant_keys] = result[relevant_keys]
+            distiset_results.append(record)
+        dataframe = pd.DataFrame(distiset_results)
+    else:
+        distiset_results = final_conversations
+        dataframe = pd.DataFrame(distiset_results)
+        dataframe["messages"] = dataframe["messages"].apply(lambda x: json.dumps(x))
+
+    progress(1.0, desc="Dataset generation completed")
+    return dataframe
+
+
+def generate_dataset(
+    input_type: str,
+    dataframe: pd.DataFrame,
+    system_prompt: str,
+    document_column: str,
+    num_turns: int = 1,
+    num_rows: int = 10,
+    temperature: float = 0.9,
+    temperature_completion: Union[float, None] = None,
+    is_sample: bool = False,
+    progress=gr.Progress(),
+) -> pd.DataFrame:
+    if input_type == "prompt-input":
+        dataframe = generate_dataset_from_prompt(
+            system_prompt=system_prompt,
+            num_turns=num_turns,
+            num_rows=num_rows,
+            temperature=temperature,
+            temperature_completion=temperature_completion,
+            is_sample=is_sample,
+        )
+    else:
+        dataframe = generate_dataset_from_seed(
+            dataframe=dataframe,
+            document_column=document_column,
+            num_turns=num_turns,
+            num_rows=num_rows,
+            temperature=temperature,
+            temperature_completion=temperature_completion,
+            is_sample=is_sample,
+        )
+    return dataframe
+
+
 def push_dataset_to_hub(
     dataframe: pd.DataFrame,
     org_name: str,
@@ -251,23 +473,48 @@ def push_dataset_to_hub(
 def push_dataset(
     org_name: str,
     repo_name: str,
+    private: bool,
+    original_repo_id: str,
+    file_paths: list[str],
+    input_type: str,
     system_prompt: str,
+    document_column: str,
     num_turns: int = 1,
     num_rows: int = 10,
-    private: bool = False,
     temperature: float = 0.9,
+    temperature_completion: Union[float, None] = None,
     pipeline_code: str = "",
     oauth_token: Union[gr.OAuthToken, None] = None,
     progress=gr.Progress(),
 ) -> pd.DataFrame:
+    if input_type == "prompt-input":
+        dataframe = _get_dataframe()
+    else:
+        dataframe, _ = load_dataset_file(
+            repo_id=original_repo_id,
+            file_paths=file_paths,
+            input_type=input_type,
+            num_rows=num_rows,
+            token=oauth_token,
+        )
+    progress(0.5, desc="Generating dataset")
     dataframe = generate_dataset(
+        input_type=input_type,
+        dataframe=dataframe,
         system_prompt=system_prompt,
+        document_column=document_column,
         num_turns=num_turns,
         num_rows=num_rows,
         temperature=temperature,
+        temperature_completion=temperature_completion
     )
     push_dataset_to_hub(
-        dataframe, org_name, repo_name, oauth_token, private, pipeline_code
+        dataframe=dataframe,
+        org_name=org_name,
+        repo_name=repo_name,
+        oauth_token=oauth_token,
+        private=private,
+        pipeline_code=pipeline_code,
     )
     try:
         progress(0.1, desc="Setting up user and workspace")
@@ -390,6 +637,28 @@ def push_dataset(
     return ""
 
 
+def show_system_prompt_visibility():
+    return {system_prompt: gr.Textbox(visible=True)}
+
+
+def hide_system_prompt_visibility():
+    return {system_prompt: gr.Textbox(visible=False)}
+
+
+def show_document_column_visibility():
+    return {document_column: gr.Dropdown(visible=True)}
+
+
+def hide_document_column_visibility():
+    return {
+        document_column: gr.Dropdown(
+            choices=["Load your data first in step 1."],
+            value="Load your data first in step 1.",
+            visible=False,
+        )
+    }
+
+
 def show_pipeline_code_visibility():
     return {pipeline_code_ui: gr.Accordion(visible=True)}
 
@@ -398,6 +667,11 @@ def hide_pipeline_code_visibility():
     return {pipeline_code_ui: gr.Accordion(visible=False)}
 
 
+def show_temperature_completion():
+    if MODEL != MODEL_COMPLETION:
+        return {temperature_completion: gr.Slider(value=0.9, visible=True)}
+
+
 ######################
 # Gradio UI
 ######################
@@ -417,29 +691,85 @@ with gr.Blocks() as app:
                 )
             )
         else:
-            gr.Markdown(value="## 1. Describe the dataset you want")
-            with gr.Row():
+            gr.Markdown("## 1. Select your input")
+            with gr.Row(equal_height=False):
                 with gr.Column(scale=2):
-                    dataset_description = gr.Textbox(
-                        label="Dataset description",
-                        placeholder="Give a precise description of your desired dataset.",
-                    )
-                    with gr.Row():
-                        clear_btn_part = gr.Button(
-                            "Clear",
-                            variant="secondary",
-                        )
-                        load_btn = gr.Button(
-                            "Create",
-                            variant="primary",
-                        )
-                with gr.Column(scale=3):
-                    examples = gr.Examples(
-                        examples=DEFAULT_DATASET_DESCRIPTIONS,
-                        inputs=[dataset_description],
-                        cache_examples=False,
-                        label="Examples",
+                    input_type = gr.Dropdown(
+                        label="Input type",
+                        choices=["prompt-input", "dataset-input", "file-input"],
+                        value="prompt-input",
+                        multiselect=False,
+                        visible=False,
                     )
+                    with gr.Tab("Generate from prompt") as tab_prompt_input:
+                        with gr.Row(equal_height=False):
+                            with gr.Column(scale=2):
+                                dataset_description = gr.Textbox(
+                                    label="Dataset description",
+                                    placeholder="Give a precise description of your desired dataset.",
+                                )
+                                with gr.Row():
+                                    clear_prompt_btn_part = gr.Button(
+                                        "Clear", variant="secondary"
+                                    )
+                                    load_prompt_btn = gr.Button(
+                                        "Create", variant="primary"
+                                    )
+                            with gr.Column(scale=3):
+                                examples = gr.Examples(
+                                    examples=DEFAULT_DATASET_DESCRIPTIONS,
+                                    inputs=[dataset_description],
+                                    cache_examples=False,
+                                    label="Examples",
+                                )
+                    with gr.Tab("Load from Hub") as tab_dataset_input:
+                        with gr.Row(equal_height=False):
+                            with gr.Column(scale=2):
+                                search_in = HuggingfaceHubSearch(
+                                    label="Search",
+                                    placeholder="Search for a dataset",
+                                    search_type="dataset",
+                                    sumbit_on_select=True,
+                                )
+                                with gr.Row():
+                                    clear_dataset_btn_part = gr.Button(
+                                        "Clear", variant="secondary"
+                                    )
+                                    load_dataset_btn = gr.Button(
+                                        "Load", variant="primary"
+                                    )
+                            with gr.Column(scale=3):
+                                examples = gr.Examples(
+                                    examples=[
+                                        "charris/wikipedia_sample",
+                                        "plaguss/argilla_sdk_docs_raw_unstructured",
+                                        "BeIR/hotpotqa-generated-queries",
+                                    ],
+                                    label="Example datasets",
+                                    fn=lambda x: x,
+                                    inputs=[search_in],
+                                    run_on_click=True,
+                                )
+                                search_out = gr.HTML(
+                                    label="Dataset preview", visible=False
+                                )
+                    with gr.Tab("Load your file") as tab_file_input:
+                        with gr.Row(equal_height=False):
+                            with gr.Column(scale=2):
+                                file_in = gr.File(
+                                    label="Upload your file. Supported formats: .md, .txt, .docx, .pdf",
+                                    file_count="multiple",
+                                    file_types=[".md", ".txt", ".docx", ".pdf"],
+                                )
+                                with gr.Row():
+                                    clear_file_btn_part = gr.Button(
+                                        "Clear", variant="secondary"
+                                    )
+                                    load_file_btn = gr.Button("Load", variant="primary")
+                            with gr.Column(scale=3):
+                                file_out = gr.HTML(
+                                    label="Dataset preview", visible=False
+                                )
 
             gr.HTML(value="<hr>")
             gr.Markdown(value="## 2. Configure your dataset")
@@ -449,6 +779,16 @@ with gr.Blocks() as app:
                         label="System prompt",
                         placeholder="You are a helpful assistant.",
                     )
+                    document_column = gr.Dropdown(
+                        label="Document Column",
+                        info="Select the document column to generate the RAG dataset",
+                        choices=["Load your data first in step 1."],
+                        value="Load your data first in step 1.",
+                        interactive=False,
+                        multiselect=False,
+                        allow_custom_value=False,
+                        visible=False,
+                    )
                     num_turns = gr.Number(
                         value=1,
                         label="Number of turns in the conversation",
@@ -489,11 +829,20 @@ with gr.Blocks() as app:
                     temperature = gr.Slider(
                         label="Temperature",
                         minimum=0.1,
-                        maximum=1,
+                        maximum=1.5,
                         value=0.9,
                         step=0.1,
                         interactive=True,
                     )
+                    temperature_completion = gr.Slider(
+                        label="Temperature for completion",
+                        minimum=0.1,
+                        maximum=1.5,
+                        value=None,
+                        step=0.1,
+                        interactive=True,
+                        visible=False,
+                    )
                     private = gr.Checkbox(
                         label="Private dataset",
                         value=False,
@@ -514,7 +863,10 @@ with gr.Blocks() as app:
                         visible=False,
                     ) as pipeline_code_ui:
                         code = generate_pipeline_code(
+                            repo_id=search_in.value,
+                            input_type=input_type.value,
                             system_prompt=system_prompt.value,
+                            document_column=document_column.value,
                             num_turns=num_turns.value,
                             num_rows=num_rows.value,
                         )
@@ -524,77 +876,138 @@ with gr.Blocks() as app:
                             label="Distilabel Pipeline Code",
                         )
 
-            load_btn.click(
-                fn=generate_system_prompt,
-                inputs=[dataset_description],
-                outputs=[system_prompt],
-                show_progress=True,
-            ).then(
-                fn=generate_sample_dataset,
-                inputs=[system_prompt, num_turns],
-                outputs=[dataframe],
-                show_progress=True,
-            )
+    tab_prompt_input.select(
+        fn=lambda: "prompt-input",
+        inputs=[],
+        outputs=[input_type],
+    ).then(fn=show_system_prompt_visibility, inputs=[], outputs=[system_prompt]).then(
+        fn=hide_document_column_visibility, inputs=[], outputs=[document_column]
+    )
 
-            btn_apply_to_sample_dataset.click(
-                fn=generate_sample_dataset,
-                inputs=[system_prompt, num_turns],
-                outputs=[dataframe],
-                show_progress=True,
-            )
+    tab_dataset_input.select(
+        fn=lambda: "dataset-input",
+        inputs=[],
+        outputs=[input_type],
+    ).then(fn=hide_system_prompt_visibility, inputs=[], outputs=[system_prompt]).then(
+        fn=show_document_column_visibility, inputs=[], outputs=[document_column]
+    )
 
-            btn_push_to_hub.click(
-                fn=validate_argilla_user_workspace_dataset,
-                inputs=[repo_name],
-                outputs=[success_message],
-                show_progress=True,
-            ).then(
-                fn=validate_push_to_hub,
-                inputs=[org_name, repo_name],
-                outputs=[success_message],
-                show_progress=True,
-            ).success(
-                fn=hide_success_message,
-                outputs=[success_message],
-                show_progress=True,
-            ).success(
-                fn=hide_pipeline_code_visibility,
-                inputs=[],
-                outputs=[pipeline_code_ui],
-                show_progress=True,
-            ).success(
-                fn=push_dataset,
-                inputs=[
-                    org_name,
-                    repo_name,
-                    system_prompt,
-                    num_turns,
-                    num_rows,
-                    private,
-                    temperature,
-                    pipeline_code,
-                ],
-                outputs=[success_message],
-                show_progress=True,
-            ).success(
-                fn=show_success_message,
-                inputs=[org_name, repo_name],
-                outputs=[success_message],
-            ).success(
-                fn=generate_pipeline_code,
-                inputs=[system_prompt, num_turns, num_rows],
-                outputs=[pipeline_code],
-            ).success(
-                fn=show_pipeline_code_visibility,
-                inputs=[],
-                outputs=[pipeline_code_ui],
-            )
-            gr.on(
-                triggers=[clear_btn_part.click, clear_btn_full.click],
-                fn=lambda _: ("", "", 1, _get_dataframe()),
-                inputs=[dataframe],
-                outputs=[system_prompt, num_turns, dataframe],
-            )
-            app.load(fn=get_org_dropdown, outputs=[org_name])
-        app.load(fn=get_random_repo_name, outputs=[repo_name])
-        app.load(fn=swap_visibility, outputs=main_ui)
+    tab_file_input.select(
+        fn=lambda: "file-input",
+        inputs=[],
+        outputs=[input_type],
+    ).then(fn=hide_system_prompt_visibility, inputs=[], outputs=[system_prompt]).then(
+        fn=show_document_column_visibility, inputs=[], outputs=[document_column]
+    )
+
+    search_in.submit(
+        fn=lambda df: pd.DataFrame(columns=df.columns),
+        inputs=[dataframe],
+        outputs=[dataframe],
+    )
+
+    load_prompt_btn.click(
+        fn=generate_system_prompt,
+        inputs=[dataset_description],
+        outputs=[system_prompt],
+    ).success(
+        fn=generate_sample_dataset,
+        inputs=[
+            search_in,
+            file_in,
+            input_type,
+            system_prompt,
+            document_column,
+            num_turns,
+            num_rows,
+        ],
+        outputs=dataframe,
+    )
+
+    gr.on(
+        triggers=[load_dataset_btn.click, load_file_btn.click],
+        fn=load_dataset_file,
+        inputs=[search_in, file_in, input_type],
+        outputs=[dataframe, document_column],
+    )
+
+    btn_apply_to_sample_dataset.click(
+        fn=generate_sample_dataset,
+        inputs=[
+            search_in,
+            file_in,
+            input_type,
+            system_prompt,
+            document_column,
+            num_turns,
+            num_rows,
+        ],
+        outputs=dataframe,
+    )
+
+    btn_push_to_hub.click(
+        fn=validate_argilla_user_workspace_dataset,
+        inputs=[repo_name],
+        outputs=[success_message],
+    ).then(
+        fn=validate_push_to_hub,
+        inputs=[org_name, repo_name],
+        outputs=[success_message],
+    ).success(
+        fn=hide_success_message,
+        outputs=[success_message],
+    ).success(
+        fn=hide_pipeline_code_visibility,
+        inputs=[],
+        outputs=[pipeline_code_ui],
+    ).success(
+        fn=push_dataset,
+        inputs=[
+            org_name,
+            repo_name,
+            private,
+            search_in,
+            file_in,
+            input_type,
+            system_prompt,
+            document_column,
+            num_turns,
+            num_rows,
+            temperature,
+            temperature_completion,
+            pipeline_code,
+        ],
+        outputs=[success_message],
+    ).success(
+        fn=show_success_message,
+        inputs=[org_name, repo_name],
+        outputs=[success_message],
+    ).success(
+        fn=generate_pipeline_code,
+        inputs=[
+            search_in,
+            input_type,
+            system_prompt,
+            document_column,
+            num_turns,
+            num_rows,
+        ],
+        outputs=[pipeline_code],
+    ).success(
+        fn=show_pipeline_code_visibility,
+        inputs=[],
+        outputs=[pipeline_code_ui],
+    )
+
+    clear_dataset_btn_part.click(fn=lambda: "", inputs=[], outputs=[search_in])
+    clear_file_btn_part.click(fn=lambda: None, inputs=[], outputs=[file_in])
+    clear_prompt_btn_part.click(fn=lambda: "", inputs=[], outputs=[dataset_description])
+    clear_btn_full.click(
+        fn=lambda df: ("", "", [], _get_dataframe()),
+        inputs=[dataframe],
+        outputs=[system_prompt, document_column, num_turns, dataframe],
+    )
+    app.load(fn=swap_visibility, outputs=main_ui)
+    app.load(fn=get_org_dropdown, outputs=[org_name])
+    app.load(fn=get_random_repo_name, outputs=[repo_name])
+    app.load(fn=show_temperature_completion, outputs=[temperature_completion])

src/synthetic_dataset_generator/apps/eval.py

@@ -15,7 +15,7 @@ from datasets import (
 from distilabel.distiset import Distiset
 from gradio.oauth import OAuthToken  #
 from gradio_huggingfacehub_search import HuggingfaceHubSearch
-from huggingface_hub import HfApi, repo_exists
+from huggingface_hub import HfApi
 
 from synthetic_dataset_generator.apps.base import (
     combine_datasets,
@@ -130,9 +130,9 @@ def load_dataset_from_hub(
             choices=response_valid_columns,
             label="Response column",
             value=col_response,
-            interactive=False
-            if col_response == "No valid response columns found."
-            else True,
+            interactive=(
+                False if col_response == "No valid response columns found." else True
+            ),
         ),
         prompt_template,
         structured_output,
@@ -831,16 +831,13 @@ with gr.Blocks() as app:
         fn=validate_argilla_user_workspace_dataset,
         inputs=[repo_name],
         outputs=[success_message],
-        show_progress=True,
     ).then(
         fn=validate_push_to_hub,
         inputs=[org_name, repo_name],
         outputs=[success_message],
-        show_progress=True,
     ).success(
         fn=hide_success_message,
         outputs=[success_message],
-        show_progress=True,
     ).success(
         fn=hide_pipeline_code_visibility,
         inputs=[],
@@ -862,7 +859,6 @@ with gr.Blocks() as app:
             pipeline_code,
         ],
         outputs=[success_message],
-        show_progress=True,
     ).success(
         fn=show_success_message,
         inputs=[org_name, repo_name],
@@ -882,13 +878,14 @@ with gr.Blocks() as app:
         outputs=[pipeline_code_ui],
     )
 
-    clear_btn_part.click(fn=lambda : "", inputs=[], outputs=[search_in])
+    clear_btn_part.click(fn=lambda: "", inputs=[], outputs=[search_in])
     clear_btn_full.click(
         fn=lambda df: ("", "", pd.DataFrame(columns=df.columns)),
         inputs=[dataframe],
         outputs=[
             instruction_instruction_response,
             response_instruction_response,
+            dataframe,
         ],
     )
 

src/synthetic_dataset_generator/apps/rag.py

@@ -1,37 +1,30 @@
 import os
 import random
 import uuid
-from tqdm import tqdm
 from typing import Union
 
 import argilla as rg
 import gradio as gr
 import nltk
 import pandas as pd
-from datasets import (
-    Dataset,
-    get_dataset_config_names,
-    get_dataset_split_names,
-    load_dataset,
-)
+from datasets import Dataset
 from distilabel.distiset import Distiset
 from gradio.oauth import OAuthToken
 from gradio_huggingfacehub_search import HuggingfaceHubSearch
 from huggingface_hub import HfApi
-from unstructured.chunking.title import chunk_by_title
-from unstructured.partition.auto import partition
 
 from synthetic_dataset_generator.apps.base import (
     combine_datasets,
-    get_iframe,
     hide_success_message,
+    load_dataset_from_hub,
+    preprocess_input_data,
     push_pipeline_code_to_hub,
     show_success_message,
     test_max_num_rows,
     validate_argilla_user_workspace_dataset,
     validate_push_to_hub,
 )
-from synthetic_dataset_generator.constants import DEFAULT_BATCH_SIZE
+from synthetic_dataset_generator.constants import DEFAULT_BATCH_SIZE, MODEL, MODEL_COMPLETION
 from synthetic_dataset_generator.pipelines.base import get_rewritten_prompts
 from synthetic_dataset_generator.pipelines.embeddings import (
     get_embeddings,
@@ -39,11 +32,11 @@ from synthetic_dataset_generator.pipelines.embeddings import (
 )
 from synthetic_dataset_generator.pipelines.rag import (
     DEFAULT_DATASET_DESCRIPTIONS,
+    generate_pipeline_code,
     get_chunks_generator,
     get_prompt_generator,
-    generate_pipeline_code,
-    get_sentence_pair_generator,
     get_response_generator,
+    get_sentence_pair_generator,
 )
 from synthetic_dataset_generator.utils import (
     column_to_list,
@@ -58,80 +51,8 @@ nltk.data.path.append("./nltk_data")
 nltk.download("punkt_tab", download_dir="./nltk_data")
 nltk.download("averaged_perceptron_tagger_eng", download_dir="./nltk_data")
 
-def _get_valid_columns(dataframe: pd.DataFrame):
-    doc_valid_columns = []
-
-    for col in dataframe.columns:
-        sample_val = dataframe[col].iloc[0]
-        if isinstance(sample_val, str):
-            doc_valid_columns.append(col)
-
-    return doc_valid_columns
-
-
-def _load_dataset_from_hub(
-    repo_id: str,
-    num_rows: int = 10,
-    token: Union[OAuthToken, None] = None,
-    progress=gr.Progress(track_tqdm=True),
-):
-    if not repo_id:
-        raise gr.Error("Hub repo id is required")
-    subsets = get_dataset_config_names(repo_id, token=token)
-    splits = get_dataset_split_names(repo_id, subsets[0], token=token)
-    ds = load_dataset(repo_id, subsets[0], split=splits[0], token=token, streaming=True)
-    rows = []
-    for idx, row in enumerate(tqdm(ds, desc="Loading the dataset", total=num_rows)):
-        rows.append(row)
-        if idx == num_rows:
-            break
-    ds = Dataset.from_list(rows)
-    dataframe = ds.to_pandas()
-    doc_valid_columns = _get_valid_columns(dataframe)
-    col_doc = doc_valid_columns[0] if doc_valid_columns else ""
-    return (
-        dataframe,
-        gr.Dropdown(
-            choices=doc_valid_columns,
-            label="Documents column",
-            value=col_doc,
-            interactive=(False if col_doc == "" else True),
-            multiselect=False,
-        ),
-    )
-
-
-def _preprocess_input_data(file_paths, num_rows, progress=gr.Progress(track_tqdm=True)):
-    data = {}
-    total_chunks = 0
 
-    for file_path in tqdm(file_paths, desc="Processing files", total=len(file_paths)):
-        partitioned_file = partition(filename=file_path)
-        chunks = [str(chunk) for chunk in chunk_by_title(partitioned_file)]
-        data[file_path] = chunks
-        total_chunks += len(chunks)
-        if total_chunks >= num_rows:
-            break
-
-    dataframe = pd.DataFrame.from_records(
-        [(k, v) for k, values in data.items() for v in values],
-        columns=["filename", "chunks"],
-    )
-    col_doc = "chunks"
-
-    return (
-        dataframe,
-        gr.Dropdown(
-            choices=["chunks"],
-            label="Documents column",
-            value=col_doc,
-            interactive=(False if col_doc == "" else True),
-            multiselect=False,
-        ),
-    )
-
-
-def generate_system_prompt(dataset_description, progress=gr.Progress()):
+def generate_system_prompt(dataset_description: str, progress=gr.Progress()):
     progress(0.1, desc="Initializing")
     generate_description = get_prompt_generator()
     progress(0.5, desc="Generating")
@@ -158,9 +79,48 @@ def load_dataset_file(
 ):
     progress(0.1, desc="Loading the source data")
     if input_type == "dataset-input":
-        return _load_dataset_from_hub(repo_id, num_rows, token)
+        return load_dataset_from_hub(repo_id=repo_id, num_rows=num_rows, token=token)
+    else:
+        return preprocess_input_data(file_paths=file_paths, num_rows=num_rows)
+
+
+def generate_sample_dataset(
+    repo_id: str,
+    file_paths: list[str],
+    input_type: str,
+    system_prompt: str,
+    document_column: str,
+    retrieval_reranking: list[str],
+    num_rows: str,
+    oauth_token: Union[OAuthToken, None],
+    progress=gr.Progress(),
+):
+    retrieval = "Retrieval" in retrieval_reranking
+    reranking = "Reranking" in retrieval_reranking
+
+    if input_type == "prompt-input":
+        dataframe = pd.DataFrame(columns=["context", "question", "response"])
     else:
-        return _preprocess_input_data(file_paths, num_rows)
+        dataframe, _ = load_dataset_file(
+            repo_id=repo_id,
+            file_paths=file_paths,
+            input_type=input_type,
+            num_rows=num_rows,
+            token=oauth_token,
+        )
+    progress(0.5, desc="Generating dataset")
+    dataframe = generate_dataset(
+        input_type=input_type,
+        dataframe=dataframe,
+        system_prompt=system_prompt,
+        document_column=document_column,
+        retrieval=retrieval,
+        reranking=reranking,
+        num_rows=10,
+        is_sample=True,
+    )
+    progress(1.0, desc="Sample dataset generated")
+    return dataframe
 
 
 def generate_dataset(
@@ -172,6 +132,7 @@ def generate_dataset(
     reranking: bool = False,
     num_rows: int = 10,
     temperature: float = 0.7,
+    temperature_completion: Union[float, None] = None,
     is_sample: bool = False,
     progress=gr.Progress(),
 ):
@@ -195,7 +156,7 @@ def generate_dataset(
         is_sample=is_sample,
     )
     response_generator = get_response_generator(
-        temperature=temperature, is_sample=is_sample
+        temperature = temperature_completion or temperature , is_sample=is_sample
     )
     if reranking:
         reranking_generator = get_sentence_pair_generator(
@@ -320,44 +281,6 @@ def generate_dataset(
     return dataframe
 
 
-def generate_sample_dataset(
-    repo_id: str,
-    file_paths: list[str],
-    input_type: str,
-    system_prompt: str,
-    document_column: str,
-    retrieval_reranking: list[str],
-    num_rows: str,
-    oauth_token: Union[OAuthToken, None],
-    progress=gr.Progress(),
-):
-    retrieval = "Retrieval" in retrieval_reranking
-    reranking = "Reranking" in retrieval_reranking
-
-    if input_type == "prompt-input":
-        dataframe = pd.DataFrame(columns=["context", "question", "response"])
-    else:
-        dataframe, _ = load_dataset_file(
-            repo_id=repo_id,
-            file_paths=file_paths,
-            input_type=input_type,
-            num_rows=num_rows,
-            token=oauth_token,
-        )
-    progress(0.5, desc="Generating dataset")
-    dataframe = generate_dataset(
-        input_type=input_type,
-        dataframe=dataframe,
-        system_prompt=system_prompt,
-        document_column=document_column,
-        retrieval=retrieval,
-        reranking=reranking,
-        num_rows=10,
-        is_sample=True,
-    )
-    return dataframe
-
-
 def push_dataset_to_hub(
     dataframe: pd.DataFrame,
     org_name: str,
@@ -398,6 +321,7 @@ def push_dataset(
     retrieval_reranking: list[str],
     num_rows: int,
     temperature: float,
+    temperature_completion: float,
     pipeline_code: str,
     oauth_token: Union[gr.OAuthToken, None] = None,
     progress=gr.Progress(),
@@ -425,15 +349,14 @@ def push_dataset(
         reranking=reranking,
         num_rows=num_rows,
         temperature=temperature,
+        temperature_completion=temperature_completion,
         is_sample=True,
     )
     push_dataset_to_hub(
         dataframe, org_name, repo_name, oauth_token, private, pipeline_code
     )
     dataframe = dataframe[
-        dataframe.applymap(
-            lambda x: str(x).strip() if pd.notna(x) else x
-        ).apply(
+        dataframe.applymap(lambda x: str(x).strip() if pd.notna(x) else x).apply(
             lambda row: row.notna().all() and (row != "").all(), axis=1
         )
     ]
@@ -593,6 +516,11 @@ def hide_pipeline_code_visibility():
     return {pipeline_code_ui: gr.Accordion(visible=False)}
 
 
+def show_temperature_completion():
+    if MODEL != MODEL_COMPLETION:
+        return {temperature_completion: gr.Slider(value=0.9, visible=True)}
+
+
 ######################
 # Gradio UI
 ######################
@@ -674,40 +602,37 @@ with gr.Blocks() as app:
 
         gr.HTML(value="<hr>")
         gr.Markdown(value="## 2. Configure your task")
-        with gr.Row(equal_height=True):
-            with gr.Row(equal_height=False):
-                with gr.Column(scale=2):
-                    system_prompt = gr.Textbox(
-                        label="System prompt",
-                        placeholder="You are a helpful assistant.",
-                        visible=False,
-                    )
-                    document_column = gr.Dropdown(
-                        label="Document Column",
-                        info="Select the document column to generate the RAG dataset",
-                        choices=["Load your data first in step 1."],
-                        value="Load your data first in step 1.",
-                        interactive=False,
-                        multiselect=False,
-                        allow_custom_value=False,
-                    )
-                    retrieval_reranking = gr.CheckboxGroup(
-                        choices=[("Retrieval", "Retrieval"), ("Reranking", "Reranking")],
-                        type="value",
-                        label="Data for RAG",
-                        info="Indicate the additional data you want to generate for RAG.",
-                    )
-                    with gr.Row():
-                        clear_btn_full = gr.Button("Clear", variant="secondary")
-                        btn_apply_to_sample_dataset = gr.Button(
-                            "Save", variant="primary"
-                        )
-                with gr.Column(scale=3):
-                    dataframe = gr.Dataframe(
-                        headers=["context", "question", "response"],
-                        wrap=True,
-                        interactive=False,
-                    )
+        with gr.Row(equal_height=False):
+            with gr.Column(scale=2):
+                system_prompt = gr.Textbox(
+                    label="System prompt",
+                    placeholder="You are a helpful assistant.",
+                    visible=False,
+                )
+                document_column = gr.Dropdown(
+                    label="Document Column",
+                    info="Select the document column to generate the RAG dataset",
+                    choices=["Load your data first in step 1."],
+                    value="Load your data first in step 1.",
+                    interactive=False,
+                    multiselect=False,
+                    allow_custom_value=False,
+                )
+                retrieval_reranking = gr.CheckboxGroup(
+                    choices=[("Retrieval", "Retrieval"), ("Reranking", "Reranking")],
+                    type="value",
+                    label="Data for RAG",
+                    info="Indicate the additional data you want to generate for RAG.",
+                )
+                with gr.Row():
+                    clear_btn_full = gr.Button("Clear", variant="secondary")
+                    btn_apply_to_sample_dataset = gr.Button("Save", variant="primary")
+            with gr.Column(scale=3):
+                dataframe = gr.Dataframe(
+                    headers=["context", "question", "response"],
+                    wrap=True,
+                    interactive=False,
+                )
 
         gr.HTML(value="<hr>")
         gr.Markdown(value="## 3. Generate your dataset")
@@ -729,11 +654,20 @@ with gr.Blocks() as app:
                 temperature = gr.Slider(
                     label="Temperature",
                     minimum=0.1,
-                    maximum=1,
+                    maximum=1.5,
                     value=0.7,
                     step=0.1,
                     interactive=True,
                 )
+                temperature_completion = gr.Slider(
+                    label="Temperature for completion",
+                    minimum=0.1,
+                    maximum=1.5,
+                    value=None,
+                    step=0.1,
+                    interactive=True,
+                    visible=False,
+                )
                 private = gr.Checkbox(
                     label="Private dataset",
                     value=False,
@@ -753,7 +687,6 @@ with gr.Blocks() as app:
                 ) as pipeline_code_ui:
                     code = generate_pipeline_code(
                         repo_id=search_in.value,
-                        file_paths=file_in.value,
                         input_type=input_type.value,
                         system_prompt=system_prompt.value,
                         document_column=document_column.value,
@@ -790,35 +723,23 @@ with gr.Blocks() as app:
         fn=hide_document_column_visibility, inputs=[], outputs=[document_column]
     )
 
-    search_in.submit(fn=get_iframe, inputs=search_in, outputs=search_out).then(
+    search_in.submit(
         fn=lambda df: pd.DataFrame(columns=df.columns),
         inputs=[dataframe],
         outputs=[dataframe],
     )
 
-    load_dataset_btn.click(
+    gr.on(
+        triggers=[load_dataset_btn.click, load_file_btn.click],
         fn=load_dataset_file,
         inputs=[search_in, file_in, input_type],
-        outputs=[
-            dataframe,
-            document_column,
-        ],
-    )
-
-    load_file_btn.click(
-        fn=load_dataset_file,
-        inputs=[search_in, file_in, input_type],
-        outputs=[
-            dataframe,
-            document_column,
-        ],
+        outputs=[dataframe, document_column],
     )
 
     load_prompt_btn.click(
         fn=generate_system_prompt,
         inputs=[dataset_description],
         outputs=[system_prompt],
-        show_progress=True,
     ).success(
         fn=generate_sample_dataset,
         inputs=[
@@ -851,16 +772,13 @@ with gr.Blocks() as app:
         fn=validate_argilla_user_workspace_dataset,
         inputs=[repo_name],
         outputs=[success_message],
-        show_progress=True,
     ).then(
         fn=validate_push_to_hub,
         inputs=[org_name, repo_name],
         outputs=[success_message],
-        show_progress=True,
     ).success(
         fn=hide_success_message,
         outputs=[success_message],
-        show_progress=True,
     ).success(
         fn=hide_pipeline_code_visibility,
         inputs=[],
@@ -879,10 +797,10 @@ with gr.Blocks() as app:
             retrieval_reranking,
             num_rows,
             temperature,
+            temperature_completion,
             pipeline_code,
         ],
         outputs=[success_message],
-        show_progress=True,
     ).success(
         fn=show_success_message,
         inputs=[org_name, repo_name],
@@ -891,7 +809,6 @@ with gr.Blocks() as app:
         fn=generate_pipeline_code,
         inputs=[
             search_in,
-            file_in,
             input_type,
             system_prompt,
             document_column,
@@ -905,11 +822,9 @@ with gr.Blocks() as app:
         outputs=[pipeline_code_ui],
     )
 
-    clear_dataset_btn_part.click(fn=lambda : "", inputs=[], outputs=[search_in])
+    clear_dataset_btn_part.click(fn=lambda: "", inputs=[], outputs=[search_in])
     clear_file_btn_part.click(fn=lambda: None, inputs=[], outputs=[file_in])
-    clear_prompt_btn_part.click(
-        fn=lambda : "", inputs=[], outputs=[dataset_description]
-    )
+    clear_prompt_btn_part.click(fn=lambda: "", inputs=[], outputs=[dataset_description])
     clear_btn_full.click(
         fn=lambda df: ("", [], pd.DataFrame(columns=df.columns)),
         inputs=[dataframe],
@@ -919,3 +834,4 @@ with gr.Blocks() as app:
     app.load(fn=swap_visibility, outputs=main_ui)
     app.load(fn=get_org_dropdown, outputs=[org_name])
     app.load(fn=get_random_repo_name, outputs=[repo_name])
+    app.load(fn=show_temperature_completion, outputs=[temperature_completion])

src/synthetic_dataset_generator/apps/textcat.py

@@ -49,7 +49,7 @@ def _get_dataframe():
     )
 
 
-def generate_system_prompt(dataset_description, progress=gr.Progress()):
+def generate_system_prompt(dataset_description: str, progress=gr.Progress()):
     progress(0.0, desc="Starting")
     progress(0.3, desc="Initializing")
     generate_description = get_prompt_generator()
@@ -71,7 +71,12 @@ def generate_system_prompt(dataset_description, progress=gr.Progress()):
 
 
 def generate_sample_dataset(
-    system_prompt, difficulty, clarity, labels, multi_label, progress=gr.Progress()
+    system_prompt: str,
+    difficulty: str,
+    clarity: str,
+    labels: List[str],
+    multi_label: bool,
+    progress=gr.Progress(),
 ):
     dataframe = generate_dataset(
         system_prompt=system_prompt,
@@ -294,14 +299,14 @@ def push_dataset(
         temperature=temperature,
     )
     push_dataset_to_hub(
-        dataframe,
-        org_name,
-        repo_name,
-        multi_label,
-        labels,
-        oauth_token,
-        private,
-        pipeline_code,
+        dataframe=dataframe,
+        org_name=org_name,
+        repo_name=repo_name,
+        multi_label=multi_label,
+        labels=labels,
+        oauth_token=oauth_token,
+        private=private,
+        pipeline_code=pipeline_code,
     )
 
     dataframe = dataframe[
@@ -453,62 +458,59 @@ with gr.Blocks() as app:
 
         gr.HTML("<hr>")
         gr.Markdown("## 2. Configure your dataset")
-        with gr.Row(equal_height=True):
-            with gr.Row(equal_height=False):
-                with gr.Column(scale=2):
-                    system_prompt = gr.Textbox(
-                        label="System prompt",
-                        placeholder="You are a helpful assistant.",
-                        visible=True,
-                    )
-                    labels = gr.Dropdown(
-                        choices=[],
-                        allow_custom_value=True,
-                        interactive=True,
-                        label="Labels",
-                        multiselect=True,
-                        info="Add the labels to classify the text.",
-                    )
-                    multi_label = gr.Checkbox(
-                        label="Multi-label",
-                        value=False,
-                        interactive=True,
-                        info="If checked, the text will be classified into multiple labels.",
-                    )
-                    clarity = gr.Dropdown(
-                        choices=[
-                            ("Clear", "clear"),
-                            (
-                                "Understandable",
-                                "understandable with some effort",
-                            ),
-                            ("Ambiguous", "ambiguous"),
-                            ("Mixed", "mixed"),
-                        ],
-                        value="mixed",
-                        label="Clarity",
-                        info="Set how easily the correct label or labels can be identified.",
-                        interactive=True,
-                    )
-                    difficulty = gr.Dropdown(
-                        choices=[
-                            ("High School", "high school"),
-                            ("College", "college"),
-                            ("PhD", "PhD"),
-                            ("Mixed", "mixed"),
-                        ],
-                        value="high school",
-                        label="Difficulty",
-                        info="Select the comprehension level for the text. Ensure it matches the task context.",
-                        interactive=True,
-                    )
-                    with gr.Row():
-                        clear_btn_full = gr.Button("Clear", variant="secondary")
-                        btn_apply_to_sample_dataset = gr.Button(
-                            "Save", variant="primary"
-                        )
-                with gr.Column(scale=3):
-                    dataframe = _get_dataframe()
+        with gr.Row(equal_height=False):
+            with gr.Column(scale=2):
+                system_prompt = gr.Textbox(
+                    label="System prompt",
+                    placeholder="You are a helpful assistant.",
+                    visible=True,
+                )
+                labels = gr.Dropdown(
+                    choices=[],
+                    allow_custom_value=True,
+                    interactive=True,
+                    label="Labels",
+                    multiselect=True,
+                    info="Add the labels to classify the text.",
+                )
+                multi_label = gr.Checkbox(
+                    label="Multi-label",
+                    value=False,
+                    interactive=True,
+                    info="If checked, the text will be classified into multiple labels.",
+                )
+                clarity = gr.Dropdown(
+                    choices=[
+                        ("Clear", "clear"),
+                        (
+                            "Understandable",
+                            "understandable with some effort",
+                        ),
+                        ("Ambiguous", "ambiguous"),
+                        ("Mixed", "mixed"),
+                    ],
+                    value="mixed",
+                    label="Clarity",
+                    info="Set how easily the correct label or labels can be identified.",
+                    interactive=True,
+                )
+                difficulty = gr.Dropdown(
+                    choices=[
+                        ("High School", "high school"),
+                        ("College", "college"),
+                        ("PhD", "PhD"),
+                        ("Mixed", "mixed"),
+                    ],
+                    value="high school",
+                    label="Difficulty",
+                    info="Select the comprehension level for the text. Ensure it matches the task context.",
+                    interactive=True,
+                )
+                with gr.Row():
+                    clear_btn_full = gr.Button("Clear", variant="secondary")
+                    btn_apply_to_sample_dataset = gr.Button("Save", variant="primary")
+            with gr.Column(scale=3):
+                dataframe = _get_dataframe()
 
         gr.HTML("<hr>")
         gr.Markdown("## 3. Generate your dataset")
@@ -530,7 +532,7 @@ with gr.Blocks() as app:
                 temperature = gr.Slider(
                     label="Temperature",
                     minimum=0.1,
-                    maximum=1,
+                    maximum=1.5,
                     value=0.8,
                     step=0.1,
                     interactive=True,
@@ -570,45 +572,37 @@ with gr.Blocks() as app:
         fn=generate_system_prompt,
         inputs=[dataset_description],
         outputs=[system_prompt, labels],
-        show_progress=True,
     ).then(
         fn=generate_sample_dataset,
         inputs=[system_prompt, difficulty, clarity, labels, multi_label],
         outputs=[dataframe],
-        show_progress=True,
     )
 
     btn_apply_to_sample_dataset.click(
         fn=validate_input_labels,
         inputs=[labels],
         outputs=[labels],
-        show_progress=True,
     ).success(
         fn=generate_sample_dataset,
         inputs=[system_prompt, difficulty, clarity, labels, multi_label],
         outputs=[dataframe],
-        show_progress=True,
     )
 
     btn_push_to_hub.click(
         fn=validate_argilla_user_workspace_dataset,
         inputs=[repo_name],
         outputs=[success_message],
-        show_progress=True,
     ).then(
         fn=validate_push_to_hub,
         inputs=[org_name, repo_name],
         outputs=[success_message],
-        show_progress=True,
     ).success(
         fn=validate_input_labels,
         inputs=[labels],
         outputs=[labels],
-        show_progress=True,
     ).success(
         fn=hide_success_message,
         outputs=[success_message],
-        show_progress=True,
     ).success(
         fn=hide_pipeline_code_visibility,
         inputs=[],
@@ -629,7 +623,6 @@ with gr.Blocks() as app:
             pipeline_code,
         ],
         outputs=[success_message],
-        show_progress=True,
     ).success(
         fn=show_success_message,
         inputs=[org_name, repo_name],
@@ -657,6 +650,7 @@ with gr.Blocks() as app:
             "",
             "",
             [],
+            "",
             _get_dataframe(),
         ),
         inputs=[dataframe],

src/synthetic_dataset_generator/constants.py

@@ -3,10 +3,6 @@ import warnings
 
 import argilla as rg
 
-# Tasks
-TEXTCAT_TASK = "text_classification"
-SFT_TASK = "supervised_fine_tuning"
-
 # Inference
 MAX_NUM_TOKENS = int(os.getenv("MAX_NUM_TOKENS", 2048))
 MAX_NUM_ROWS = int(os.getenv("MAX_NUM_ROWS", 1000))
@@ -20,28 +16,56 @@ OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL")
 HUGGINGFACE_BASE_URL = os.getenv("HUGGINGFACE_BASE_URL")
 VLLM_BASE_URL = os.getenv("VLLM_BASE_URL")
 
-# check if model is set correctly
-if HUGGINGFACE_BASE_URL and MODEL:
-    raise ValueError(
-        "`HUGGINGFACE_BASE_URL` and `MODEL` cannot be set at the same time. Use a model id for serverless inference and a base URL dedicated to Hugging Face Inference Endpoints."
-    )
-if not MODEL:
-    if OPENAI_BASE_URL or OLLAMA_BASE_URL or VLLM_BASE_URL:
-        raise ValueError("`MODEL` is not set. Please provide a model id for inference.")
-
-# Check if multiple base URLs are provided
-base_urls = [
-    url
-    for url in [OPENAI_BASE_URL, OLLAMA_BASE_URL, HUGGINGFACE_BASE_URL, VLLM_BASE_URL]
-    if url
+# Just used in case of selecting a different model for completions
+MODEL_COMPLETION = os.getenv("MODEL_COMPLETION", MODEL)
+TOKENIZER_ID_COMPLETION = os.getenv("TOKENIZER_ID_COMPLETION", TOKENIZER_ID)
+OPENAI_BASE_URL_COMPLETION = os.getenv("OPENAI_BASE_URL_COMPLETION", OPENAI_BASE_URL)
+OLLAMA_BASE_URL_COMPLETION = os.getenv("OLLAMA_BASE_URL_COMPLETION", OLLAMA_BASE_URL)
+HUGGINGFACE_BASE_URL_COMPLETION = os.getenv(
+    "HUGGINGFACE_BASE_URL_COMPLETION", HUGGINGFACE_BASE_URL
+)
+VLLM_BASE_URL_COMPLETION = os.getenv("VLLM_BASE_URL_COMPLETION", VLLM_BASE_URL)
+
+base_urls = [OPENAI_BASE_URL, OLLAMA_BASE_URL, HUGGINGFACE_BASE_URL, VLLM_BASE_URL]
+base_urls_completion = [
+    OPENAI_BASE_URL_COMPLETION,
+    OLLAMA_BASE_URL_COMPLETION,
+    HUGGINGFACE_BASE_URL_COMPLETION,
+    VLLM_BASE_URL_COMPLETION,
 ]
-if len(base_urls) > 1:
-    raise ValueError(
-        f"Multiple base URLs provided: {', '.join(base_urls)}. Only one base URL can be set at a time."
-    )
-BASE_URL = OPENAI_BASE_URL or OLLAMA_BASE_URL or HUGGINGFACE_BASE_URL or VLLM_BASE_URL
 
 
+# Validate the configuration of the model and base URLs.
+def validate_configuration(base_urls, model, env_context=""):
+    huggingface_url = base_urls[2]
+    if huggingface_url and model:
+        raise ValueError(
+            f"`HUGGINGFACE_BASE_URL{env_context}` and `MODEL{env_context}` cannot be set at the same time. "
+            "Use a model id for serverless inference and a base URL dedicated to Hugging Face Inference Endpoints."
+        )
+
+    if not model and any(base_urls):
+        raise ValueError(
+            f"`MODEL{env_context}` is not set. Please provide a model id for inference."
+        )
+
+    active_urls = [url for url in base_urls if url]
+    if len(active_urls) > 1:
+        raise ValueError(
+            f"Multiple base URLs are provided: {', '.join(active_urls)}. "
+            "Only one base URL can be set at a time."
+        )
+validate_configuration(base_urls, MODEL)
+validate_configuration(base_urls_completion, MODEL_COMPLETION, "_COMPLETION")
+
+BASE_URL = OPENAI_BASE_URL or OLLAMA_BASE_URL or HUGGINGFACE_BASE_URL or VLLM_BASE_URL
+BASE_URL_COMPLETION = (
+    OPENAI_BASE_URL_COMPLETION
+    or OLLAMA_BASE_URL_COMPLETION
+    or HUGGINGFACE_BASE_URL_COMPLETION
+    or VLLM_BASE_URL_COMPLETION
+)
+
 # API Keys
 HF_TOKEN = os.getenv("HF_TOKEN")
 if not HF_TOKEN:

src/synthetic_dataset_generator/pipelines/base.py

@@ -8,11 +8,17 @@ from synthetic_dataset_generator.constants import (
     API_KEYS,
     DEFAULT_BATCH_SIZE,
     HUGGINGFACE_BASE_URL,
+    HUGGINGFACE_BASE_URL_COMPLETION,
     MODEL,
+    MODEL_COMPLETION,
     OLLAMA_BASE_URL,
+    OLLAMA_BASE_URL_COMPLETION,
     OPENAI_BASE_URL,
+    OPENAI_BASE_URL_COMPLETION,
     TOKENIZER_ID,
+    TOKENIZER_ID_COMPLETION,
     VLLM_BASE_URL,
+    VLLM_BASE_URL_COMPLETION,
 )
 
 TOKEN_INDEX = 0
@@ -73,12 +79,20 @@ def _get_llm_class() -> str:
         return "InferenceEndpointsLLM"
 
 
-def _get_llm(use_magpie_template=False, **kwargs):
+def _get_llm(
+    structured_output: dict = None,
+    use_magpie_template: str = False,
+    is_completion: bool = False,
+    **kwargs,
+):
+    model = MODEL_COMPLETION if is_completion else MODEL
+    tokenizer_id = TOKENIZER_ID_COMPLETION if is_completion else TOKENIZER_ID or model
     if OPENAI_BASE_URL:
         llm = OpenAILLM(
-            model=MODEL,
-            base_url=OPENAI_BASE_URL,
+            model=model,
+            base_url=OPENAI_BASE_URL_COMPLETION if is_completion else OPENAI_BASE_URL,
             api_key=_get_next_api_key(),
+            structured_output=structured_output,
             **kwargs,
         )
         if "generation_kwargs" in kwargs:
@@ -108,19 +122,25 @@ def _get_llm(use_magpie_template=False, **kwargs):
             kwargs["generation_kwargs"] = {}
             kwargs["generation_kwargs"]["options"] = options
         llm = OllamaLLM(
-            model=MODEL,
-            host=OLLAMA_BASE_URL,
-            tokenizer_id=TOKENIZER_ID or MODEL,
+            model=model,
+            host=OLLAMA_BASE_URL_COMPLETION if is_completion else OLLAMA_BASE_URL,
+            tokenizer_id=tokenizer_id,
             use_magpie_template=use_magpie_template,
+            structured_output=structured_output,
             **kwargs,
         )
     elif HUGGINGFACE_BASE_URL:
         kwargs["generation_kwargs"]["do_sample"] = True
         llm = InferenceEndpointsLLM(
             api_key=_get_next_api_key(),
-            base_url=HUGGINGFACE_BASE_URL,
-            tokenizer_id=TOKENIZER_ID or MODEL,
+            base_url=(
+                HUGGINGFACE_BASE_URL_COMPLETION
+                if is_completion
+                else HUGGINGFACE_BASE_URL
+            ),
+            tokenizer_id=tokenizer_id,
             use_magpie_template=use_magpie_template,
+            structured_output=structured_output,
             **kwargs,
         )
     elif VLLM_BASE_URL:
@@ -128,19 +148,21 @@ def _get_llm(use_magpie_template=False, **kwargs):
             if "do_sample" in kwargs["generation_kwargs"]:
                 del kwargs["generation_kwargs"]["do_sample"]
         llm = ClientvLLM(
-            base_url=VLLM_BASE_URL,
-            model=MODEL,
-            tokenizer=TOKENIZER_ID or MODEL,
+            base_url=VLLM_BASE_URL_COMPLETION if is_completion else VLLM_BASE_URL,
+            model=model,
+            tokenizer=tokenizer_id,
             api_key=_get_next_api_key(),
             use_magpie_template=use_magpie_template,
+            structured_output=structured_output,
             **kwargs,
         )
     else:
         llm = InferenceEndpointsLLM(
             api_key=_get_next_api_key(),
-            tokenizer_id=TOKENIZER_ID or MODEL,
-            model_id=MODEL,
+            tokenizer_id=tokenizer_id,
+            model_id=model,
             use_magpie_template=use_magpie_template,
+            structured_output=structured_output,
             **kwargs,
         )
 

src/synthetic_dataset_generator/pipelines/chat.py

@@ -1,4 +1,10 @@
-from distilabel.steps.tasks import ChatGeneration, Magpie, TextGeneration
+from datasets import get_dataset_config_names, get_dataset_split_names
+from distilabel.steps.tasks import (
+    ChatGeneration,
+    Magpie,
+    GenerateSentencePair,
+    TextGeneration,
+)
 
 from synthetic_dataset_generator.constants import (
     MAGPIE_PRE_QUERY_TEMPLATE,
@@ -118,6 +124,18 @@ The prompt you write should follow the same style and structure as the following
 User dataset description:
 """
 
+FOLLOW_UP_TEMPLATE = """Conversation:
+{% for message in messages %}
+    {% if message.role == "user" %}
+User Question: {{ message.content }}
+    {% elif message.role == "assistant" %}
+Assistant Response: {{ message.content }}
+    {% endif %}
+{% endfor %}
+
+Please generate the next logical user message in this conversation. Do not include any other information or 'User Question' in your response.
+""".rstrip()
+
 DEFAULT_DATASET_DESCRIPTIONS = [
     "rude customer assistant for a phone company",
     "assistant that solves math puzzles using python",
@@ -140,7 +158,7 @@ else:
     ]
 
 
-def _get_output_mappings(num_turns):
+def _get_output_mappings(num_turns: int):
     if num_turns == 1:
         return {"instruction": "prompt", "response": "completion"}
     else:
@@ -162,7 +180,7 @@ def get_prompt_generator():
     return prompt_generator
 
 
-def get_magpie_generator(system_prompt, num_turns, temperature, is_sample):
+def get_magpie_generator(num_turns: int, temperature: float, is_sample: bool):
     input_mappings = _get_output_mappings(num_turns)
     output_mappings = input_mappings.copy()
     if num_turns == 1:
@@ -203,14 +221,31 @@ def get_magpie_generator(system_prompt, num_turns, temperature, is_sample):
     return magpie_generator
 
 
-def get_response_generator(system_prompt, num_turns, temperature, is_sample):
+def get_sentence_pair_generator(temperature: float, is_sample: bool):
+    generation_kwargs = {
+        "temperature": temperature,
+        "max_new_tokens": 256 if is_sample else MAX_NUM_TOKENS,
+    }
+    sentence_pair_generator = GenerateSentencePair(
+        llm=_get_llm(generation_kwargs=generation_kwargs),
+        triplet=False,
+        action="query",
+        hard_negative=True,
+    )
+    sentence_pair_generator.load()
+    return sentence_pair_generator
+
+
+def get_response_generator(
+    system_prompt: str, num_turns: int, temperature: float, is_sample: bool
+):
     if num_turns == 1:
         generation_kwargs = {
             "temperature": temperature,
             "max_new_tokens": 256 if is_sample else int(MAX_NUM_TOKENS * 0.5),
         }
         response_generator = TextGeneration(
-            llm=_get_llm(generation_kwargs=generation_kwargs),
+            llm=_get_llm(is_completion=True, generation_kwargs=generation_kwargs),
             system_prompt=system_prompt,
             output_mappings={"generation": "completion"},
             input_mappings={"instruction": "prompt"},
@@ -221,7 +256,7 @@ def get_response_generator(system_prompt, num_turns, temperature, is_sample):
             "max_new_tokens": MAX_NUM_TOKENS,
         }
         response_generator = ChatGeneration(
-            llm=_get_llm(generation_kwargs=generation_kwargs),
+            llm=_get_llm(is_completion=True, generation_kwargs=generation_kwargs),
             output_mappings={"generation": "completion"},
             input_mappings={"conversation": "messages"},
         )
@@ -229,36 +264,236 @@ def get_response_generator(system_prompt, num_turns, temperature, is_sample):
     return response_generator
 
 
-def generate_pipeline_code(system_prompt, num_turns, num_rows):
+def get_follow_up_generator(type: str, temperature: float, is_sample: bool):
+    if type == "instruction":
+        generation_kwargs = {
+            "temperature": temperature,
+            "max_new_tokens": 256 if is_sample else int(MAX_NUM_TOKENS * 0.5),
+        }
+        follow_up_generator = TextGeneration(
+            llm=_get_llm(generation_kwargs=generation_kwargs),
+            template=FOLLOW_UP_TEMPLATE,
+            columns=["messages"],
+        )
+    else:
+        generation_kwargs = {
+            "temperature": temperature,
+            "max_new_tokens": MAX_NUM_TOKENS,
+        }
+        follow_up_generator = ChatGeneration(
+            llm=_get_llm(is_completion=True, generation_kwargs=generation_kwargs),
+        )
+    follow_up_generator.load()
+    return follow_up_generator
+
+def generate_pipeline_code_system_prompt(
+    system_prompt: str,
+    num_turns: int,
+    num_rows: int,
+):
     input_mappings = _get_output_mappings(num_turns)
+    code = f"""
+    # Requirements: `pip install distilabel[hf-inference-endpoints]`
+    import os
+    from distilabel.pipeline import Pipeline
+    from distilabel.steps import KeepColumns
+    from distilabel.steps.tasks import MagpieGenerator
+    from distilabel.llms import {_get_llm_class()}
+
+    SYSTEM_PROMPT = "{system_prompt}"
+
+    with Pipeline(name="sft") as pipeline:
+        magpie = MagpieGenerator(
+            llm={_get_llm_class()}.from_dict(
+                {_get_llm().dump()}
+            ),
+            n_turns={num_turns},
+            num_rows={num_rows},
+            batch_size=1,
+            system_prompt=SYSTEM_PROMPT,
+            output_mappings={input_mappings},
+        )
+        keep_columns = KeepColumns(
+            columns={list(input_mappings.values())} + ["model_name"],
+        )
+        magpie.connect(keep_columns)
 
+    if __name__ == "__main__":
+        distiset = pipeline.run()
+    """
+    return code
+
+def generate_pipeline_code_seed(
+    repo_id: str,
+    subset: str,
+    split: str,
+    input_type: str,
+    document_column: str,
+    num_turns: int,
+    num_rows: int,
+):
     code = f"""
 # Requirements: `pip install distilabel[hf-inference-endpoints]`
-import os
+from distilabel.models import {_get_llm_class()}
 from distilabel.pipeline import Pipeline
-from distilabel.steps import KeepColumns
-from distilabel.steps.tasks import MagpieGenerator
-from distilabel.llms import {_get_llm_class()}
+from distilabel.steps import KeepColumns{", LoadDataFromDicts" if input_type != "dataset-input"  else ""}{", LoadDataFromHub" if input_type == "dataset-input" else ""}{", StepInput, step" if num_turns > 1 else ""}
+from distilabel.steps.tasks import GenerateSentencePair, TextGeneration {", ChatGeneration" if num_turns > 1 else ""}
+"""
 
-SYSTEM_PROMPT = "{system_prompt}"
+    if num_turns > 1:
+        code += """
+FOLLOW_UP_TEMPLATE = '''Conversation:
+{{% for message in messages %}}
+    {{% if message.role == "user" %}}
+User Question: {{{{ message.content }}}}
+    {{% elif message.role == "assistant" %}}
+Assistant Response: {{{{ message.content }}}}
+    {{% endif %}}
+{{% endfor %}}
+
+Please generate the next logical user message in this conversation. Do not include any other information or 'User Question' in your response.
+'''.rstrip()
+
+@step(inputs=["prompt", "completion"], outputs=["messages"])
+def PrepareMessages(*inputs: StepInput) -> StepOutput:
+    for input in inputs:
+        for item in input:
+            item["messages"] = [
+                {"role": "user", "content": item["prompt"]},
+                {"role": "assistant", "content": item["completion"]},
+            ]
+        yield input
+
+
+@step(inputs=["messages", "generation"], outputs=["messages"])
+def FormatMessagesInstruction(*inputs: StepInput) -> StepOutput:
+    for input in inputs:
+        for item in input:
+            item["messages"].append({"role": "user", "content": item["generation"]})
+        yield input
+
+
+@step(inputs=["messages", "generation"], outputs=["messages"])
+def FormatMessagesResponse(*inputs: StepInput) -> StepOutput:
+    for input in inputs:
+        for item in input:
+            item["messages"].append({"role": "assistant", "content": item["generation"]})
+        yield input
+"""
+
+    if input_type == "dataset-input":
+        code += f"""
+with Pipeline(name="sft") as pipeline:
+    load_the_dataset = LoadDataFromHub(
+        repo_id='{repo_id}',
+        config='{subset}',
+        split='{split}',
+        num_examples={num_rows},
+        batch_size=2,
+        output_mappings={{'{document_column}':'anchor'}},
+    )
+    """
+
+    else: 
+        code += """
+data = process_and_chunk_files(files=[files])
 
 with Pipeline(name="sft") as pipeline:
-    magpie = MagpieGenerator(
+    load_the_dataset = LoadDataFromDicts(
+        data = data
+    )
+"""
+    code += f"""
+    instruction_generator = GenerateSentencePair(
+        name="instruction_generation",
+        triplet=False,
+        hard_negative=True,
+        action="query",
+        llm={_get_llm_class()}.from_dict(
+            {_get_llm().dump()}
+        ),
+        input_batch_size=10,
+        output_mappings={{"positive": "prompt"}},
+    )
+
+    response_generator = TextGeneration(
+        name="response_generation",
+        llm={_get_llm_class()}.from_dict(
+            {_get_llm().dump()}
+        ),
+        input_batch_size=10,
+        input_mappings={{"instruction": "prompt"}},
+        output_mappings={{"generation": "completion"}},
+    )
+    """
+
+    if num_turns > 1:
+        code += """
+    prepare_messages = PrepareMessages()
+    """
+
+        for i in range(num_turns - 1):
+            code += f"""
+    follow_up_instruction_{i} = TextGeneration(
         llm={_get_llm_class()}.from_dict(
             {_get_llm().dump()}
         ),
-        n_turns={num_turns},
-        num_rows={num_rows},
-        batch_size=1,
-        system_prompt=SYSTEM_PROMPT,
-        output_mappings={input_mappings},
+        template=FOLLOW_UP_TEMPLATE,
+        columns=["messages"],
     )
-    keep_columns = KeepColumns(
-        columns={list(input_mappings.values())} + ["model_name"],
+    format_instruction_{i} = FormatMessagesInstruction()
+    follow_up_response_{i} = ChatGeneration(
+        llm={_get_llm_class()}.from_dict(
+            {_get_llm().dump()}
+        ),
     )
-    magpie.connect(keep_columns)
+    format_response_{i} = FormatMessagesResponse()
+    """
+
+    if num_turns > 1:
+        code += """
+        keep_columns = KeepColumns(columns=["messages"])
+        """
+        code += "load_the_dataset >> instruction_generator >> response_generator >> prepare_messages"
+
+        for i in range(1, num_turns + 1):
+            code += f" >> follow_up_instruction_{i} >> format_instruction_{i} >> follow_up_response_{i} >> format_response_{i}"
 
+        code += " >> keep_columns"
+
+    code += """
 if __name__ == "__main__":
     distiset = pipeline.run()
+)
 """
     return code
+
+def generate_pipeline_code(
+    repo_id: str,
+    input_type: str,
+    system_prompt: str,
+    document_column: str,
+    num_turns: int,
+    num_rows: int,
+):
+    if input_type == "dataset-input" and repo_id is not None:
+        subset = get_dataset_config_names(repo_id)[0]
+        split = get_dataset_split_names(repo_id, subset)[0]
+    else:
+        subset = "default"
+        split = "train"
+    if input_type == "prompt-type":
+        return generate_pipeline_code_system_prompt(
+            system_prompt=system_prompt,
+            num_turns=num_turns,
+            num_rows=num_rows,
+        )
+    return generate_pipeline_code_seed(
+        repo_id=repo_id,
+        subset=subset,
+        split=split,
+        input_type=input_type,
+        document_column=document_column,
+        num_turns=num_turns,
+        num_rows=num_rows,
+    )

src/synthetic_dataset_generator/pipelines/eval.py

@@ -1,3 +1,5 @@
+from typing import List
+
 from datasets import get_dataset_config_names, get_dataset_split_names
 from distilabel.models import InferenceEndpointsLLM
 from distilabel.steps.tasks import (
@@ -10,7 +12,7 @@ from synthetic_dataset_generator.pipelines.base import _get_next_api_key
 from synthetic_dataset_generator.utils import extract_column_names
 
 
-def get_ultrafeedback_evaluator(aspect, is_sample):
+def get_ultrafeedback_evaluator(aspect: str, is_sample: bool):
     ultrafeedback_evaluator = UltraFeedback(
         llm=InferenceEndpointsLLM(
             model_id=MODEL,
@@ -27,7 +29,9 @@ def get_ultrafeedback_evaluator(aspect, is_sample):
     return ultrafeedback_evaluator
 
 
-def get_custom_evaluator(prompt_template, structured_output, columns, is_sample):
+def get_custom_evaluator(
+    prompt_template: str, structured_output: dict, columns: List[str], is_sample: bool
+):
     custom_evaluator = TextGeneration(
         llm=InferenceEndpointsLLM(
             model_id=MODEL,
@@ -47,7 +51,13 @@ def get_custom_evaluator(prompt_template, structured_output, columns, is_sample)
 
 
 def generate_ultrafeedback_pipeline_code(
-    repo_id, subset, split, aspects, instruction_column, response_columns, num_rows
+    repo_id: str,
+    subset: str,
+    split: str,
+    aspects: List[str],
+    instruction_column: str,
+    response_columns: str,
+    num_rows: int,
 ):
     if len(aspects) == 1:
         code = f"""

src/synthetic_dataset_generator/pipelines/rag.py

@@ -1,7 +1,3 @@
-import os
-
-from typing import List
-
 from datasets import get_dataset_config_names, get_dataset_split_names
 from distilabel.steps.tasks import (
     GenerateSentencePair,
@@ -87,7 +83,7 @@ def get_prompt_generator():
     return text_generator
 
 
-def get_chunks_generator(temperature, is_sample):
+def get_chunks_generator(temperature: float, is_sample: bool):
     generation_kwargs = {
         "temperature": temperature,
         "max_new_tokens": MAX_NUM_TOKENS if is_sample else 256,
@@ -104,7 +100,7 @@ def get_chunks_generator(temperature, is_sample):
     return text_generator
 
 
-def get_sentence_pair_generator(action, triplet, temperature, is_sample):
+def get_sentence_pair_generator(action: str, triplet: bool, temperature: float, is_sample: bool):
     generation_kwargs = {
         "temperature": temperature,
         "max_new_tokens": 256 if is_sample else MAX_NUM_TOKENS,
@@ -119,13 +115,13 @@ def get_sentence_pair_generator(action, triplet, temperature, is_sample):
     return sentence_pair_generator
 
 
-def get_response_generator(temperature, is_sample):
+def get_response_generator(temperature: float, is_sample: bool):
     generation_kwargs = {
         "temperature": temperature,
         "max_new_tokens": MAX_NUM_TOKENS if is_sample else 256,
     }
     text_generator = TextGeneration(
-        llm=_get_llm(generation_kwargs=generation_kwargs),
+        llm=_get_llm(is_completion=True, generation_kwargs=generation_kwargs),
         system_prompt=SYSTEM_PROMPT_RAG,
         template=RAG_TEMPLATE,
         columns=["context", "question"],
@@ -138,7 +134,6 @@ def get_response_generator(temperature, is_sample):
 
 def generate_pipeline_code(
     repo_id: str,
-    file_paths: List[str],
     input_type: str,
     system_prompt: str,
     document_column: str,
@@ -293,10 +288,7 @@ with Pipeline(name="rag") as pipeline:
 
     pipeline += """
     if __name__ == "__main__":
-        distiset = pipeline.run(use_cache=False)
-        print(distiset)
-        if distiset:
-            print(distiset["default"]["train"][0])
+        distiset = pipeline.run()
     """
 
     return base_code + pipeline

src/synthetic_dataset_generator/pipelines/textcat.py

@@ -85,7 +85,9 @@ def get_prompt_generator():
     return prompt_generator
 
 
-def get_textcat_generator(difficulty, clarity, temperature, is_sample):
+def get_textcat_generator(
+    difficulty: str, clarity: str, temperature: float, is_sample: bool
+):
     generation_kwargs = {
         "temperature": temperature,
         "max_new_tokens": 256 if is_sample else MAX_NUM_TOKENS,
@@ -102,12 +104,12 @@ def get_textcat_generator(difficulty, clarity, temperature, is_sample):
     return textcat_generator
 
 
-def get_labeller_generator(system_prompt, labels, multi_label):
+def get_labeller_generator(system_prompt: str, labels: List[str], multi_label: bool):
     generation_kwargs = {
         "temperature": 0.01,
         "max_new_tokens": MAX_NUM_TOKENS,
     }
-    llm = _get_llm(generation_kwargs=generation_kwargs)
+    llm = _get_llm(is_completion=True, generation_kwargs=generation_kwargs)
     labeller_generator = TextClassification(
         llm=llm,
         context=system_prompt,