update sft and use input parameters

src/distilabel_dataset_generator/apps/sft.py

@@ -24,7 +24,6 @@ from src.distilabel_dataset_generator.pipelines.embeddings import (
 )
 from src.distilabel_dataset_generator.pipelines.sft import (
     DEFAULT_DATASET_DESCRIPTIONS,
-    PROMPT_CREATION_PROMPT,
     generate_pipeline_code,
     get_magpie_generator,
     get_prompt_generator,
@@ -55,7 +54,6 @@ def convert_dataframe_messages(dataframe: pd.DataFrame) -> pd.DataFrame:
 
 def generate_system_prompt(dataset_description, temperature, progress=gr.Progress()):
     progress(0.0, desc="Generating system prompt")
-
     progress(0.3, desc="Initializing text generation")
     generate_description = get_prompt_generator(temperature)
     progress(0.7, desc="Generating system prompt")
@@ -63,20 +61,19 @@ def generate_system_prompt(dataset_description, temperature, progress=gr.Progres
         generate_description.process(
             [
                 {
-                    "system_prompt": PROMPT_CREATION_PROMPT,
                     "instruction": dataset_description,
                 }
             ]
         )
     )[0]["generation"]
     progress(1.0, desc="System prompt generated")
-    return result, pd.DataFrame()
+    return result
 
 
-def generate_sample_dataset(system_prompt, progress=gr.Progress()):
+def generate_sample_dataset(system_prompt, num_turns, progress=gr.Progress()):
     df = generate_dataset(
         system_prompt=system_prompt,
-        num_turns=1,
+        num_turns=num_turns,
         num_rows=10,
         progress=progress,
         is_sample=True,
@@ -92,10 +89,8 @@ def generate_dataset(
     progress=gr.Progress(),
 ) -> pd.DataFrame:
     progress(0.0, desc="(1/2) Generating instructions")
-    magpie_generator = get_magpie_generator(
-        num_turns, num_rows, system_prompt, is_sample
-    )
-    response_generator = get_response_generator(num_turns, system_prompt, is_sample)
+    magpie_generator = get_magpie_generator(system_prompt, num_turns, is_sample)
+    response_generator = get_response_generator(system_prompt, num_turns, is_sample)
     total_steps: int = num_rows * 2
     batch_size = DEFAULT_BATCH_SIZE
 
@@ -405,7 +400,7 @@ with gr.Blocks(css=_LOGGED_OUT_CSS) as app:
                 )
             with gr.Column(scale=3):
                 dataframe = gr.Dataframe(
-                    headers=["prompt", "completion"], wrap=True, height=300
+                    headers=["prompt", "completion"], wrap=True, height=500, interactive=False
                 )
 
         gr.HTML(value="<hr>")
@@ -450,15 +445,21 @@ with gr.Blocks(css=_LOGGED_OUT_CSS) as app:
                         label="Distilabel Pipeline Code",
                     )
 
-    gr.on(
-        triggers=[load_btn.click, btn_apply_to_sample_dataset.click],
+    load_btn.click(
         fn=generate_system_prompt,
         inputs=[dataset_description, temperature],
-        outputs=[system_prompt, dataframe],
+        outputs=[system_prompt],
         show_progress=True,
     ).then(
         fn=generate_sample_dataset,
-        inputs=[system_prompt],
+        inputs=[system_prompt, num_turns],
+        outputs=[dataframe],
+        show_progress=True,
+    )
+    
+    btn_apply_to_sample_dataset.click(
+        fn=generate_sample_dataset,
+        inputs=[system_prompt, num_turns],
         outputs=[dataframe],
         show_progress=True,
     )

src/distilabel_dataset_generator/pipelines/sft.py

@@ -138,51 +138,6 @@ def _get_output_mappings(num_turns):
         return {"conversation": "messages"}
 
 
-def generate_pipeline_code(system_prompt, num_turns, num_rows):
-    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 InferenceEndpointsLLM
-
-MODEL = "{MODEL}"
-SYSTEM_PROMPT = "{system_prompt}"
-os.environ["HF_TOKEN"] = "hf_xxx" # https://huggingface.co/settings/tokens/new?ownUserPermissions=repo.content.read&ownUserPermissions=repo.write&globalPermissions=inference.serverless.write&canReadGatedRepos=true&tokenType=fineGrained
-
-with Pipeline(name="sft") as pipeline:
-    magpie = MagpieGenerator(
-        llm=InferenceEndpointsLLM(
-            model_id=MODEL,
-            tokenizer_id=MODEL,
-            magpie_pre_query_template="llama3",
-            generation_kwargs={{
-                "temperature": 0.9,
-                "do_sample": True,
-                "max_new_tokens": 2048,
-                "stop_sequences": {_STOP_SEQUENCES}
-            }},
-            api_key=os.environ["HF_TOKEN"],
-        ),
-        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 get_prompt_generator(temperature):
     prompt_generator = TextGeneration(
         llm=InferenceEndpointsLLM(
@@ -195,13 +150,14 @@ def get_prompt_generator(temperature):
                 "do_sample": True,
             },
         ),
+        system_prompt=PROMPT_CREATION_PROMPT,
         use_system_prompt=True,
     )
     prompt_generator.load()
     return prompt_generator
 
 
-def get_magpie_generator(num_turns, num_rows, system_prompt, is_sample):
+def get_magpie_generator(system_prompt, num_turns, is_sample):
     input_mappings = _get_output_mappings(num_turns)
     output_mappings = input_mappings.copy()
     if num_turns == 1:
@@ -246,7 +202,7 @@ def get_magpie_generator(num_turns, num_rows, system_prompt, is_sample):
     return magpie_generator
 
 
-def get_response_generator(num_turns, system_prompt, is_sample):
+def get_response_generator(system_prompt, num_turns, is_sample):
     if num_turns == 1:
         response_generator = TextGeneration(
             llm=InferenceEndpointsLLM(
@@ -278,3 +234,48 @@ def get_response_generator(num_turns, system_prompt, is_sample):
         )
     response_generator.load()
     return response_generator
+
+
+def generate_pipeline_code(system_prompt, num_turns, num_rows):
+    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 InferenceEndpointsLLM
+
+MODEL = "{MODEL}"
+SYSTEM_PROMPT = "{system_prompt}"
+os.environ["HF_TOKEN"] = "hf_xxx" # https://huggingface.co/settings/tokens/new?ownUserPermissions=repo.content.read&ownUserPermissions=repo.write&globalPermissions=inference.serverless.write&canReadGatedRepos=true&tokenType=fineGrained
+
+with Pipeline(name="sft") as pipeline:
+    magpie = MagpieGenerator(
+        llm=InferenceEndpointsLLM(
+            model_id=MODEL,
+            tokenizer_id=MODEL,
+            magpie_pre_query_template="llama3",
+            generation_kwargs={{
+                "temperature": 0.9,
+                "do_sample": True,
+                "max_new_tokens": 2048,
+                "stop_sequences": {_STOP_SEQUENCES}
+            }},
+            api_key=os.environ["HF_TOKEN"],
+        ),
+        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