Add pipeline option to generate_data()

v0.0.5 will include support for other pipeline types. This trivial change is to update the `generate_data()` arguments ahead of time so we can get a CLI option in place to set this option prior to new types being made available. Signed-off-by: Russell Bryant <[email protected]>
instructlab · Jun 30, 2024 · 499c732 · 499c732
1 parent 78d54f5
commit 499c732
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/instructlab/sdg/generate_data.py b/src/instructlab/sdg/generate_data.py
@@ -468,11 +468,15 @@ def generate_data(
     tls_client_cert: Optional[str] = None,
     tls_client_key: Optional[str] = None,
     tls_client_passwd: Optional[str] = None,
+    pipeline: Optional[str] = None,
 ):
     seed_instruction_data = []
     machine_seed_instruction_data = []
     generate_start = time.time()
 
+    if pipeline and pipeline != "simple":
+        raise SystemExit("Error: Only 'simple' pipeline is supported.")
+
     if not os.path.exists(output_dir):
         os.mkdir(output_dir)