Merge pull request #28 from argilla-io/bug/fix-bugs

src/synthetic_dataset_generator/apps/base.py

@@ -64,7 +64,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 +93,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(

src/synthetic_dataset_generator/apps/chat.py

@@ -60,7 +60,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,7 +77,7 @@ def generate_system_prompt(dataset_description, progress=gr.Progress()):
     return result
 
 
-def generate_sample_dataset(system_prompt, num_turns, progress=gr.Progress()):
+def generate_sample_dataset(system_prompt: str, num_turns: int, progress=gr.Progress()):
     progress(0.1, desc="Generating sample dataset")
     dataframe = generate_dataset(
         system_prompt=system_prompt,
@@ -109,7 +109,7 @@ def generate_dataset(
     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
+        num_turns, temperature, is_sample
     )
     response_generator = get_response_generator(
         system_prompt, num_turns, temperature, is_sample
@@ -267,7 +267,12 @@ def push_dataset(
         temperature=temperature,
     )
     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")
@@ -524,77 +529,77 @@ 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,
-            )
+    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,
+    )
 
-            btn_apply_to_sample_dataset.click(
-                fn=generate_sample_dataset,
-                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,
+    )
 
-            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)
+    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=[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

@@ -889,6 +889,7 @@ with gr.Blocks() as app:
         outputs=[
             instruction_instruction_response,
             response_instruction_response,
+            dataframe
         ],
     )
 

src/synthetic_dataset_generator/apps/rag.py

@@ -76,7 +76,7 @@ def _load_dataset_from_hub(
     progress=gr.Progress(track_tqdm=True),
 ):
     if not repo_id:
-        raise gr.Error("Hub repo id is required")
+        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)
@@ -101,7 +101,10 @@ def _load_dataset_from_hub(
     )
 
 
-def _preprocess_input_data(file_paths, num_rows, progress=gr.Progress(track_tqdm=True)):
+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
 
@@ -131,7 +134,7 @@ def _preprocess_input_data(file_paths, num_rows, progress=gr.Progress(track_tqdm
     )
 
 
-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")
@@ -753,7 +756,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,
@@ -891,7 +893,6 @@ with gr.Blocks() as app:
         fn=generate_pipeline_code,
         inputs=[
             search_in,
-            file_in,
             input_type,
             system_prompt,
             document_column,

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[
@@ -657,6 +662,7 @@ with gr.Blocks() as app:
             "",
             "",
             [],
+            "",
             _get_dataframe(),
         ),
         inputs=[dataframe],

src/synthetic_dataset_generator/pipelines/chat.py

@@ -140,7 +140,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 +162,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,7 +203,9 @@ 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_response_generator(
+    system_prompt: str, num_turns: int, temperature: float, is_sample: bool
+):
     if num_turns == 1:
         generation_kwargs = {
             "temperature": temperature,
@@ -229,7 +231,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):
+def generate_pipeline_code(system_prompt: str, num_turns: int, num_rows: int):
     input_mappings = _get_output_mappings(num_turns)
 
     code = f"""

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

@@ -87,7 +87,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 +104,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,7 +119,7 @@ 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,
@@ -138,7 +138,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,

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,7 +104,7 @@ 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,