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?

[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

[RobotSail]

I think SDG has a logger class, should we be using that here instead of print?

[RobotSail]

It seems like __DEVICE is only used to pick a default for this function. I would recommend using the Python convention of setting this to None as a default, and then falling back to the initialization as it was via:

if not device:
  device = "cuda" if torch.cuda.is_available() else "cpu"

Also - I recommend going through torch.device rather than using a str, but if it works then it works.

[eshwarprasadS]

Oh definitely! good catch, thanks for pointing these out.

[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.