add system prompt rewriter for more dynamic generation SFT

src/synthetic_dataset_generator/apps/sft.py

@@ -1,4 +1,5 @@
 import ast
+import random
 import uuid
 from typing import Dict, List, Union
 
@@ -32,6 +33,7 @@ from synthetic_dataset_generator.pipelines.sft import (
     generate_pipeline_code,
     get_magpie_generator,
     get_prompt_generator,
+    get_prompt_rewriter,
     get_response_generator,
 )
 from synthetic_dataset_generator.utils import (
@@ -103,6 +105,7 @@ 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
     )
@@ -112,6 +115,16 @@ def generate_dataset(
     total_steps: int = num_rows * 2
     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 / 50))
+    ]
+    batch = list(prompt_rewriter.process(inputs=inputs))
+    prompt_rewrites = [entry["generation"] for entry in batch[0]] + [system_prompt]
+
     # create instructions
     n_processed = 0
     magpie_results = []
@@ -123,7 +136,8 @@ def generate_dataset(
         )
         remaining_rows = num_rows - n_processed
         batch_size = min(batch_size, remaining_rows)
-        inputs = [{"system_prompt": system_prompt} for _ in range(batch_size)]
+        rewritten_system_prompt = random.choice(prompt_rewrites)
+        inputs = [{"system_prompt": rewritten_system_prompt} for _ in range(batch_size)]
         batch = list(magpie_generator.process(inputs=inputs))
         magpie_results.extend(batch[0])
         n_processed += batch_size
@@ -487,7 +501,7 @@ with gr.Blocks() as app:
                 with gr.Column(scale=3):
                     success_message = gr.Markdown(
                         visible=True,
-                        height=100,  # don't remove this otherwise progress is not visible
+                        min_height=100,  # don't remove this otherwise progress is not visible
                     )
                     with gr.Accordion(
                         "Customize your pipeline with distilabel",
@@ -543,6 +557,7 @@ with gr.Blocks() as app:
             fn=hide_pipeline_code_visibility,
             inputs=[],
             outputs=[pipeline_code_ui],
+            show_progress=True,
         ).success(
             fn=push_dataset,
             inputs=[

src/synthetic_dataset_generator/pipelines/sft.py

@@ -175,12 +175,11 @@ def get_magpie_generator(system_prompt, num_turns, temperature, is_sample):
                 generation_kwargs={
                     "temperature": temperature,
                     "do_sample": True,
-                    "max_new_tokens": 256 if is_sample else MAX_NUM_TOKENS,
+                    "max_new_tokens": 256 if is_sample else int(MAX_NUM_TOKENS * 0.25),
                     "stop_sequences": _STOP_SEQUENCES,
                 },
             ),
             n_turns=num_turns,
-            system_prompt=system_prompt,
             output_mappings=output_mappings,
             only_instruction=True,
         )
@@ -195,19 +194,34 @@ def get_magpie_generator(system_prompt, num_turns, temperature, is_sample):
                 generation_kwargs={
                     "temperature": temperature,
                     "do_sample": True,
-                    "max_new_tokens": 256 if is_sample else MAX_NUM_TOKENS,
+                    "max_new_tokens": 256 if is_sample else int(MAX_NUM_TOKENS * 0.5),
                     "stop_sequences": _STOP_SEQUENCES,
                 },
             ),
             end_with_user=True,
             n_turns=num_turns,
-            system_prompt=system_prompt,
             output_mappings=output_mappings,
         )
     magpie_generator.load()
     return magpie_generator
 
 
+def get_prompt_rewriter():
+    prompt_rewriter = TextGeneration(
+        llm=InferenceEndpointsLLM(
+            model_id=MODEL,
+            tokenizer_id=MODEL,
+            base_url=BASE_URL,
+            api_key=_get_next_api_key(),
+            generation_kwargs={
+                "temperature": 1,
+            },
+        ),
+    )
+    prompt_rewriter.load()
+    return prompt_rewriter
+
+
 def get_response_generator(system_prompt, num_turns, temperature, is_sample):
     if num_turns == 1:
         response_generator = TextGeneration(
@@ -218,7 +232,7 @@ def get_response_generator(system_prompt, num_turns, temperature, is_sample):
                 api_key=_get_next_api_key(),
                 generation_kwargs={
                     "temperature": temperature,
-                    "max_new_tokens": 256 if is_sample else MAX_NUM_TOKENS,
+                    "max_new_tokens": 256 if is_sample else int(MAX_NUM_TOKENS * 0.5),
                 },
             ),
             system_prompt=system_prompt,