add lacking typing

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")

src/synthetic_dataset_generator/apps/rag.py

@@ -101,7 +101,7 @@ 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)):
     data = {}
     total_chunks = 0
 
@@ -131,7 +131,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 +753,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 +890,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[

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,