Subset Selection Integration

[eshwarprasadS]

Minimal Implementation of Subset Selection

This resolves #541

Important Incoming Changes

• Added h5py>=3.12.1 to requirements
• Added submodlib source to requirements
• Added subset_selection.py
• Added encoders dir housing bge and arctic encoders
• Added subset_selection_utils.py under src/instructlab/sdg/utils/
• Support for snowflake-arctic-embed-l-v2.0 enabled

[RobotSail]

This is what we were doing previously with GPTDolomite and Dolomite Engine. We will likely need to move this to live in its own repo as we did for GPTDolomite so we can do builds.

[eshwarprasadS]

We have contacted the authors and we expect submodlib to have a public release soon, so, I think we can fall back to moving it into its own repo, if that does not go through?

[khaledsulayman]

@RobotSail is it necessary to move it to its own repo? We're already resorting to forking over some code from another repo in this PR, so I'd hesitate to bring code in from another source that we're also committing to maintain

[RobotSail]

@khaledsulayman We would likely maintain it to the same extent that we do for GPTDolomite -- i.e., it's a temporary workaround that we can get rid of once there is a more official and steady package to depend on. Otherwise, if that never happens, we can always just copy the new code over into the repo if we wanted to.

[tiran]

PyPI does not accept wheels that have a dependency on a package with a link. You will not be able to upload SDG wheels with this change.

[dhellmann]

We are planning to add a check to the downstream builds that would prevent a dependency with a git URL like this from building, because we can't ensure that it is safe to pull code from an arbitrary git branch. Please work with the submodlib authors to produce a real release of that code so we do not have to express the dependency using a git URL.

[bbrowning]

Has this been resolved and can we replace this dependency with a real release yet? If not, is there progress on this in that upstream repository?

[abhi1092]

@eshwarprasadS can we also add a functional test? Something like this.

#!/usr/bin/env python3
from instructlab.sdg.subset_selection import subset_datasets


if __name__ == "__main__":
    dataset_files = ["<dataset_path>"]
    subset_config = {
    "instruction": "conversation",
    "query_description": "conversation",
    "templates": {
      "message": "{% for msg in messages if msg.role != 'system' %}{{ msg.role }}: {{ msg.content }}\n{% endfor %}",
        "conversation": "{% for conv in conversation %}{{ conv.from }}: {{ conv.value }}\n{% endfor %}",
    },
    "batch_size": 100000,
    "num_folds": 25
    ,
    "subset_sizes": [0.97],
    "seed": 42,
    "template_name": "text",
    "combine_files": False,
    "encoder_type": "bge",
    "encoder_model": "BAAI/bge-base-en"
  }
    subset_datasets(input_files=dataset_files, **subset_config)

Also, are you planning on adding the artic snowflake model?

[eshwarprasadS]

@eshwarprasadS can we also add a functional test? Something like this.

#!/usr/bin/env python3
from instructlab.sdg.subset_selection import subset_datasets


if __name__ == "__main__":
    dataset_files = ["<dataset_path>"]
    subset_config = {
    "instruction": "conversation",
    "query_description": "conversation",
    "templates": {
      "message": "{% for msg in messages if msg.role != 'system' %}{{ msg.role }}: {{ msg.content }}\n{% endfor %}",
        "conversation": "{% for conv in conversation %}{{ conv.from }}: {{ conv.value }}\n{% endfor %}",
    },
    "batch_size": 100000,
    "num_folds": 25
    ,
    "subset_sizes": [0.97],
    "seed": 42,
    "template_name": "text",
    "combine_files": False,
    "encoder_type": "bge",
    "encoder_model": "BAAI/bge-base-en"
  }
    subset_datasets(input_files=dataset_files, **subset_config)

Also, are you planning on adding the artic snowflake model?

Thanks @abhi1092 for the comments. Yes, I am planning to add a host of unit and functional tests in the coming week for the feature, thanks for the suggestion.

And yes, I am planning to incorporate the arctic-snowflake encoder class in encoders.py and change default encoder to that one. I seem to have gotten hold of the implementation for that from the author, so we should have it in here soon.

[RobotSail]

I'll need to go through in further detail but here are some comments so far

[eshwarprasadS]

I'll need to go through in further detail but here are some comments so far

Thanks a lot for looking at this! Good parts of this implementation was adapted directly from the upstream repository, so this really helps to have extra eyes on it.

[bbrowning]

Note that submodlib does not build properly - see the e2e-medium job logs, with a snippet copied here:

  writing manifest file '/home/tmp/pip-modern-metadata-ljrr7qrp/submodlib_py.egg-info/SOURCES.txt'
  creating '/home/tmp/pip-modern-metadata-ljrr7qrp/submodlib_py-0.0.1.dist-info'
  Preparing metadata (pyproject.toml): finished with status 'done'
  WARNING: Generating metadata for package submodlib produced metadata for project name submodlib-py. Fix your #egg=submodlib fragments.
Discarding git+https://github.com/decile-team/submodlib.git: Requested submodlib-py from git+https://github.com/decile-team/submodlib.git (from instructlab-sdg==0.7.1.dev124) has inconsistent name: expected 'submodlib', but metadata has 'submodlib-py'

...

ERROR: Could not find a version that satisfies the requirement submodlib (unavailable) (from instructlab-sdg) (from versions: none)
ERROR: No matching distribution found for submodlib (unavailable)

[bbrowning]

I briefly looked at the failing Mac tests, and it looks like submodlib didn't release anything that works with Macs. Their only release is binary-only and I only see wheels that will work for some x86_64 linux systems. I'm not sure how they created or uploaded this release, but it doesn't appear to be done in the typical way where pypi can build from source for systems that don't have precompiled wheels. We either need to change this back to building from source OR adjust this requirement to only pull in submodlib for linux systems.

[eshwarprasadS]

I briefly looked at the failing Mac tests, and it looks like submodlib didn't release anything that works with Macs. Their only release is binary-only and I only see wheels that will work for some x86_64 linux systems. I'm not sure how they created or uploaded this release, but it doesn't appear to be done in the typical way where pypi can build from source for systems that don't have precompiled wheels. We either need to change this back to building from source OR adjust this requirement to only pull in submodlib for linux systems.

Yes @bbrowning, I just realized that the failures are localized to MacOS. I was thinking along the lines of the python versions, since we had an issue with the wheels only working for 3.10 last week.

[bbrowning]

I adjusted the submodlib requirement to only get pulled in for Linux, since that's the only OS its pypi binary is built for. And, I marked the functional test that uses it as one that requires GPU, since the test uses cuda stuff directly. That got CI green, so things should be good to merge if the others involved are comfortable with this.

[aakankshaduggal]

Thanks, lgtm!

[eshwarprasadS]

Thanks @bbrowning for adjusting the requirements. Thanks everyone for the reviews along the way as well.

[RobotSail]

Congrats on getting this in @eshwarprasadS !!