change package name to data_selection

Author: sangmichaelxie

README.md

@@ -1,6 +1,57 @@
 # Data Selection for Language Models via Importance Resampling (DSIR)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+[![arXiv](https://img.shields.io/badge/arXiv-2305.10429-00ff00.svg)](https://arxiv.org/abs/2302.03169)
 
-This repository contains pre-filtered datasets and code for selecting relevant language model training data from The Pile.
+This repository contains the [DSIR](https://arxiv.org/abs/2302.03169) data selection tool for selecting relevant language model training data from any raw data source given a target dataset, as well as pre-filtered datasets and some pretrained models.
+
+DSIR is built for:
+- fast, large-scale (trillion-token scale) data selection from large raw text datasets (Pile, RefinedWeb, RedPajama, ...)
+- selecting data that is distributed like a given target dataset (domain-specific data, Wikipedia, ...). Relevance and diversity are balanced automatically.
+
+Compute needed:
+- 1 CPU node
+- a large amount of RAM (at least few hundred GB)
+- a high number of cores (parallelism on a file level. For best performance, use as many CPU cores as data files)
+
+![DSIR figure](fig1.png)
+
+Code related to the DSIR paper's experiments are in the `experimental/` directory.
+
+## Quickstart
+
+Install from pip:
+```
+pip install data-selection
+```
+
+Install from source by cloning this repo and installing via pip:
+```
+git clone [email protected]:/p-lambda/dsir
+pip install ./dsir
+```
+
+To select data, initialize a `HashedNgramDSIR` object and call the following functions:
+```
+from data_selection import HashedNgramDSIR
+
+raw_datasets = [<list of paths>]
+target_datasets = [<list of paths>]
+
+dsir = HashedNgramDSIR(raw_datasets, num_proc=30)
+dsir.fit_importance_estimator(target_datasets)
+dsir.compute_importance_weights()
+dsir.resample(out_dir='resampled', num_to_sample=1000000, cache_dir='/scr/resampled_cache')
+```
+This will save 1M examples in `jsonl` files inside an output directory named `resampled`. The files will first be written to `cache_dir` and moved to `out_dir` upon completion.
+
+The `dsir` intermediate results (after `fit_importance_estimator` and `compute_importance_weights`) can be saved and loaded for later use, for example to resample a different number of examples:
+```
+dsir.save('dsir_params')
+
+# later on
+dsir.load('dsir_params')
+dsir.resample(out_dir='resampled', num_to_sample=10000000, cache_dir='/scr/resampled_cache')
+```
 
 ## Pre-filtered datasets
 Note: previous versions of the datasets had a small validation and test split (50000 examples each), but we concatenated these onto the end of the train set (in the order validation, then test) to better align with the paper. The datasets should be further shuffled during preprocessing before training.
@@ -52,32 +103,12 @@ In the table below, `{dataset}` can be replaced with one of `{ag, amazon, citati
 | heuristiccls-roberta-continuedpretrain-{dataset} | Link format: `https://huggingface.co/sangmichaelxie/dsir-roberta-continuedpretrain-{dataset}` | 6.4B tokens (25M examples) | 256 | 25000 | roberta-base | roberta-base | RoBERTa model with continued pretraining on data selected by heurstic classification with target={dataset} |
 | randomselect-roberta-continuedpretrain | [Link](https://huggingface.co/sangmichaelxie/randomselect-roberta-continuedpretrain) | 6.4B tokens (25M examples) | 256 | 25000 | roberta-base | roberta-base | RoBERTa model with continued pretraining on random subset of The Pile |
 
-## Code for data selection
-
-To select your own subset of The Pile, all you need is a small set of target examples representing the kind of data you want to select.
-This target dataset should be in jsonl format -- it can also be a dataset from HuggingFace Datasets. Note that our current workflow requires about 2TB of storage space --- we're working on reducing this! All the code should be run from the outer `dsir` directory.
-1. Create a virtualenv using `requirements.txt`: `virtualenv .venv; source .venv/bin/activate; pip install -r requirements.txt`
-2. Download The Pile to `PILE_PATH` and change the corresponding variables in `config.sh`.
-3. Run preprocessing on The Pile: Run `bash preprocessing/run_slurm.sh`. You can also run `bash preprocessing/run.sh` directly using the arguments in `preprocessing/run_slurm.sh`. This only needs to be run once. 
-4. Precompute quality filter stats: Run `bash preprocessing/quality_scores/run_slurm_quality_stats.sh`. After this, run `bash preprocessing/quality_scores/run_merge_quality_scores.sh`. This only needs to be run once. (We're working on streamlining steps 3 and 4. Stay tuned!) 
-5. Run DSIR: For an example, run `bash data_selection/run_cmds.sh`. For new target datasets, some information about which fields in the dataset to use should be placed in the `dsname_to_args` dictionary at the top of the `data_selection/dsir_pipeline.py` file. If you wish to retrieve from custom subsets of the Pile (for example, only select data from one chunk of the Pile), you will need to tweak one part of the code, in the main part of the script (an example is provided of how to do so as a comment). Many of the steps in DSIR can be cached and will only run the first time. For example, resampling a different number of examples with the same target dataset uses cached importance weights.
-
-## Code for pretraining and GLUE evaluation
-
-We provide scripts for training BERT-style masked language models on the selected data and evaluating it on GLUE in the `train` and `glue_eval` directories, respectively. All code should be run from the outer `dsir` directory.
-1. Install further dependencies using `train/requirements.txt`: `pip install -r train/requirements.txt`
-2. Change the `PRETRAIN_OUTPUT_DIR` variable in `config.sh`.
-3. Write a job command in `train/run_slurm.sh`. An example command in this file. You will need to change the path to the training data. If you want to skip preprocessing (if it's already done), set the first of two boolean variables to `false`. By setting both to `true`, there will be two jobs launched: one for preprocessing and one for pretraining. The pretraining job should take about 50 hours on 4 RTX 3090 GPUs. Kick off the jobs by running `bash train/run_slurm.sh`.
-4. Evaluate the trained model by editing the evaluation job command in `glue_eval/run_eval_exps.sh` with the path to the model checkpoint. This script runs 5 seeds for each GLUE dataset. The results and finetuned models will be saved a new `finetune_runs` directory inside the pretrained model checkpoint directory. Kick off the jobs by running `bash glue_exps/run_eval_exps.sh`.
-5. Read the GLUE results by running `python read_glue_results.py --results_dir </path/to/checkpoint>/finetune_runs` in the `glue_eval` directory.
-
-
 ## Citation Information
 Paper: <https://arxiv.org/abs/2302.03169>
 ```
 @article{xie2023data,
   author = {Sang Michael Xie and Shibani Santurkar and Tengyu Ma and Percy Liang},
-  journal = {arXiv preprint arXiv:2302.03169},
+  journal = {Advances in Neural Information Processing Systems (NeurIPS)},
   title = {Data Selection for Language Models via Importance Resampling},
   year = {2023},
 }

data_selection/__init__.py

@@ -0,0 +1,2 @@
+from .base import DSIR
+from .hashed_ngram_dsir import HashedNgramDSIR

dsir/base.py data_selection/base.py

@@ -9,7 +9,7 @@
 from datasets import load_dataset
 from tqdm import tqdm
 
-from dsir.utils import parallelize
+from data_selection.utils import parallelize
 
 
 def default_load_dataset_fn(path: str) -> Iterable[Dict]:

dsir/hashed_ngram_dsir.py data_selection/hashed_ngram_dsir.py

@@ -10,8 +10,8 @@
 from nltk import ngrams as get_ngrams
 import numpy as np
 
-from dsir.base import DSIR, default_load_dataset_fn, default_parse_example_fn
-from dsir.utils import parallelize
+from data_selection.base import DSIR, default_load_dataset_fn, default_parse_example_fn
+from data_selection.utils import parallelize
 
 
 wpt = WordPunctTokenizer()

dsir/utils.py data_selection/utils.py

@@ -0,0 +1,33 @@
+# Code for the DSIR paper
+This directory has the code for preprocessing, data selection, pretraining, and fine-tuning for the experiments in the DSIR paper. Pre-filtered datasets and pre-trained models from the paper are linked in the README at the outer directory.
+
+## Code for data selection
+
+To select your own subset of The Pile, all you need is a small set of target examples representing the kind of data you want to select.
+This target dataset should be in jsonl format -- it can also be a dataset from HuggingFace Datasets. Note that our current workflow requires about 2TB of storage space --- we're working on reducing this! All the code should be run from the `experimental/` directory.
+1. Create a virtualenv using `requirements.txt`: `virtualenv .venv; source .venv/bin/activate; pip install -r requirements.txt`
+2. Download The Pile to `PILE_PATH` and change the corresponding variables in `config.sh`.
+3. Run preprocessing on The Pile: Run `bash preprocessing/run_slurm.sh`. You can also run `bash preprocessing/run.sh` directly using the arguments in `preprocessing/run_slurm.sh`. This only needs to be run once. 
+4. Precompute quality filter stats: Run `bash preprocessing/quality_scores/run_slurm_quality_stats.sh`. After this, run `bash preprocessing/quality_scores/run_merge_quality_scores.sh`. This only needs to be run once. (We're working on streamlining steps 3 and 4. Stay tuned!) 
+5. Run DSIR: For an example, run `bash data_selection/run_cmds.sh`. For new target datasets, some information about which fields in the dataset to use should be placed in the `dsname_to_args` dictionary at the top of the `data_selection/dsir_pipeline.py` file. If you wish to retrieve from custom subsets of the Pile (for example, only select data from one chunk of the Pile), you will need to tweak one part of the code, in the main part of the script (an example is provided of how to do so as a comment). Many of the steps in DSIR can be cached and will only run the first time. For example, resampling a different number of examples with the same target dataset uses cached importance weights.
+
+## Code for pretraining and GLUE evaluation
+
+We provide scripts for training BERT-style masked language models on the selected data and evaluating it on GLUE in the `train` and `glue_eval` directories, respectively. All code should be run from the `experimental/` directory.
+1. Install further dependencies using `train/requirements.txt`: `pip install -r train/requirements.txt`
+2. Change the `PRETRAIN_OUTPUT_DIR` variable in `config.sh`.
+3. Write a job command in `train/run_slurm.sh`. An example command in this file. You will need to change the path to the training data. If you want to skip preprocessing (if it's already done), set the first of two boolean variables to `false`. By setting both to `true`, there will be two jobs launched: one for preprocessing and one for pretraining. The pretraining job should take about 50 hours on 4 RTX 3090 GPUs. Kick off the jobs by running `bash train/run_slurm.sh`.
+4. Evaluate the trained model by editing the evaluation job command in `glue_eval/run_eval_exps.sh` with the path to the model checkpoint. This script runs 5 seeds for each GLUE dataset. The results and finetuned models will be saved a new `finetune_runs` directory inside the pretrained model checkpoint directory. Kick off the jobs by running `bash glue_exps/run_eval_exps.sh`.
+5. Read the GLUE results by running `python read_glue_results.py --results_dir </path/to/checkpoint>/finetune_runs` in the `glue_eval` directory.
+
+## Citation Information
+Paper: <https://arxiv.org/abs/2302.03169>
+```
+@article{xie2023data,
+  author = {Sang Michael Xie and Shibani Santurkar and Tengyu Ma and Percy Liang},
+  journal = {arXiv preprint arXiv:2302.03169},
+  title = {Data Selection for Language Models via Importance Resampling},
+  year = {2023},
+}
+```
+

dsir/__init__.py

@@ -0,0 +1,13 @@
+#!/bin/bash
+
+CACHE='/path/to/cachedir'
+ROOT_DIR='/path/to/dsir/experimental'
+VIRTUAL_ENV='/path/to/.env'
+PILE_PATH='/path/to/pile'
+DSIR_OUTPUT_DIR='/path/to/outputdir'
+PRETRAIN_OUTPUT_DIR='/path/to/model_outputdir'
+WORD_VECTORS_PATH='/path/to/pretrained_fasttext_wordvecs.vec'
+# Slurm
+cluster_info='--partition <PARTITION_NAME>'
+
+source ${VIRTUAL_ENV}/bin/activate

experimental/README.md

@@ -0,0 +1,34 @@
+[build-system]
+requires      = ["setuptools>=61.0.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "data-selection"
+version = "0.0.1"
+authors = [
+  { name="Sang Michael Xie", email="[email protected]" },
+]
+description = "Data Selection with Importance Resampling"
+readme = "README.md"
+requires-python = ">=3.6"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+
+license = { file = "LICENSE" }
+keywords = ["data selection", "importance resampling", "dsir", "nlp", "language models"]
+dependencies = [
+    'numpy>=1.21.6',
+    'datasets>=2.13.2',
+    'tqdm>=4.62.3',
+    'joblib>=1.1.0',
+    'nltk>=3.7',
+]
+
+[project.optional-dependencies]
+dev = ["pytest"]
+
+[project.urls]
+"Homepage" = "https://github.com/p-lambda/dsir"

experimental/config.sh

@@ -1,17 +1,18 @@
 from setuptools import setup, find_packages
 
-setup(name='dsir',
-      version='0.0.1',
-      description='Data Selection with Importance Resampling',
-      url='https://github.com/p-lambda/dsir',
-      author='Sang Michael Xie',
-      author_email='[email protected]',
-      packages=find_packages('.'),
-      install_requires=[
-        'numpy>=1.21.6',
-        'datasets>=2.13.2',
-        'tqdm>=4.62.3',
-        'joblib>=1.1.0',
-        'nltk>=3.7',
-      ]
-)
+if __name__ == "__main__":
+    setup(name='data-selection',
+          version='0.0.1',
+          description='Data Selection with Importance Resampling',
+          url='https://github.com/p-lambda/dsir',
+          author='Sang Michael Xie',
+          author_email='[email protected]',
+          packages=find_packages('.'),
+          install_requires=[
+            'numpy>=1.21.6',
+            'datasets>=2.13.2',
+            'tqdm>=4.62.3',
+            'joblib>=1.1.0',
+            'nltk>=3.7',
+          ]
+    )

requirements.txt experimental/requirements.txt

@@ -4,7 +4,7 @@
 import json
 import shutil
 
-from dsir.hashed_ngram_dsir import HashedNgramDSIR, hash_buckets, get_ngram_counts
+from data_selection.hashed_ngram_dsir import HashedNgramDSIR, hash_buckets, get_ngram_counts
 
 
 toy_dataset = Path(__file__).parent / "toy_pile_data.jsonl"
@@ -147,4 +147,11 @@ def test_save_load(dsir_obj):
     assert dsir_obj_2.ngrams == dsir_obj.ngrams
 
 if __name__ == "__main__":
-    test_get_ngram_counts()
+    dsir = HashedNgramDSIR(
+            raw_datasets,
+            parse_example_fn=parse_example_fn,
+            num_proc=2,
+            ngrams=2,
+            num_buckets=10000)
+
+    test_resample(dsir)

fig1.png

@@ -1,4 +1,4 @@
-from dsir import utils
+from data_selection import utils
 
 
 def job(arg):