Skip to content

Commit f38c79a

Browse files
committed
Add a functional test to exercise our split out APIs
This adds a functional test to test the entire end-to-end flow of `ilab data generate` exercised purely from the individual pieces of the split out SDG APIs. No actual LLM inference happens, and instead we just mock out all the responses for the sake of testing speed / hardware since we don't need real LLM responses to verify our APIs. Signed-off-by: Ben Browning <[email protected]>
1 parent 8336e42 commit f38c79a

File tree

3 files changed

+156
-0
lines changed

3 files changed

+156
-0
lines changed

tests/functional/__init__.py

Whitespace-only changes.

tests/functional/test_granular_api.py

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
# SPDX-License-Identifier: Apache-2.0
2+
3+
# Standard
4+
from datetime import datetime
5+
from unittest.mock import MagicMock
6+
import glob
7+
import pathlib
8+
9+
# Third Party
10+
import git
11+
12+
# First Party
13+
from instructlab.sdg import BlockRegistry
14+
from instructlab.sdg.generate_data import (
15+
generate_taxonomy,
16+
mix_datasets,
17+
postprocess_taxonomy,
18+
preprocess_taxonomy,
19+
)
20+
21+
# Local
22+
from ..mockllmblock import MockLLMBlock
23+
24+
25+
def _clone_instructlab_taxonomy(taxonomy_dir):
26+
taxonomy_repo_url = "https://github.com/instructlab/taxonomy"
27+
taxonomy_commit = "dfa3afaf26f40f923cf758389719619ec9b1ddb1"
28+
repo = git.Repo.clone_from(taxonomy_repo_url, taxonomy_dir, no_checkout=True)
29+
repo.git.checkout(taxonomy_commit)
30+
31+
32+
def test_granular_api_end_to_end(testdata_path: pathlib.Path, tmp_path: pathlib.Path):
33+
# Registry our mock block so we can reference it in pipelines
34+
BlockRegistry.register("MockLLMBlock")(MockLLMBlock)
35+
36+
# Clone a taxonomy and edit 1 file in it
37+
taxonomy_dir = tmp_path.joinpath("taxonomy")
38+
_clone_instructlab_taxonomy(taxonomy_dir)
39+
changed_qna_yaml = taxonomy_dir.joinpath(
40+
"knowledge", "science", "animals", "birds", "black_capped_chickadee", "qna.yaml"
41+
)
42+
with open(changed_qna_yaml, "a", encoding="utf-8") as file:
43+
file.write("")
44+
45+
pipeline_dir = testdata_path.joinpath("mock_pipelines")
46+
date_suffix = datetime.now().replace(microsecond=0).isoformat().replace(":", "_")
47+
48+
preprocessed_dir = tmp_path.joinpath("preprocessed")
49+
preprocess_taxonomy(
50+
taxonomy_dir=taxonomy_dir,
51+
output_dir=preprocessed_dir,
52+
)
53+
chickadee_docs = glob.glob(
54+
str(
55+
preprocessed_dir.joinpath(
56+
"documents", "knowledge_science_*", "chickadee.md"
57+
)
58+
)
59+
)
60+
assert chickadee_docs
61+
chickadee_samples_path = preprocessed_dir.joinpath(
62+
"knowledge_science_animals_birds_black_capped_chickadee.jsonl"
63+
)
64+
assert chickadee_samples_path.is_file()
65+
66+
client = MagicMock()
67+
client.server_supports_batched = False
68+
generated_dir = tmp_path.joinpath("generated")
69+
generate_taxonomy(
70+
client=client,
71+
input_dir=preprocessed_dir,
72+
output_dir=generated_dir,
73+
pipeline=pipeline_dir,
74+
)
75+
generated_chickadee_samples_path = generated_dir.joinpath(
76+
"knowledge_science_animals_birds_black_capped_chickadee.jsonl"
77+
)
78+
assert generated_chickadee_samples_path.is_file()
79+
80+
postprocessed_dir = tmp_path.joinpath("postprocessed")
81+
postprocess_taxonomy(
82+
input_dir=generated_dir,
83+
output_dir=postprocessed_dir,
84+
date_suffix=date_suffix,
85+
pipeline=pipeline_dir,
86+
)
87+
knowledge_recipe_file = postprocessed_dir.joinpath(
88+
f"knowledge_recipe_{date_suffix}.yaml"
89+
)
90+
assert knowledge_recipe_file.is_file()
91+
skills_recipe_file = postprocessed_dir.joinpath(f"skills_recipe_{date_suffix}.yaml")
92+
assert skills_recipe_file.is_file()
93+
94+
mixed_skills_output_file = (
95+
f"{postprocessed_dir}/skills_train_msgs_{date_suffix}.jsonl"
96+
)
97+
mix_datasets(
98+
recipe_file=f"{postprocessed_dir}/skills_recipe_{date_suffix}.yaml",
99+
output_file=mixed_skills_output_file,
100+
)
101+
assert pathlib.Path(mixed_skills_output_file).is_file()

tests/mockllmblock.py

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
# SPDX-License-Identifier: Apache-2.0
2+
3+
# Standard
4+
import random
5+
import string
6+
7+
# Third Party
8+
from datasets import Dataset
9+
10+
# First Party
11+
from instructlab.sdg import LLMBlock
12+
13+
14+
def _random_string(size):
15+
return "".join(random.choices(string.ascii_lowercase, k=size))
16+
17+
18+
def _add_mocked_cols(sample, block_name):
19+
match block_name:
20+
case "gen_questions" | "gen_grounded_questions":
21+
sample["question"] = f"Is this a question {_random_string(8)}?"
22+
case "eval_questions" | "eval_grounded_questions":
23+
sample["evaluation"] = "This is an evaluation."
24+
sample["score"] = "1"
25+
case "gen_responses" | "gen_grounded_responses":
26+
sample["response"] = "This is a response."
27+
case "evaluate_qa_pair" | "evaluate_grounded_qa_pair":
28+
sample["evaluation"] = "This is an evaluation."
29+
sample["score"] = "2"
30+
case "gen_contexts":
31+
sample["context"] = f"This is a context {_random_string(8)}."
32+
case "gen_spellcheck":
33+
sample["spellcheck"] = sample["document"]
34+
case "gen_knowledge":
35+
sample["question"] = f"Is this a question {_random_string(8)}?"
36+
sample["response"] = "This is a response."
37+
case "eval_faithfulness_qa_pair":
38+
sample["explanation"] = "This is an explanation."
39+
sample["judgment"] = "YES"
40+
case "eval_relevancy_qa_pair":
41+
sample["feedback"] = "This is some feedback."
42+
sample["score"] = "2"
43+
case "eval_verify_question":
44+
sample["explanation"] = "This is an explanation."
45+
sample["rating"] = "1"
46+
case _:
47+
raise Exception(
48+
f"Received an un-mocked LLMBlock: {block_name}. Add code in {__file__} to handle this block."
49+
)
50+
return sample
51+
52+
53+
class MockLLMBlock(LLMBlock):
54+
def generate(self, samples: Dataset):
55+
return samples.map(_add_mocked_cols, fn_kwargs={"block_name": self.block_name})

0 commit comments

Comments
 (0)