Add RAG generation (#19)

* feat: add pipelines for RAG

* feat: add the app for rag

* feat: update required dependencies

* feat: update add rag tab

* fix: correct unstructured dependency

* fix: fix layout, arg and input errors

* fix: rename generator

* feat: generate without ground data

* organize dependencies and add pydantic version

* update with feedback

* fix errors and update with latest changes

* remove temperature as value ini pipeline code

* fix pipeline errors

* update TODOs and review prompt generator

* Update READMES

* remove unneeded imports

* use llm.dump() instead of llm.model_dump()

* fix pipeline code

* apply feedback

* fix progress, messages and clear buttons

README.md

@@ -34,6 +34,7 @@ Supported Tasks:
 
 - Text Classification
 - Chat Data for Supervised Fine-Tuning
+- Retrieval Augmented Generation
 
 This tool simplifies the process of creating custom datasets, enabling you to:
 

pyproject.toml

@@ -18,13 +18,15 @@ readme = "README.md"
 license = {text = "Apache 2"}
 
 dependencies = [
+    "argilla>=2.4.0,<3.0.0",
     "distilabel[argilla,hf-inference-endpoints,hf-transformers,instructor,llama-cpp,ollama,openai,outlines,vllm] @ git+https://github.com/argilla-io/distilabel.git@develop",
     "gradio[oauth]>=5.4.0,<6.0.0",
-    "transformers>=4.44.2,<5.0.0",
-    "sentence-transformers>=3.2.0,<4.0.0",
-    "model2vec>=0.2.4,<1.0.0",
     "gradio-huggingfacehub-search>=0.0.12,<1.0.0",
-    "argilla>=2.4.0,<3.0.0",
+    "model2vec>=0.2.4,<1.0.0",
+    "pydantic>=2.10.5,<3.0.0",
+    "sentence-transformers>=3.2.0,<4.0.0",
+    "transformers>=4.44.2,<5.0.0",
+    "unstructured[md,pdf,docx]>=0.16.0,<1.0.0",
 ]
 
 [build-system]

src/synthetic_dataset_generator/_distiset.py

@@ -92,7 +92,22 @@ class CustomDistisetWithAdditionalTag(distilabel.distiset.Distiset):
         elif ("prompt" in columns and "completion" in columns) or (
             "messages" in columns
         ):
-            task_categories: list[str] = ["text-generation", "text2text-generation"]
+            task_categories: list[str] = [
+                "text-generation",
+                "text2text-generation",
+                "question-answering",
+            ]
+        elif "context" in columns and "question" in columns and "response" in columns:
+            task_categories: list[str] = [
+                "text-generation",
+                "text2text-generation",
+                "text-retrieval",
+                "question-answering"
+            ]
+            if (
+                "positive_retrieval" in columns and "negative_retrieval" in columns
+            ) or ("positive_reranking" in columns and "negative_reranking" in columns):
+                task_categories.append("sentence-similarity")
         else:
             task_categories: list[str] = []
             gr.Info(

src/synthetic_dataset_generator/app.py

@@ -1,6 +1,7 @@
 from synthetic_dataset_generator._tabbedinterface import TabbedInterface
 
 # from synthetic_dataset_generator.apps.eval import app as eval_app
+from synthetic_dataset_generator.apps.rag import app as rag_app
 from synthetic_dataset_generator.apps.about import app as about_app
 from synthetic_dataset_generator.apps.chat import app as chat_app
 from synthetic_dataset_generator.apps.textcat import app as textcat_app
@@ -22,8 +23,8 @@ button[role="tab"][aria-selected="true"]:hover {border-color: var(--button-prima
 image = """<br><img src="https://raw.githubusercontent.com/argilla-io/synthetic-data-generator/main/assets/logo.svg" alt="Synthetic Data Generator Logo" style="display: block; margin-left: auto; margin-right: auto; width: clamp(50%, 400px, 100%)"/>"""
 
 demo = TabbedInterface(
-    [textcat_app, chat_app, about_app],
-    ["Text Classification", "Chat Data", "About"],
+    [textcat_app, chat_app, rag_app, about_app],
+    ["Text Classification", "Chat Data", "RAG", "About"],
     css=css,
     title=image,
     theme=theme,

src/synthetic_dataset_generator/apps/about.py

@@ -7,7 +7,7 @@ with gr.Blocks() as app:
 
         Introducing the Synthetic Data Generator, a user-friendly application that takes a no-code approach to creating custom datasets with Large Language Models (LLMs). The best part: A simple step-by-step process, making dataset creation a non-technical breeze, allowing anyone to create datasets and models in minutes and without any code.
 
-        The synthetic data generator takes your custom prompt and returns a dataset for your use case, using a synthetic data pipeline. In the background this is powered by [distilabel](https://distilabel.argilla.io/latest/) and the [free Hugging Face text-generation API](https://huggingface.co/docs/api-inference/en/index) but we don’t need to worry about these complexities and we can focus on using the UI.
+        The synthetic data generator takes your custom prompt and returns a dataset for your use case, using a synthetic data pipeline. In the background this is powered by [distilabel](https://distilabel.argilla.io/latest/) and the [free Hugging Face text-generation API](https://huggingface.co/docs/api-inference/en/index) but we don't need to worry about these complexities and we can focus on using the UI.
 
         - Read more in [our announcement blog post](https://huggingface.co/blog/synthetic-data-generator)
         - Find the library on [GitHub](https://github.com/argilla-io/synthetic-data-generator)

src/synthetic_dataset_generator/apps/base.py

@@ -6,7 +6,7 @@ import argilla as rg
 import gradio as gr
 from datasets import Dataset, concatenate_datasets, load_dataset
 from gradio import OAuthToken
-from huggingface_hub import HfApi, upload_file
+from huggingface_hub import HfApi, upload_file, repo_exists
 
 from synthetic_dataset_generator.constants import MAX_NUM_ROWS
 from synthetic_dataset_generator.utils import get_argilla_client
@@ -18,7 +18,7 @@ def validate_argilla_user_workspace_dataset(
     oauth_token: Union[OAuthToken, None] = None,
     progress=gr.Progress(),
 ) -> str:
-    progress(0, desc="Validating dataset configuration")
+    progress(0.1, desc="Validating dataset configuration")
     hf_user = HfApi().whoami(token=oauth_token.token)["name"]
     client = get_argilla_client()
     if dataset_name is None or dataset_name == "":
@@ -38,6 +38,7 @@ def validate_argilla_user_workspace_dataset(
     dataset = client.datasets(name=dataset_name, workspace=hf_user)
     if dataset and not add_to_existing_dataset:
         raise gr.Error(f"Dataset {dataset_name} already exists")
+    progress(1.0, desc="Dataset configuration validated")
     return ""
 
 
@@ -159,3 +160,22 @@ def test_max_num_rows(num_rows: int) -> int:
             f"Number of rows is larger than the configured maximum. Setting number of rows to {MAX_NUM_ROWS}. Set environment variable `MAX_NUM_ROWS` to change this behavior."
         )
     return num_rows
+
+
+def get_iframe(hub_repo_id: str) -> str:
+    if not hub_repo_id:
+        return ""
+
+    if not repo_exists(repo_id=hub_repo_id, repo_type="dataset"):
+        return ""
+
+    url = f"https://huggingface.co/datasets/{hub_repo_id}/embed/viewer"
+    iframe = f"""
+    <iframe
+        src="{url}"
+        frameborder="0"
+        width="100%"
+        height="600px"
+    ></iframe>
+    """
+    return iframe

src/synthetic_dataset_generator/apps/chat.py

@@ -25,7 +25,7 @@ from synthetic_dataset_generator.constants import (
     MODEL,
     SFT_AVAILABLE,
 )
-from synthetic_dataset_generator.pipelines.base import get_rewriten_prompts
+from synthetic_dataset_generator.pipelines.base import get_rewritten_prompts
 from synthetic_dataset_generator.pipelines.chat import (
     DEFAULT_DATASET_DESCRIPTIONS,
     generate_pipeline_code,
@@ -61,10 +61,9 @@ def convert_dataframe_messages(dataframe: pd.DataFrame) -> pd.DataFrame:
 
 
 def generate_system_prompt(dataset_description, progress=gr.Progress()):
-    progress(0.0, desc="Starting")
-    progress(0.3, desc="Initializing")
+    progress(0.1, desc="Initializing")
     generate_description = get_prompt_generator()
-    progress(0.7, desc="Generating")
+    progress(0.5, desc="Generating")
     result = next(
         generate_description.process(
             [
@@ -79,6 +78,7 @@ def generate_system_prompt(dataset_description, progress=gr.Progress()):
 
 
 def generate_sample_dataset(system_prompt, num_turns, progress=gr.Progress()):
+    progress(0.1, desc="Generating sample dataset")
     dataframe = generate_dataset(
         system_prompt=system_prompt,
         num_turns=num_turns,
@@ -86,6 +86,7 @@ def generate_sample_dataset(system_prompt, num_turns, progress=gr.Progress()):
         progress=progress,
         is_sample=True,
     )
+    progress(1.0, desc="Sample dataset generated")
     return dataframe
 
 
@@ -117,7 +118,7 @@ def generate_dataset(
     batch_size = DEFAULT_BATCH_SIZE
 
     # create prompt rewrites
-    prompt_rewrites = get_rewriten_prompts(system_prompt, num_rows)
+    prompt_rewrites = get_rewritten_prompts(system_prompt, num_rows)
 
     # create instructions
     n_processed = 0
@@ -274,6 +275,7 @@ def push_dataset(
         client = get_argilla_client()
         if client is None:
             return ""
+        progress(0.5, desc="Creating dataset in Argilla")
         if "messages" in dataframe.columns:
             settings = rg.Settings(
                 fields=[
@@ -370,7 +372,6 @@ def push_dataset(
             dataframe["completion_length"] = dataframe["completion"].apply(len)
             dataframe["prompt_embeddings"] = get_embeddings(dataframe["prompt"])
 
-        progress(0.5, desc="Creating dataset")
         rg_dataset = client.datasets(name=repo_name, workspace=hf_user)
         if rg_dataset is None:
             rg_dataset = rg.Dataset(
@@ -516,7 +517,6 @@ with gr.Blocks() as app:
                             system_prompt=system_prompt.value,
                             num_turns=num_turns.value,
                             num_rows=num_rows.value,
-                            temperature=temperature.value,
                         )
                         pipeline_code = gr.Code(
                             value=code,
@@ -582,7 +582,7 @@ with gr.Blocks() as app:
                 outputs=[success_message],
             ).success(
                 fn=generate_pipeline_code,
-                inputs=[system_prompt, num_turns, num_rows, temperature],
+                inputs=[system_prompt, num_turns, num_rows],
                 outputs=[pipeline_code],
             ).success(
                 fn=show_pipeline_code_visibility,
@@ -593,7 +593,7 @@ with gr.Blocks() as app:
                 triggers=[clear_btn_part.click, clear_btn_full.click],
                 fn=lambda _: ("", "", 1, _get_dataframe()),
                 inputs=[dataframe],
-                outputs=[dataset_description, system_prompt, num_turns, 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])

src/synthetic_dataset_generator/apps/eval.py

@@ -19,6 +19,7 @@ from huggingface_hub import HfApi, repo_exists
 
 from synthetic_dataset_generator.apps.base import (
     combine_datasets,
+    get_iframe,
     hide_success_message,
     push_pipeline_code_to_hub,
     show_success_message,
@@ -48,25 +49,6 @@ from synthetic_dataset_generator.utils import (
 )
 
 
-def get_iframe(hub_repo_id: str) -> str:
-    if not hub_repo_id:
-        return ""
-
-    if not repo_exists(repo_id=hub_repo_id, repo_type="dataset"):
-        return ""
-
-    url = f"https://huggingface.co/datasets/{hub_repo_id}/embed/viewer"
-    iframe = f"""
-    <iframe
-        src="{url}"
-        frameborder="0"
-        width="100%"
-        height="600px"
-    ></iframe>
-    """
-    return iframe
-
-
 def get_valid_columns(dataframe: pd.DataFrame):
     instruction_valid_columns = []
     response_valid_columns = []
@@ -357,11 +339,15 @@ def push_dataset_to_hub(
     oauth_token: Union[gr.OAuthToken, None],
     private: bool,
     pipeline_code: str,
+    progress=gr.Progress(),
 ):
+    progress(0.0, desc="Validating")
     repo_id = validate_push_to_hub(org_name, repo_name)
+    progress(0.5, desc="Creating dataset")
     dataset = Dataset.from_pandas(dataframe)
     dataset = combine_datasets(repo_id, dataset, oauth_token)
     distiset = Distiset({"default": dataset})
+    progress(0.9, desc="Pushing dataset")
     distiset.push_to_hub(
         repo_id=repo_id,
         private=private,
@@ -370,6 +356,8 @@ def push_dataset_to_hub(
         create_pr=False,
     )
     push_pipeline_code_to_hub(pipeline_code, org_name, repo_name, oauth_token)
+    progress(1.0, desc="Dataset pushed")
+    return dataframe
 
 
 def push_dataset(
@@ -408,6 +396,7 @@ def push_dataset(
         client = get_argilla_client()
         if client is None:
             return ""
+        progress(0.5, desc="Creating dataset in Argilla")
         if eval_type == "chat-eval":
             num_generations = len((dataframe["generations"][0]))
             fields = [
@@ -488,7 +477,6 @@ def push_dataset(
                 dataframe["instruction"].to_list()
             )
 
-            progress(0.5, desc="Creating dataset")
             rg_dataset = client.datasets(name=repo_name, workspace=hf_user)
             if rg_dataset is None:
                 rg_dataset = rg.Dataset(
@@ -628,7 +616,6 @@ def push_dataset(
                 dataframe[f"{column}_length"] = dataframe[column].apply(len)
                 dataframe[f"{column}_embeddings"] = get_embeddings(dataframe[column])
 
-            progress(0.5, desc="Creating dataset")
             rg_dataset = client.datasets(name=repo_name, workspace=hf_user)
             if rg_dataset is None:
                 rg_dataset = rg.Dataset(
@@ -895,12 +882,11 @@ with gr.Blocks() as app:
         outputs=[pipeline_code_ui],
     )
 
-    clear_btn_part.click(fn=lambda x: "", 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=[
-            search_in,
             instruction_instruction_response,
             response_instruction_response,
         ],

src/synthetic_dataset_generator/apps/rag.py

@@ -0,0 +1,896 @@
+import random
+import uuid
+from tqdm import tqdm
+from typing import Union
+
+import argilla as rg
+import gradio as gr
+import pandas as pd
+from datasets import (
+    Dataset,
+    get_dataset_config_names,
+    get_dataset_split_names,
+    load_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,
+    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.pipelines.base import get_rewritten_prompts
+from synthetic_dataset_generator.pipelines.embeddings import (
+    get_embeddings,
+    get_sentence_embedding_dimensions,
+)
+from synthetic_dataset_generator.pipelines.rag import (
+    DEFAULT_DATASET_DESCRIPTIONS,
+    get_chunks_generator,
+    get_prompt_generator,
+    generate_pipeline_code,
+    get_sentence_pair_generator,
+    get_response_generator,
+)
+from synthetic_dataset_generator.utils import (
+    column_to_list,
+    get_argilla_client,
+    get_org_dropdown,
+    get_random_repo_name,
+    swap_visibility,
+)
+
+
+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=["chucks"],
+            label="Documents column",
+            value=col_doc,
+            interactive=(False if col_doc == "" else True),
+            multiselect=False,
+        ),
+    )
+
+
+def generate_system_prompt(dataset_description, progress=gr.Progress()):
+    progress(0.1, desc="Initializing")
+    generate_description = get_prompt_generator()
+    progress(0.5, desc="Generating")
+    result = next(
+        generate_description.process(
+            [
+                {
+                    "instruction": dataset_description,
+                }
+            ]
+        )
+    )[0]["generation"]
+    progress(1.0, desc="Prompt generated")
+    return result
+
+
+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, num_rows, token)
+    else:
+        return _preprocess_input_data(file_paths, num_rows)
+
+
+def generate_dataset(
+    input_type: str,
+    dataframe: pd.DataFrame,
+    system_prompt: str,
+    document_column: str,
+    retrieval: bool = False,
+    reranking: bool = False,
+    num_rows: int = 10,
+    temperature: float = 0.7,
+    is_sample: bool = False,
+    progress=gr.Progress(),
+):
+    num_rows = test_max_num_rows(num_rows)
+    progress(0.0, desc="Generating questions")
+    if input_type == "prompt-input":
+        chunk_generator = get_chunks_generator(
+            temperature=temperature, is_sample=is_sample
+        )
+    else:
+        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)
+            )
+
+    retrieval_generator = get_sentence_pair_generator(
+        action="query",
+        triplet=True if retrieval else False,
+        temperature=temperature,
+        is_sample=is_sample,
+    )
+    response_generator = get_response_generator(
+        temperature=temperature, is_sample=is_sample
+    )
+    if reranking:
+        reranking_generator = get_sentence_pair_generator(
+            action="semantically-similar",
+            triplet=True,
+            temperature=temperature,
+            is_sample=is_sample,
+        )
+    steps = 2 + sum([1 if reranking else 0, 1 if input_type == "prompt-type" else 0])
+    total_steps: int = num_rows * steps
+    step_progress = round(1 / steps, 2)
+    batch_size = DEFAULT_BATCH_SIZE
+
+    # generate chunks
+    if input_type == "prompt-input":
+        n_processed = 0
+        chunk_results = []
+        rewritten_system_prompts = get_rewritten_prompts(system_prompt, num_rows)
+        while n_processed < num_rows:
+            progress(
+                step_progress * n_processed / num_rows,
+                total=total_steps,
+                desc="Generating chunks",
+            )
+            remaining_rows = num_rows - n_processed
+            batch_size = min(batch_size, remaining_rows)
+            inputs = [
+                {"task": random.choice(rewritten_system_prompts)}
+                for _ in range(batch_size)
+            ]
+            chunks = list(chunk_generator.process(inputs=inputs))
+            chunk_results.extend(chunks[0])
+            n_processed += batch_size
+            random.seed(a=random.randint(0, 2**32 - 1))
+        document_data = [chunk["generation"] for chunk in chunk_results]
+        progress(step_progress, desc="Generating chunks")
+
+    # generate questions
+    n_processed = 0
+    retrieval_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)
+        inputs = [
+            {"anchor": document}
+            for document in document_data[n_processed : n_processed + batch_size]
+        ]
+        questions = list(retrieval_generator.process(inputs=inputs))
+        retrieval_results.extend(questions[0])
+        n_processed += batch_size
+    for result in retrieval_results:
+        result["context"] = result["anchor"]
+        if retrieval:
+            result["question"] = result["positive"]
+            result["positive_retrieval"] = result.pop("positive")
+            result["negative_retrieval"] = result.pop("negative")
+        else:
+            result["question"] = result.pop("positive")
+
+    progress(step_progress, desc="Generating questions")
+
+    # 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 = retrieval_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["response"] = result["generation"]
+    progress(step_progress, desc="Generating responses")
+
+    # generate reranking
+    if reranking:
+        n_processed = 0
+        reranking_results = []
+        while n_processed < num_rows:
+            progress(
+                step_progress * n_processed / num_rows,
+                total=total_steps,
+                desc="Generating reranking data",
+            )
+            batch = response_results[n_processed : n_processed + batch_size]
+            batch = list(reranking_generator.process(inputs=batch))
+            reranking_results.extend(batch[0])
+            n_processed += batch_size
+        for result in reranking_results:
+            result["positive_reranking"] = result.pop("positive")
+            result["negative_reranking"] = result.pop("negative")
+    progress(
+        1,
+        total=total_steps,
+        desc="Creating dataset",
+    )
+
+    # create distiset
+    distiset_results = []
+    source_results = reranking_results if reranking else response_results
+    base_keys = ["context", "question", "response"]
+    retrieval_keys = ["positive_retrieval", "negative_retrieval"] if retrieval else []
+    reranking_keys = ["positive_reranking", "negative_reranking"] if reranking else []
+    relevant_keys = base_keys + retrieval_keys + reranking_keys
+
+    for result in source_results:
+        record = {key: result.get(key) for key in relevant_keys if key in result}
+        distiset_results.append(record)
+
+    dataframe = pd.DataFrame(distiset_results)
+
+    progress(1.0, desc="Dataset generation completed")
+    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,
+    repo_name: str,
+    oauth_token: Union[gr.OAuthToken, None],
+    private: bool,
+    pipeline_code: str,
+    progress=gr.Progress(),
+):
+    progress(0.0, desc="Validating")
+    repo_id = validate_push_to_hub(org_name, repo_name)
+    progress(0.5, desc="Creating dataset")
+    dataset = Dataset.from_pandas(dataframe)
+    dataset = combine_datasets(repo_id, dataset, oauth_token)
+    distiset = Distiset({"default": dataset})
+    progress(0.9, desc="Pushing dataset")
+    distiset.push_to_hub(
+        repo_id=repo_id,
+        private=private,
+        include_script=False,
+        token=oauth_token.token,
+        create_pr=False,
+    )
+    push_pipeline_code_to_hub(pipeline_code, org_name, repo_name, oauth_token)
+    progress(1.0, desc="Dataset pushed")
+    return dataframe
+
+
+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,
+    retrieval_reranking: list[str],
+    num_rows: int,
+    temperature: float,
+    pipeline_code: str,
+    oauth_token: Union[gr.OAuthToken, None] = None,
+    progress=gr.Progress(),
+) -> pd.DataFrame:
+    retrieval = "Retrieval" in retrieval_reranking
+    reranking = "Reranking" in retrieval_reranking
+
+    if input_type != "prompt-input":
+        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,
+        retrieval=retrieval,
+        reranking=reranking,
+        num_rows=num_rows,
+        temperature=temperature,
+        is_sample=True,
+    )
+    push_dataset_to_hub(
+        dataframe, org_name, repo_name, oauth_token, private, pipeline_code
+    )
+    try:
+        progress(0.1, desc="Setting up user and workspace")
+        hf_user = HfApi().whoami(token=oauth_token.token)["name"]
+        client = get_argilla_client()
+        if client is None:
+            return ""
+
+        progress(0.5, desc="Creating dataset in Argilla")
+        fields = [
+            rg.TextField(
+                name="context",
+                title="Context",
+                description="Context for the generation",
+            ),
+            rg.ChatField(
+                name="chat",
+                title="Chat",
+                description="User and assistant conversation based on the context",
+            ),
+        ]
+        for item in ["positive", "negative"]:
+            if retrieval:
+                fields.append(
+                    rg.TextField(
+                        name=f"{item}_retrieval",
+                        title=f"{item.capitalize()} retrieval",
+                        description=f"The {item} query for retrieval",
+                    )
+                )
+            if reranking:
+                fields.append(
+                    rg.TextField(
+                        name=f"{item}_reranking",
+                        title=f"{item.capitalize()} reranking",
+                        description=f"The {item} query for reranking",
+                    )
+                )
+
+        questions = [
+            rg.LabelQuestion(
+                name="relevant",
+                title="Are the question and response relevant to the given context?",
+                labels=["yes", "no"],
+            ),
+            rg.LabelQuestion(
+                name="is_response_correct",
+                title="Is the response correct?",
+                labels=["yes", "no"],
+            ),
+        ]
+        for item in ["positive", "negative"]:
+            if retrieval:
+                questions.append(
+                    rg.LabelQuestion(
+                        name=f"is_{item}_retrieval_relevant",
+                        title=f"Is the {item} retrieval relevant?",
+                        labels=["yes", "no"],
+                        required=False,
+                    )
+                )
+            if reranking:
+                questions.append(
+                    rg.LabelQuestion(
+                        name=f"is_{item}_reranking_relevant",
+                        title=f"Is the {item} reranking relevant?",
+                        labels=["yes", "no"],
+                        required=False,
+                    )
+                )
+        metadata = [
+            rg.IntegerMetadataProperty(
+                name=f"{item}_length", title=f"{item.capitalize()} length"
+            )
+            for item in ["context", "question", "response"]
+        ]
+
+        vectors = [
+            rg.VectorField(
+                name=f"{item}_embeddings",
+                dimensions=get_sentence_embedding_dimensions(),
+            )
+            for item in ["context", "question", "response"]
+        ]
+        settings = rg.Settings(
+            fields=fields,
+            questions=questions,
+            metadata=metadata,
+            vectors=vectors,
+            guidelines="Please review the conversation and provide an evaluation.",
+        )
+
+        dataframe["chat"] = dataframe.apply(
+            lambda row: [
+                {"role": "user", "content": row["question"]},
+                {"role": "assistant", "content": row["response"]},
+            ],
+            axis=1,
+        )
+
+        for item in ["context", "question", "response"]:
+            dataframe[f"{item}_length"] = dataframe[item].apply(len)
+            dataframe[f"{item}_embeddings"] = get_embeddings(dataframe[item].to_list())
+
+        rg_dataset = client.datasets(name=repo_name, workspace=hf_user)
+        if rg_dataset is None:
+            rg_dataset = rg.Dataset(
+                name=repo_name,
+                workspace=hf_user,
+                settings=settings,
+                client=client,
+            )
+            rg_dataset = rg_dataset.create()
+
+        progress(0.7, desc="Pushing dataset to Argilla")
+        hf_dataset = Dataset.from_pandas(dataframe)
+        rg_dataset.records.log(records=hf_dataset)
+        progress(1.0, desc="Dataset pushed to Argilla")
+    except Exception as e:
+        raise gr.Error(f"Error pushing dataset to Argilla: {e}")
+    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(visible=False)}
+
+
+def show_pipeline_code_visibility():
+    return {pipeline_code_ui: gr.Accordion(visible=True)}
+
+
+def hide_pipeline_code_visibility():
+    return {pipeline_code_ui: gr.Accordion(visible=False)}
+
+
+######################
+# Gradio UI
+######################
+
+
+with gr.Blocks() as app:
+    with gr.Column() as main_ui:
+        gr.Markdown("## 1. Select your input")
+        with gr.Row(equal_height=False):
+            with gr.Column(scale=2):
+                input_type = gr.Dropdown(
+                    label="Input type",
+                    choices=["dataset-input", "file-input", "prompt-input"],
+                    value="dataset-input",
+                    multiselect=False,
+                    visible=False,
+                )
+                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)
+                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",
+                            )
+
+        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,
+                    )
+
+        gr.HTML(value="<hr>")
+        gr.Markdown(value="## 3. Generate your dataset")
+        with gr.Row(equal_height=False):
+            with gr.Column(scale=2):
+                org_name = get_org_dropdown()
+                repo_name = gr.Textbox(
+                    label="Repo name",
+                    placeholder="dataset_name",
+                    value=f"my-distiset-{str(uuid.uuid4())[:8]}",
+                    interactive=True,
+                )
+                num_rows = gr.Number(
+                    label="Number of rows",
+                    value=10,
+                    interactive=True,
+                    scale=1,
+                )
+                temperature = gr.Slider(
+                    label="Temperature",
+                    minimum=0.1,
+                    maximum=1,
+                    value=0.7,
+                    step=0.1,
+                    interactive=True,
+                )
+                private = gr.Checkbox(
+                    label="Private dataset",
+                    value=False,
+                    interactive=True,
+                    scale=1,
+                )
+                btn_push_to_hub = gr.Button("Push to Hub", variant="primary", scale=2)
+            with gr.Column(scale=3):
+                success_message = gr.Markdown(
+                    visible=True,
+                    min_height=100,  # don't remove this otherwise progress is not visible
+                )
+                with gr.Accordion(
+                    "Customize your pipeline with distilabel",
+                    open=False,
+                    visible=False,
+                ) 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,
+                        retrieval_reranking=retrieval_reranking.value,
+                        num_rows=num_rows.value,
+                    )
+                    pipeline_code = gr.Code(
+                        value=code,
+                        language="python",
+                        label="Distilabel Pipeline Code",
+                    )
+
+    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]
+    )
+
+    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]
+    )
+
+    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]
+    )
+
+    search_in.submit(fn=get_iframe, inputs=search_in, outputs=search_out).then(
+        fn=lambda df: pd.DataFrame(columns=df.columns),
+        inputs=[dataframe],
+        outputs=[dataframe],
+    )
+
+    load_dataset_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,
+        ],
+    )
+
+    load_prompt_btn.click(
+        fn=generate_system_prompt,
+        inputs=[dataset_description],
+        outputs=[system_prompt],
+        show_progress=True,
+    ).success(
+        fn=generate_sample_dataset,
+        inputs=[
+            search_in,
+            file_in,
+            input_type,
+            system_prompt,
+            document_column,
+            retrieval_reranking,
+            num_rows,
+        ],
+        outputs=dataframe,
+    )
+
+    btn_apply_to_sample_dataset.click(
+        fn=generate_sample_dataset,
+        inputs=[
+            search_in,
+            file_in,
+            input_type,
+            system_prompt,
+            document_column,
+            retrieval_reranking,
+            num_rows,
+        ],
+        outputs=dataframe,
+    )
+
+    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],
+    ).success(
+        fn=push_dataset,
+        inputs=[
+            org_name,
+            repo_name,
+            private,
+            search_in,
+            file_in,
+            input_type,
+            system_prompt,
+            document_column,
+            retrieval_reranking,
+            num_rows,
+            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=[
+            search_in,
+            file_in,
+            input_type,
+            system_prompt,
+            document_column,
+            retrieval_reranking,
+            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: ("", [], pd.DataFrame(columns=df.columns)),
+        inputs=[dataframe],
+        outputs=[document_column, retrieval_reranking, 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])

src/synthetic_dataset_generator/apps/textcat.py

@@ -20,7 +20,7 @@ from synthetic_dataset_generator.apps.base import (
     validate_push_to_hub,
 )
 from synthetic_dataset_generator.constants import DEFAULT_BATCH_SIZE
-from synthetic_dataset_generator.pipelines.base import get_rewriten_prompts
+from synthetic_dataset_generator.pipelines.base import get_rewritten_prompts
 from synthetic_dataset_generator.pipelines.embeddings import (
     get_embeddings,
     get_sentence_embedding_dimensions,
@@ -120,7 +120,7 @@ def generate_dataset(
     # create text classification data
     n_processed = 0
     textcat_results = []
-    rewritten_system_prompts = get_rewriten_prompts(system_prompt, num_rows)
+    rewritten_system_prompts = get_rewritten_prompts(system_prompt, num_rows)
     while n_processed < num_rows:
         progress(
             2 * 0.5 * n_processed / num_rows,
@@ -314,6 +314,7 @@ def push_dataset(
         if client is None:
             return ""
         labels = get_preprocess_labels(labels)
+        progress(0.5, desc="Creating dataset in Argilla")
         settings = rg.Settings(
             fields=[
                 rg.TextField(
@@ -354,7 +355,6 @@ def push_dataset(
         dataframe["text_length"] = dataframe["text"].apply(len)
         dataframe["text_embeddings"] = get_embeddings(dataframe["text"].to_list())
 
-        progress(0.5, desc="Creating dataset")
         rg_dataset = client.datasets(name=repo_name, workspace=hf_user)
         if rg_dataset is None:
             rg_dataset = rg.Dataset(
@@ -559,7 +559,6 @@ with gr.Blocks() as app:
                         labels=labels.value,
                         num_labels=len(labels.value) if multi_label.value else 1,
                         num_rows=num_rows.value,
-                        temperature=temperature.value,
                     )
                     pipeline_code = gr.Code(
                         value=code,
@@ -644,7 +643,6 @@ with gr.Blocks() as app:
             labels,
             multi_label,
             num_rows,
-            temperature,
         ],
         outputs=[pipeline_code],
     ).success(
@@ -662,7 +660,7 @@ with gr.Blocks() as app:
             _get_dataframe(),
         ),
         inputs=[dataframe],
-        outputs=[dataset_description, system_prompt, labels, dataframe],
+        outputs=[dataset_description, system_prompt, labels, multi_label, dataframe],
     )
 
     app.load(fn=swap_visibility, outputs=main_ui)

src/synthetic_dataset_generator/pipelines/base.py

@@ -39,7 +39,7 @@ def _get_prompt_rewriter():
     return prompt_rewriter
 
 
-def get_rewriten_prompts(prompt: str, num_rows: int):
+def get_rewritten_prompts(prompt: str, num_rows: int):
     prompt_rewriter = _get_prompt_rewriter()
     # create prompt rewrites
     inputs = [

src/synthetic_dataset_generator/pipelines/chat.py

@@ -227,7 +227,7 @@ def get_response_generator(system_prompt, num_turns, temperature, is_sample):
     return response_generator
 
 
-def generate_pipeline_code(system_prompt, num_turns, num_rows, temperature):
+def generate_pipeline_code(system_prompt, num_turns, num_rows):
     input_mappings = _get_output_mappings(num_turns)
 
     code = f"""
@@ -242,7 +242,9 @@ SYSTEM_PROMPT = "{system_prompt}"
 
 with Pipeline(name="sft") as pipeline:
     magpie = MagpieGenerator(
-        llm={_get_llm_class()}.from_dict({_get_llm().model_dump()}),
+        llm={_get_llm_class()}.from_dict(
+            {_get_llm().dump()}
+        ),
         n_turns={num_turns},
         num_rows={num_rows},
         batch_size=1,

src/synthetic_dataset_generator/pipelines/rag.py

@@ -0,0 +1,302 @@
+import os
+
+from typing import List
+
+from datasets import get_dataset_config_names, get_dataset_split_names
+from distilabel.steps.tasks import (
+    GenerateSentencePair,
+    TextGeneration,
+)
+
+from synthetic_dataset_generator.constants import MAX_NUM_TOKENS
+from synthetic_dataset_generator.pipelines.base import _get_llm, _get_llm_class
+
+DEFAULT_DATASET_DESCRIPTIONS = [
+    "A dataset to retrieve information from legal documents.",
+    "A dataset to search for economical techniques.",
+]
+
+PROMPT_CREATION_PROMPT = """
+
+You are an AI assistant specialized in designing retrieval-augmented generation (RAG) tasks for dataset creation.
+
+Your task is to generate a well-structured and descriptive prompt based on the provided dataset description and company context. Respond with only the generated prompt and nothing else.
+
+The prompt should closely follow the style and structure of the example prompts below. Ensure that you include all relevant details from the dataset description and reflect the company context accurately.
+
+Description: A dataset to retrieve information from legal documents.
+Output: A dataset to retrieve information from a collection of legal documents related to the US law system and the status of contracts.
+
+Description: A dataset to search for economical techniques.
+Output: A dataset to search for economical techniques and strategies for the European market and the financial sector.
+
+Description: A dataset covering FAQ questions for a tech company called Argilla that sells technology datasets within the open-source Natural Language Processing space.
+Output: A dataset covering FAQ questions for a tech company called Argilla that sells technology datasets within the open-source Natural Language Processing space.
+
+Description:
+"""
+
+SYSTEM_PROMPT_CHUCKS = """
+You are a helpful and knowledgeable AI assistant. Your task is to generate concise and informative text chunks relevant to the given retrieval task.
+
+Ensure the text chunks are:
+- Focused and directly related to the retrieval task.
+- Clear, truthful, and based on your general knowledge.
+
+Do not include or reference the retrieval task itself in the generated chunks.
+"""
+
+CHUNKS_TEMPLATE = """You have been assigned to generate text chunks based on the following retrieval task: {{ task }}.
+
+Provide only the text chunks without explaining your process or reasoning.
+
+Ensure the chunks are clear, accurate, and directly relevant to the task.
+
+Use your general knowledge to create informative and precise outputs.
+"""
+
+SYSTEM_PROMPT_RAG = """
+You are a helpful AI assistant. Your task is to answer the following question based on the provided document.
+
+If the answer is not explicitly stated in the document, use your knowledge to provide the most relevant and accurate answer possible.
+
+If you cannot answer the question based on the given information, state that clearly.
+"""
+
+RAG_TEMPLATE = """Document:
+{{ context }}
+
+Question: {{ question }}
+
+Please provide a clear and concise answer to the question based on the information in the document:
+""".rstrip()
+
+
+def get_prompt_generator():
+    generation_kwargs = {
+        "temperature": 0.8,
+        "max_new_tokens": MAX_NUM_TOKENS,
+    }
+    text_generator = TextGeneration(
+        llm=_get_llm(generation_kwargs=generation_kwargs),
+        system_prompt=PROMPT_CREATION_PROMPT,
+        use_system_prompt=True,
+    )
+
+    text_generator.load()
+    return text_generator
+
+
+def get_chunks_generator(temperature, is_sample):
+    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),
+        system_prompt=SYSTEM_PROMPT_CHUCKS,
+        template=CHUNKS_TEMPLATE,
+        columns=["task"],
+        use_system_prompt=True,
+    )
+
+    text_generator.load()
+    return text_generator
+
+
+def get_sentence_pair_generator(action, triplet, temperature, is_sample):
+    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=triplet,
+        action=action,
+        hard_negative=True,
+    )
+    sentence_pair_generator.load()
+    return sentence_pair_generator
+
+
+def get_response_generator(temperature, is_sample):
+    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),
+        system_prompt=SYSTEM_PROMPT_RAG,
+        template=RAG_TEMPLATE,
+        columns=["context", "question"],
+        use_system_prompt=True,
+    )
+
+    text_generator.load()
+    return text_generator
+
+
+def generate_pipeline_code(
+    repo_id: str,
+    file_paths: List[str],
+    input_type: str,
+    system_prompt: str,
+    document_column: str,
+    retrieval_reranking: list[str],
+    num_rows: int = 10,
+) -> str:
+    if repo_id is None:
+        subset = "default"
+        split = "train"
+    else:
+        subset = get_dataset_config_names(repo_id)[0]
+        split = get_dataset_split_names(repo_id, subset)[0]
+    retrieval = "Retrieval" in retrieval_reranking
+    reranking = "Reranking" in retrieval_reranking
+    base_code = f"""
+# Requirements: `pip install distilabel[hf-inference-endpoints]`
+{"import random" if input_type == "prompt-input" else ""}
+from distilabel.models import {_get_llm_class()}
+from distilabel.pipeline import Pipeline
+from distilabel.steps import KeepColumns{", LoadDataFromDicts" if input_type != "dataset-input"  else ""}{", LoadDataFromHub" if input_type == "dataset-input" else ""}{", CombineOutputs" if retrieval and reranking else ""}
+from distilabel.steps.tasks import GenerateSentencePair, TextGeneration {", GenerateTextRetrievalData" if input_type == "prompt-input" else ""}
+
+SYSTEM_PROMPT_RAG = '''
+You are a helpful AI assistant. Your task is to answer the following question based on the provided document.
+
+If the answer is not explicitly stated in the document, use your knowledge to provide the most relevant and accurate answer possible.
+
+If you cannot answer the question based on the given information, state that clearly.
+'''
+
+RAG_TEMPLATE = '''Document:
+{{{{ filename }}}}
+
+Question: {{{{ question }}}}
+
+Please provide a clear and concise answer to the question based on the information in the document:
+'''.rstrip()
+"""
+
+    if input_type == "file-input":
+        base_code += """
+data = process_and_chunk_files(files=[files])
+"""
+
+    if input_type == "prompt-input":
+        pipeline = f"""
+TASK_SYSTEM_PROMPT =  '''
+
+{system_prompt}    
+''' 
+
+with Pipeline(name="rag") as pipeline:
+
+    task_generator = LoadDataFromDicts(data=[{{"task": TASK_SYSTEM_PROMPT}}])
+
+    sentence_similarity_generation = GenerateTextRetrievalData(
+        llm={_get_llm_class()}.from_dict(
+            {_get_llm().dump()}
+        ),
+        seed=random.randint(0, 2**32 - 1),
+        query_type="common",
+        difficulty="high school",
+        clarity="clear",
+        num_generations={num_rows},
+        output_mappings={{"positive_document": "anchor"}},
+    )
+
+    keep_columns_prompt = KeepColumns(
+        columns=["anchor"],
+    )
+    """
+    else:
+        pipeline = """
+with Pipeline(name="rag") as pipeline:
+"""
+        if input_type == "file-input":
+            pipeline += """
+    load_the_dataset = LoadDataFromDicts(
+        data = data,
+    )
+    """
+        else:
+            pipeline += f"""
+    load_the_dataset = LoadDataFromHub(
+        repo_id="{repo_id}",
+        config="{subset}",
+        split="{split}",
+        num_examples={num_rows},
+        batch_size=2,
+        output_mappings={{'{document_column}': 'anchor'}}
+    )
+    """
+
+    pipeline += f"""
+    generate_retrieval_pairs = GenerateSentencePair(
+        triplet={str(retrieval)},
+        hard_negative=True,
+        action="query",
+        llm={_get_llm_class()}.from_dict(
+            {_get_llm().dump()}
+        ),
+        output_mappings={{"positive": "positive_retrieval"{', "negative": "negative_retrieval"' if retrieval else ""}}},
+        input_batch_size=10,
+    )
+    """
+
+    if reranking:
+        pipeline += f"""
+    generate_reranking_pairs = GenerateSentencePair(
+        triplet=True,
+        hard_negative=True,
+        action="semantically-similar",
+        llm={_get_llm_class()}.from_dict(
+            {_get_llm().dump()}
+        ),
+        input_batch_size=10,
+        output_mappings={{"positive": "positive_reranking", "negative": "negative_reranking"}},
+    )
+    
+    combine_outputs = CombineOutputs()
+    """
+
+    pipeline += f"""
+    generate_response = TextGeneration(
+        llm={_get_llm_class()}.from_dict(
+            {_get_llm().dump()}
+        ),
+        system_prompt=SYSTEM_PROMPT_RAG,
+        template=RAG_TEMPLATE,
+        columns=["filename", "question"],
+        use_system_prompt=True,
+        input_mappings={{"filename": "anchor", "question": "positive_retrieval"}},
+        output_mappings={{"generation": "response"}},
+    )
+    
+    keep_columns = KeepColumns(
+        columns=["anchor", "positive_retrieval", "response"{', "negative_retrieval"' if retrieval else ""}{', "positive_reranking", "negative_reranking"' if reranking else ""}],
+    )
+    """
+
+    pipeline_steps = (
+        "[generate_retrieval_pairs, generate_reranking_pairs] >> combine_outputs >> generate_response >> keep_columns"
+        if reranking
+        else "generate_retrieval_pairs >> generate_response >> keep_columns"
+    )
+
+    pipeline += """
+    task_generator >> sentence_similarity_generation >> keep_columns_prompt >> {pipeline_steps}
+""".format(pipeline_steps=pipeline_steps) if input_type == "prompt-input" else """
+    load_the_dataset >> {pipeline_steps}
+""".format(pipeline_steps=pipeline_steps)
+
+    pipeline += """
+    if __name__ == "__main__":
+        distiset = pipeline.run(use_cache=False)
+        print(distiset)
+        if distiset:
+            print(distiset["default"]["train"][0])
+    """
+
+    return base_code + pipeline

src/synthetic_dataset_generator/pipelines/textcat.py

@@ -126,7 +126,6 @@ def generate_pipeline_code(
     labels: List[str] = None,
     num_labels: int = 1,
     num_rows: int = 10,
-    temperature: float = 0.9,
 ) -> str:
     labels = get_preprocess_labels(labels)
     base_code = f"""
@@ -142,10 +141,12 @@ SYSTEM_PROMPT = "{system_prompt}"
 
 with Pipeline(name="textcat") as pipeline:
 
-    task_generator = LoadDataFromDicts(data=[{{"task": TEXT_CLASSIFICATION_TASK}}])
+    task_generator = LoadDataFromDicts(data=[{{"task": SYSTEM_PROMPT}}])
 
     textcat_generation = GenerateTextClassificationData(
-        llm={_get_llm_class()}.from_dict({_get_llm().model_dump()}),
+        llm={_get_llm_class()}.from_dict(
+            {_get_llm().dump()}
+        ),
         seed=random.randint(0, 2**32 - 1),
         difficulty={None if difficulty == "mixed" else repr(difficulty)},
         clarity={None if clarity == "mixed" else repr(clarity)},
@@ -178,10 +179,12 @@ with Pipeline(name="textcat") as pipeline:
     )
 
     textcat_labeller = TextClassification(
-        llm={_get_llm_class()}.from_dict({_get_llm().model_dump()}),
+        llm={_get_llm_class()}.from_dict(
+            {_get_llm().dump()}
+        ),
         n={num_labels},
         available_labels={labels},
-        context=TEXT_CLASSIFICATION_TASK,
+        context=SYSTEM_PROMPT,
         default_label="unknown"
     )