add randomisation of system prompts for generation

README.md

@@ -87,7 +87,7 @@ 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.
 - `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/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`.
 
 SFT and Chat Data generation is only supported with Hugging Face Inference Endpoints , and you can set the following environment variables use it with models other than Llama3 and Qwen2.

examples/ollama_deployment.py

@@ -9,6 +9,9 @@ assert os.getenv("HF_TOKEN")  # push the data to huggingface
 os.environ["OLLAMA_BASE_URL"] = "http://127.0.0.1:11434/"  # ollama base url
 os.environ["MODEL"] = "llama3.1:8b-instruct-q8_0"  # model id
 os.environ["TOKENIZER_ID"] = "meta-llama/Llama-3.1-8B-Instruct"  # tokenizer id
-os.environ["MAGPIE_PRE_QUERY_TEMPLATE"] = "llama3"  # magpie template
+os.environ["MAGPIE_PRE_QUERY_TEMPLATE"] = "llama3"
+os.environ["MAX_NUM_ROWS"] = "10000"
+os.environ["DEFAULT_BATCH_SIZE"] = "5"
+os.environ["MAX_NUM_TOKENS"] = "1024"
 
 launch()

src/synthetic_dataset_generator/apps/base.py

@@ -77,10 +77,15 @@ def validate_push_to_hub(org_name, repo_name):
     return repo_id
 
 
-def combine_datasets(repo_id: str, dataset: Dataset) -> Dataset:
+def combine_datasets(
+    repo_id: str, dataset: Dataset, oauth_token: Union[OAuthToken, None]
+) -> Dataset:
     try:
         new_dataset = load_dataset(
-            repo_id, split="train", download_mode="force_redownload"
+            repo_id,
+            split="train",
+            download_mode="force_redownload",
+            token=oauth_token.token,
         )
         return concatenate_datasets([dataset, new_dataset])
     except Exception:

src/synthetic_dataset_generator/apps/chat.py

@@ -25,12 +25,12 @@ from synthetic_dataset_generator.constants import (
     MODEL,
     SFT_AVAILABLE,
 )
