Parameterize # of seed examples when converting to test data

Instead of hardcoding this to always be 3, add a parameter with a default of 3 when converting our seed examples to the test output dataset. Co-authored-by: Aakanksha Duggal <[email protected]> Signed-off-by: Ben Browning <[email protected]>
instructlab · Jan 15, 2025 · 69dd254 · 69dd254
1 parent f84902f
commit 69dd254
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/instructlab/sdg/generate_data.py b/src/instructlab/sdg/generate_data.py
@@ -123,9 +123,9 @@ def _gen_train_data(
     jldump(messages_data, output_file_messages)
 
 
-def _knowledge_seed_example_to_test_data(seed_example, system_prompt):
+def _knowledge_seed_example_to_test_data(seed_example, system_prompt, num_iterations=3):
     res = []
-    for i in range(3):
+    for i in range(num_iterations):
         idx = i + 1
         user = seed_example[f"icl_query_{idx}"] + "\n" + seed_example["icl_document"]
         test_sample = {