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 · 6c8544e · 6c8544e
1 parent 0f78465
commit 6c8544e
Show file tree

Hide file tree

Showing 2 changed files with 4 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 = {

diff --git a/tests/functional/test_granular_api.py b/tests/functional/test_granular_api.py
@@ -46,9 +46,11 @@ def test_granular_api_end_to_end(testdata_path: pathlib.Path, tmp_path: pathlib.
     date_suffix = datetime.now().replace(microsecond=0).isoformat().replace(":", "_")
 
     preprocessed_dir = tmp_path.joinpath("preprocessed")
+    teacher_model_path = testdata_path.joinpath("models/instructlab/granite-7b-lab")
     preprocess_taxonomy(
         taxonomy_dir=taxonomy_dir,
         output_dir=preprocessed_dir,
+        teacher_model_path=teacher_model_path,
     )
     chickadee_docs = glob.glob(
         str(