linting

Author: andreatgretel

examples/example.py

@@ -1,14 +1,14 @@
 from data_designer.essentials import (
+    CategorySamplerParams,
+    DataDesigner,
     DataDesignerConfigBuilder,
-    ModelConfig,
     InferenceParameters,
-    SamplerColumnConfig,
-    CategorySamplerParams,
-    SubcategorySamplerParams,
-    PersonSamplerParams,
     LLMTextColumnConfig,
+    ModelConfig,
+    PersonSamplerParams,
+    SamplerColumnConfig,
     Score,
-    DataDesigner,
+    SubcategorySamplerParams,
     ToJsonlProcessorConfig,
 )
 
@@ -40,10 +40,7 @@
 
 # Unique record ID
 config_builder.add_column(
-    name="record_id",
-    column_type="sampler",
-    sampler_type="uuid",
-    params={"short_form": True, "uppercase": True}
+    name="record_id", column_type="sampler", sampler_type="uuid", params={"short_form": True, "uppercase": True}
 )
 
 # ESI level (balanced sampling)
@@ -121,9 +118,7 @@
     SamplerColumnConfig(
         name="writing_style",
         sampler_type="category",
-        params=CategorySamplerParams(
-            values=["Draft", "Adequate", "Polished"]
-        ),
+        params=CategorySamplerParams(values=["Draft", "Adequate", "Polished"]),
     )
 )
 
@@ -159,8 +154,8 @@
         "4": "Note is well-aligned, with only minor details that might be slightly inconsistent.",
         "3": "Note is generally consistent, but some key clinical indicators are missing or don't fully match the ESI level.",
         "2": "Note shows significant inconsistency between the clinical details and the assigned ESI level.",
-        "1": "Note is clinically incoherent and does not reflect the assigned ESI level or scenario at all."
-    }
+        "1": "Note is clinically incoherent and does not reflect the assigned ESI level or scenario at all.",
+    },
 )
 
 # Rubric: ESI level complexity (reduced to 3 levels: Simple, Moderate, Complex)
@@ -170,8 +165,8 @@
     options={
         "Complex": "Note contains subtle or conflicting information, requiring clinical reasoning to distinguish between ESI levels.",
         "Moderate": "Note requires some clinical inference; indicators are present but not always immediately obvious.",
-        "Simple": "Note uses clear, direct, or textbook indicators that make the ESI level obvious."
-    }
+        "Simple": "Note uses clear, direct, or textbook indicators that make the ESI level obvious.",
+    },
 )
 
 jsonl_entry_template = {
@@ -183,22 +178,17 @@
                 f" The possible levels are: {', '.join([repr(level) for level in ESI_LEVELS])}."
                 " Carefully analyze the clinical details in the triage note, focusing on patient acuity, resource needs, and risk of rapid deterioration."
                 " Respond with only the selected ESI level description, exactly matching one of the listed possibilities. Do not provide extra text or explanation."
-            )
+            ),
         },
         {
             "role": "user",
             "content": (
                 "Triage Note: {{ content }}\n"
                 "Classify the ESI level for this note based on the provided definitions."
-                " Respond in JSON format only: { \"esi_level_description\": \"...\" }"
-            )
-        },
-        {
-            "role": "assistant",
-            "content": (
-                '{ "esi_level_description": "{{ esi_level_description }}" }'
-            )
+                ' Respond in JSON format only: { "esi_level_description": "..." }'
+            ),
         },
+        {"role": "assistant", "content": ('{ "esi_level_description": "{{ esi_level_description }}" }')},
     ],
 }
 
@@ -213,5 +203,7 @@
     )
 )
 
-dd = DataDesigner(artifact_path="./artifacts", blob_storage_path="/Users/amanoel/Data/nemotron-personas-datasets_v0.0.6")
-preview = dd.preview(config_builder, num_records=10)
+dd = DataDesigner(
+    artifact_path="./artifacts", blob_storage_path="/Users/amanoel/Data/nemotron-personas-datasets_v0.0.6"
+)
+preview = dd.preview(config_builder, num_records=10)

src/data_designer/config/processors.py

@@ -21,7 +21,7 @@ class ProcessorType(str, Enum):
 class ProcessorConfig(ConfigBase, ABC):
     build_stage: BuildStage = Field(
         default=BuildStage.POST_BATCH,
-        description=f"The stage at which the processor will run. Supported stages: {', '.join(SUPPORTED_STAGES)}"
+        description=f"The stage at which the processor will run. Supported stages: {', '.join(SUPPORTED_STAGES)}",
     )
 
     @field_validator("build_stage")
@@ -58,4 +58,4 @@ class ToJsonlProcessorConfig(ProcessorConfig):
     def validate_fraction_per_file(cls, v: dict[str, float]) -> dict[str, float]:
         if sum(v.values()) != 1:
             raise ValueError("The fractions must sum to 1.")
-        return v
+        return v

src/data_designer/engine/processing/processors/to_jsonl.py

@@ -3,10 +3,10 @@
 
 import json
 import logging
+from pathlib import Path
 import tempfile
 
 import pandas as pd
-from pathlib import Path
 
 from data_designer.config.processors import ToJsonlProcessorConfig
 from data_designer.engine.configurable_task import ConfigurableTaskMetadata
@@ -62,4 +62,4 @@ def process(self, data: pd.DataFrame, *, current_batch_number: int | None = None
 
                 self.artifact_storage.move_to_outputs(Path(temp_dir) / filename, self.config.folder_name)
 
-        return data
+        return data