We'll need to also rename tests/test_chunking.py to tests/test_docprocessor.py and update its tests to pass in the new expected values to the chunk_documents method.

Do we actually need pypdf here? We're using it to read the contents of the PDF files, but then we're never actually using the contents of PDF files downstream. If we really want to read the contents of the PDF, we could use something like https://github.com/DS4SD/docling-parse instead as that's already a dependency of docling anyway, so won't be bringing in anything extra?

Should we align this with instructlab/instructlab? Currently that's on 4.41.2 at https://github.com/instructlab/instructlab/blob/main/requirements.txt#L34

We're missing an entry for docling here to get docling itself added as a requirement. And, since docling depends on torch, we'll want to also add a dependency on torch to pin its version to the same range as instructlab/instructlab - https://github.com/instructlab/instructlab/blob/main/requirements.txt#L31 (currently torch>=2.3.0,<2.4.0)

Here's a suggested change to requirements.txt that adds docling and aligns a couple of dependencies with the CLI repo:

diff --git a/requirements.txt b/requirements.txt
index 195dddf..84bb75c 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,17 +1,20 @@
 # SPDX-License-Identifier: Apache-2.0
 click>=8.1.7,<9.0.0
 datasets>=2.18.0,<3.0.0
+docling>=1.15.0,<2.0.0
 GitPython>=3.1.42,<4.0.0
 httpx>=0.25.0,<1.0.0
 instructlab-schema>=0.4.0
 langchain-text-splitters
 openai>=1.13.3,<2.0.0
+pypdf>=5.0.0
+tabulate>=0.9.0
 # Note: this dependency goes along with langchain-text-splitters and may be
 #       removed once that one is removed.
 # do not use 8.4.0 due to a bug in the library
 # https://github.com/instructlab/instructlab/issues/1389
-pypdf>=5.0.0
-tabulate>=0.9.0
 tenacity>=8.3.0,!=8.4.0
-transformers>=4.44.2
+# align torch with instructlab/instructlab
+torch>=2.3.0,<2.4.0
+transformers>=4.41.2
 xdg-base-dirs>=6.0.1

That would be enough to at least get some of these tests running, although until test_chunking.py is adjusted it's going to just error on every run.

This path should probably be somewhere under the output_dir from generate_data.py. Since we'll create new docs here for every run, perhaps under the node_datasets_* subdirectory for that run?

This should come from the model_name parameter passed to generate_data.py's generate_data. That way we load the tokenizer for the same teacher model that's being used.

Prefer os.path.join to combine paths to string concatenation - os.path.join handles various edge cases like ensuring we don't add double slashes and more complex paths.

Prefer os.path.join to combine paths to string concatenation - os.path.join handles various edge cases like ensuring we don't add double slashes and more complex paths.

Where do we add the icl values for non-PDF files? This codepath only happens for PDFs, right?

Ok, digging into this a bit more, we shouldn't be doing anything different for icl values with PDFs, right? I'm confused why we do anything with icl_query_* here at all, as that all gets handled in taxonomy.py already. What's the intent behind the separate duplicate icl handling for pdfs?