+from synthetic_dataset_generator.pipelines.base import get_rewriten_prompts
 from synthetic_dataset_generator.pipelines.chat import (
     DEFAULT_DATASET_DESCRIPTIONS,
     generate_pipeline_code,
     get_magpie_generator,
     get_prompt_generator,
-    get_prompt_rewriter,
     get_response_generator,
 )
 from synthetic_dataset_generator.pipelines.embeddings import (
@@ -40,6 +40,7 @@ from synthetic_dataset_generator.pipelines.embeddings import (
 from synthetic_dataset_generator.utils import (
     get_argilla_client,
     get_org_dropdown,
+    get_random_repo_name,
     swap_visibility,
 )
 
@@ -106,7 +107,6 @@ def generate_dataset(
 ) -> pd.DataFrame:
     num_rows = test_max_num_rows(num_rows)
     progress(0.0, desc="(1/2) Generating instructions")
-    prompt_rewriter = get_prompt_rewriter()
     magpie_generator = get_magpie_generator(
         system_prompt, num_turns, temperature, is_sample
     )
@@ -117,14 +117,7 @@ def generate_dataset(
     batch_size = DEFAULT_BATCH_SIZE
 
     # create prompt rewrites
-    inputs = [
-        {
-            "instruction": f"Rewrite this prompt keeping the same structure but highlighting different aspects of the original without adding anything new. Original prompt: {system_prompt} Rewritten prompt: "
-        }
-        for i in range(int(num_rows / 100))
-    ]
-    batch = list(prompt_rewriter.process(inputs=inputs))
-    prompt_rewrites = [entry["generation"] for entry in batch[0]] + [system_prompt]
+    prompt_rewrites = get_rewriten_prompts(system_prompt, num_rows)
 
     # create instructions
     n_processed = 0
@@ -142,6 +135,7 @@ def generate_dataset(
         batch = list(magpie_generator.process(inputs=inputs))
         magpie_results.extend(batch[0])
         n_processed += batch_size
+        random.seed(a=random.randint(0, 2**32 - 1))
     progress(0.5, desc="(1/2) Generating instructions")
 
     # generate responses
@@ -158,6 +152,7 @@ def generate_dataset(
             responses = list(response_generator.process(inputs=batch))
             response_results.extend(responses[0])
             n_processed += batch_size
+            random.seed(a=random.randint(0, 2**32 - 1))
         for result in response_results:
             result["prompt"] = result["instruction"]
             result["completion"] = result["generation"]
@@ -178,6 +173,7 @@ def generate_dataset(
             responses = list(response_generator.process(inputs=batch))
             response_results.extend(responses[0])
             n_processed += batch_size
+            random.seed(a=random.randint(0, 2**32 - 1))
         for result in response_results:
             result["messages"].append(
                 {"role": "assistant", "content": result["generation"]}
@@ -236,7 +232,7 @@ def push_dataset_to_hub(
     dataframe = convert_dataframe_messages(dataframe)
     progress(0.7, desc="Creating dataset")
     dataset = Dataset.from_pandas(dataframe)
-    dataset = combine_datasets(repo_id, dataset)
+    dataset = combine_datasets(repo_id, dataset, oauth_token)
     progress(0.9, desc="Pushing dataset")
     distiset = Distiset({"default": dataset})
     distiset.push_to_hub(
@@ -600,4 +596,5 @@ with gr.Blocks() as app:
                 outputs=[dataset_description, 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)

src/synthetic_dataset_generator/apps/eval.py

@@ -41,6 +41,7 @@ from synthetic_dataset_generator.utils import (
     extract_column_names,
     get_argilla_client,
     get_org_dropdown,
+    get_random_repo_name,
     pad_or_truncate_list,
     process_columns,
     swap_visibility,
@@ -359,7 +360,7 @@ def push_dataset_to_hub(
 ):
     repo_id = validate_push_to_hub(org_name, repo_name)
     dataset = Dataset.from_pandas(dataframe)
-    dataset = combine_datasets(repo_id, dataset)
+    dataset = combine_datasets(repo_id, dataset, oauth_token)
     distiset = Distiset({"default": dataset})
     distiset.push_to_hub(
         repo_id=repo_id,
@@ -907,3 +908,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])

src/synthetic_dataset_generator/apps/textcat.py

@@ -20,6 +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.embeddings import (
     get_embeddings,
     get_sentence_embedding_dimensions,
@@ -35,6 +36,7 @@ from synthetic_dataset_generator.utils import (
     get_argilla_client,
     get_org_dropdown,
     get_preprocess_labels,
+    get_random_repo_name,
     swap_visibility,
 )
 
@@ -106,7 +108,7 @@ def generate_dataset(
     )
     updated_system_prompt = f"{system_prompt}. Optional labels: {', '.join(labels)}."
     if multi_label:
-        updated_system_prompt = f"{updated_system_prompt}. Only apply relevant labels. Applying less labels is better than applying too many labels."
+        updated_system_prompt = f"{updated_system_prompt}. Only apply relevant labels. Applying less labels is always better than applying too many labels."
     labeller_generator = get_labeller_generator(
         system_prompt=updated_system_prompt,
         labels=labels,
@@ -118,6 +120,7 @@ def generate_dataset(
     # create text classification data
     n_processed = 0
     textcat_results = []
+    rewritten_system_prompts = get_rewriten_prompts(system_prompt, num_rows)
     while n_processed < num_rows:
         progress(
             2 * 0.5 * n_processed / num_rows,
@@ -128,25 +131,24 @@ def generate_dataset(
         batch_size = min(batch_size, remaining_rows)
         inputs = []
         for _ in range(batch_size):
+            k = 1
             if multi_label:
                 num_labels = len(labels)
                 k = int(
                     random.betavariate(alpha=(num_labels - 1), beta=num_labels)
                     * num_labels
                 )
-            else:
-                k = 1
-
             sampled_labels = random.sample(labels, min(k, len(labels)))
             random.shuffle(sampled_labels)
             inputs.append(
                 {
-                    "task": f"{system_prompt}. The text represents the following categories: {', '.join(sampled_labels)}"
+                    "task": f"{random.choice(rewritten_system_prompts)}. The text represents the following categories: {', '.join(sampled_labels)}"
                 }
             )
         batch = list(textcat_generator.process(inputs=inputs))
         textcat_results.extend(batch[0])
         n_processed += batch_size
+        random.seed(a=random.randint(0, 2**32 - 1))
     for result in textcat_results:
         result["text"] = result["input_text"]
 
@@ -164,6 +166,7 @@ def generate_dataset(
         labels_batch = list(labeller_generator.process(inputs=batch))
         labeller_results.extend(labels_batch[0])
         n_processed += batch_size
+        random.seed(a=random.randint(0, 2**32 - 1))
     progress(
         1,
         total=total_steps,
@@ -250,7 +253,7 @@ def push_dataset_to_hub(
         dataframe.reset_index(drop=True),
         features=features,
     )
-    dataset = combine_datasets(repo_id, dataset)
+    dataset = combine_datasets(repo_id, dataset, oauth_token)
     distiset = Distiset({"default": dataset})
     progress(0.9, desc="Pushing dataset")
     distiset.push_to_hub(
@@ -662,3 +665,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])

src/synthetic_dataset_generator/constants.py

@@ -81,7 +81,7 @@ if OPENAI_BASE_URL:
 
 if not SFT_AVAILABLE:
     warnings.warn(
-        "`SFT_AVAILABLE` is set to `False`. Use Hugging Face Inference Endpoints or Ollama to generate chat data, provide a `TOKENIZER_ID` and `MAGPIE_PRE_QUERY_TEMPLATE`."
+        "`SFT_AVAILABLE` is set to `False`. Use Hugging Face Inference Endpoints or Ollama to generate chat data, provide a `TOKENIZER_ID` and `MAGPIE_PRE_QUERY_TEMPLATE`. You can also use `HUGGINGFACE_BASE_URL` to with vllm."
     )
     MAGPIE_PRE_QUERY_TEMPLATE = None
 

src/synthetic_dataset_generator/pipelines/base.py

@@ -1,8 +1,13 @@
+import math
+import random
+
 import gradio as gr
 from distilabel.llms import InferenceEndpointsLLM, OllamaLLM, OpenAILLM
+from distilabel.steps.tasks import TextGeneration
 
 from synthetic_dataset_generator.constants import (
     API_KEYS,
+    DEFAULT_BATCH_SIZE,
     HUGGINGFACE_BASE_URL,
     MAGPIE_PRE_QUERY_TEMPLATE,
     MODEL,
@@ -21,6 +26,41 @@ def _get_next_api_key():
     return api_key
 
 
+def _get_prompt_rewriter():
+    generation_kwargs = {
+        "temperature": 1,
+    }
+    system_prompt = "You are a prompt rewriter. You are given a prompt and you need to rewrite it keeping the same structure but highlighting different aspects of the original without adding anything new."
+    prompt_rewriter = TextGeneration(
+        llm=_get_llm(generation_kwargs=generation_kwargs),
+        system_prompt=system_prompt,
+        use_system_prompt=True,
+    )
+    prompt_rewriter.load()
+    return prompt_rewriter
+
+
+def get_rewriten_prompts(prompt: str, num_rows: int):
+    prompt_rewriter = _get_prompt_rewriter()
+    # create prompt rewrites
+    inputs = [
+        {"instruction": f"Original prompt: {prompt} \nRewritten prompt: "}
+        for i in range(math.floor(num_rows / 100))
+    ]
+    n_processed = 0
+    prompt_rewrites = [prompt]
+    while n_processed < num_rows:
+        batch = list(
+            prompt_rewriter.process(
+                inputs=inputs[n_processed : n_processed + DEFAULT_BATCH_SIZE]
+            )
+        )
+        prompt_rewrites += [entry["generation"] for entry in batch[0]]
+        n_processed += DEFAULT_BATCH_SIZE
+        random.seed(a=random.randint(0, 2**32 - 1))
+    return prompt_rewrites
+
+
 def _get_llm(use_magpie_template=False, **kwargs):
     if OPENAI_BASE_URL:
         llm = OpenAILLM(

src/synthetic_dataset_generator/pipelines/chat.py

@@ -203,15 +203,6 @@ def get_magpie_generator(system_prompt, num_turns, temperature, is_sample):
     return magpie_generator
 
 
-def get_prompt_rewriter():
-    generation_kwargs = {
-        "temperature": 1,
-    }
-    prompt_rewriter = TextGeneration(llm=_get_llm(generation_kwargs=generation_kwargs))
-    prompt_rewriter.load()
-    return prompt_rewriter
-
-
 def get_response_generator(system_prompt, num_turns, temperature, is_sample):
     if num_turns == 1:
         generation_kwargs = {

src/synthetic_dataset_generator/pipelines/textcat.py

@@ -94,7 +94,6 @@ def get_textcat_generator(difficulty, clarity, temperature, is_sample):
         "top_p": 0.95,
     }
     llm = _get_llm(generation_kwargs=generation_kwargs)
-
     textcat_generator = GenerateTextClassificationData(
         llm=llm,
         difficulty=None if difficulty == "mixed" else difficulty,
@@ -111,7 +110,6 @@ def get_labeller_generator(system_prompt, labels, multi_label):
         "max_new_tokens": MAX_NUM_TOKENS,
     }
     llm = _get_llm(generation_kwargs=generation_kwargs)
-
     labeller_generator = TextClassification(
         llm=llm,
         context=system_prompt,

src/synthetic_dataset_generator/utils.py

@@ -1,4 +1,5 @@
 import json
+import uuid
 import warnings
 from typing import List, Optional, Union
 
@@ -55,6 +56,10 @@ def list_orgs(oauth_token: Union[OAuthToken, None] = None):
     return organizations
 
 
+def get_random_repo_name():
+    return f"my-distiset-{str(uuid.uuid4())[:8]}"
+
+
 def get_org_dropdown(oauth_token: Union[OAuthToken, None] = None):
     if oauth_token is not None:
         orgs = list_orgs(oauth_token)