Adding support for Classifiers and Search tools

.devcontainer/postInstall.sh

@@ -2,4 +2,4 @@
 
 PATH=/home/vscode/.cargo/bin:$PATH
 cd dolma
-source /home/vscode/miniforge3/bin/activate && pip install cmake "maturin[patchelf]>=1.1,<2.0"
+source /home/vscode/miniforge3/bin/activate && pip install cmake "maturin>=1.5,<2.0"

.github/workflows/CI.yml

@@ -19,6 +19,7 @@ permissions:
 env:
   DOLMA_TESTS_SKIP_AWS: ${{ secrets.AWS_ACCESS_KEY_ID == '' && 'true' || 'false' }}
   DOLMA_TEST_S3_PREFIX: s3://dolma-tests
+  DOLMA_TEST_SKIP_LARGE_MODELS: "true"
   RUST_CHANNEL: stable
 
 jobs:

classifiers/README.md

@@ -0,0 +1,30 @@
+# Dolma Classifiers
+
+
+## Getting Started
+
+From root directory, install the package:
+
+```bash
+pip install -e classifiers
+```
+
+## Examples
+
+Run [Huggingface FineWeb classifier](https://huggingface.co/HuggingFaceFW/fineweb-edu-classifier) on S3 data:
+
+```bash
+python -m dolma_classifiers.inference \
+    -s 's3://ai2-llm/pretraining-data/sources/dclm/v0/documents/40b-split/20b-01/*zstd' \
+    -m HuggingFaceFW/fineweb-edu-classifier
+```
+
+Run [NVIDIA's Deberta quality classifier](https://huggingface.co/nvidia/quality-classifier-deberta) on S3 data with model compilation:
+
+```bash
+python -m dolma_classifiers.inference \
+    -s 's3://ai2-llm/pretraining-data/sources/dclm/v0/documents/40b-split/*/*zstd' \
+    -m nvidia/quality-classifier-deberta \
+    --model-compile \
+    --max-length 1024
+```

classifiers/pyproject.toml

@@ -0,0 +1,107 @@
+[project]
+name = "dolma-classifiers"
+version = "0.1.0"
+description = "Toolkit for easy classification of data in Dolma format."
+authors = [
+    {name = "Luca Soldaini", email = "[email protected]" }
+]
+license = {text = "Apache-2.0"}
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "msgspec",
+    "fsspec[s3]",
+    "smart_open[s3]>=7.0.4",
+    "tqdm",
+    "torch",
+    "transformers",
+    "wandb",
+    "jq"
+]
+
+[project.urls]
+"Homepage" = "https://github.com/allenai/dolma"
+"Repository" = "https://github.com/allenai/dolma"
+"Bug Tracker" = "https://github.com/allenai/dolma/issues"
+
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.setuptools.package-data]
+dolma_classifiers = ["py.typed", "*.pyi"]
+
+
+[build-system]
+build-backend = "setuptools.build_meta"
+requires = [
+    "setuptools >= 61.0.0",
+    "wheel"
+]
+
+[project.optional-dependencies]
+dev = [
+    "black>=22.6.0",
+    "isort>=5.10.1",
+    "mypy>=0.971",
+    "pytest>=5.2",
+    "ipython>=8.4.0",
+    "autopep8>=1.7.0",
+    "flake8>=5.0",
+    "ipdb>=0.13.0",
+    "flake8-pyi>=22.8.1",
+    "Flake8-pyproject>=1.1.0",
+    "pytest-asyncio>=0.15.1",
+    "pytest-cov>=2.12.1",
+    "aioresponses>=0.7.2",
+]
+
+[tool.black]
+line-length = 115
+include = '\.pyi?$'
+exclude = '''
+(
+      __pycache__
+    | \.git
+    | \.mypy_cache
+    | \.pytest_cache
+    | \.vscode
+    | \.venv
+    | \bdist\b
+    | \bdoc\b
+)
+'''
+
+[tool.isort]
+profile = "black"
+line_length = 115
+multi_line_output = 3
+
+[tool.autopep8]
+max_line_length = 115
+in-place = true
+recursive = true
+aggressive = 3
+
+[tool.mypy]
+python_version = "3.10"
+ignore_missing_imports = true
+no_site_packages = true
+allow_redefinition = false
+warn_unused_configs = true
+warn_unused_ignores = true
+warn_no_return = true
+warn_return_any = false
+warn_unreachable = true
+show_error_codes = true
+pretty = true
+
+[tool.mypy-tests]
+strict_optional = false
+
+[tool.flake8]
+per-file-ignores = [
+    '__init__.py:F401',
+    '*.pyi:E302,E305',
+    '*.py:E203'
+]

classifiers/scripts/fineweb_100b.sh

@@ -0,0 +1,45 @@
+#! /bin/bash
+
+DOCUMENTS='s3://ai2-llm/pretraining-data/sources/dclm/v0/documents/100*/*.jsonl.zstd'
+
+NUM_NODES=2
+MODEL_NAME="HuggingFaceFW/fineweb-edu-classifier"
+CLUSTER="ai2/jupiter*"
+BATCH_SIZE=1024
+PRIORITY="high"
+
+# Generate a hash for the run name by combining model name and documents
+RUN_HASH=$(echo -n "${MODEL_NAME}${DOCUMENTS}" | md5sum | awk '{print $1}')
+RUN_NAME="fineweb_classifier_${RUN_HASH:0:8}"
+
+# Set the run name as an environment variable
+export BEAKER_EXPERIMENT_NAME="${RUN_NAME}"
+
+
+gantry run \
+    --task-name "${RUN_NAME}" \
+    --description "Score ${DOCUMENTS} with ${MODEL_NAME}" \
+    --allow-dirty \
+    --workspace ai2/davidw-oe-annealing \
+    --beaker-image 'petew/olmo-torch23-gantry' \
+    --timeout -1 \
+    --show-logs \
+    --host-networking \
+    --venv 'base' \
+    --priority "${PRIORITY}" \
+    --leader-selection \
+    --gpus 8 \
+    --replicas ${NUM_NODES} \
+    --preemptible \
+    --cluster "${CLUSTER}" \
+    --budget ai2/oe-data \
+    --env LOG_FILTER_TYPE=local_rank0_only \
+    --env OMP_NUM_THREADS=8 \
+    --env BEAKER_USER_ID=$(beaker account whoami --format json | jq '.[0].name' -cr) \
+    --env-secret AWS_ACCESS_KEY_ID=lucas-AWS_ACCESS_KEY_ID \
+    --env-secret AWS_SECRET_ACCESS_KEY=lucas-AWS_SECRET_ACCESS_KEY \
+    --env-secret WANDB_API_KEY=lucas-WANDB_API_KEY \
+    --shared-memory 10GiB \
+    --install "pip install -e classifiers/" \
+    --yes \
+    -- /bin/bash -c "huggingface-cli download ${MODEL_NAME} && torchrun --nnodes "${NUM_NODES}:${NUM_NODES}" --nproc-per-node 8 --rdzv_id 12347 --rdzv_backend static --rdzv_endpoint "\${BEAKER_LEADER_REPLICA_HOSTNAME}:29400" --node_rank "\${BEAKER_REPLICA_RANK}" --rdzv_conf 'read_timeout=3600' -m dolma_classifiers.inference --source-prefix ${DOCUMENTS} --batch-size ${BATCH_SIZE} --use-wandb --wandb-project 'dolma-classifiers' --wandb-entity ai2-llm --model-name ${MODEL_NAME} --num-workers 8 --prefetch-factor 8"

classifiers/scripts/fineweb_40b.sh

@@ -0,0 +1,45 @@
+#! /bin/bash
+
+DOCUMENTS='s3://ai2-llm/pretraining-data/sources/dclm/v0/documents/40b-split/*/*zstd'
+NUM_NODES=1
+BATCH_SIZE=1024
+CLUSTER="ai2/neptune*"
+PRIORITY="high"
+MODEL_NAME="HuggingFaceFW/fineweb-edu-classifier"
+
+
+# Generate a hash for the run name by combining model name and documents
+RUN_HASH=$(echo -n "${MODEL_NAME}${DOCUMENTS}" | md5sum | awk '{print $1}')
+RUN_NAME="fineweb_classifier_${RUN_HASH:0:8}"
+
+# Set the run name as an environment variable
+export BEAKER_EXPERIMENT_NAME="${RUN_NAME}"
+
+
+gantry run \
+    --task-name "${RUN_NAME}" \
+    --description "Score ${DOCUMENTS} with ${MODEL_NAME}" \
+    --allow-dirty \
+    --workspace ai2/davidw-oe-annealing \
+    --beaker-image 'petew/olmo-torch23-gantry' \
+    --timeout -1 \
+    --show-logs \
+    --host-networking \
+    --venv 'base' \
+    --priority "${PRIORITY}" \
+    --leader-selection \
+    --gpus 8 \
+    --replicas ${NUM_NODES} \
+    --preemptible \
+    --cluster "${CLUSTER}" \
+    --budget ai2/oe-data \
+    --env LOG_FILTER_TYPE=local_rank0_only \
+    --env OMP_NUM_THREADS=8 \
+    --env BEAKER_USER_ID=$(beaker account whoami --format json | jq '.[0].name' -cr) \
+    --env-secret AWS_ACCESS_KEY_ID=lucas-AWS_ACCESS_KEY_ID \
+    --env-secret AWS_SECRET_ACCESS_KEY=lucas-AWS_SECRET_ACCESS_KEY \
+    --env-secret WANDB_API_KEY=lucas-WANDB_API_KEY \
+    --shared-memory 10GiB \
+    --install "pip install -e classifiers/" \
+    --yes \
+    -- /bin/bash -c "huggingface-cli download ${MODEL_NAME} && torchrun --nnodes "${NUM_NODES}:${NUM_NODES}" --nproc-per-node 8 --rdzv_id 12347 --rdzv_backend static --rdzv_endpoint "\${BEAKER_LEADER_REPLICA_HOSTNAME}:29400" --node_rank "\${BEAKER_REPLICA_RANK}" --rdzv_conf 'read_timeout=3600' -m dolma_classifiers.inference --source-prefix ${DOCUMENTS} --batch-size ${BATCH_SIZE} --use-wandb --wandb-project 'dolma-classifiers' --wandb-entity ai2-llm --model-name ${MODEL_NAME} --num-workers 4 --prefetch-factor 8"

classifiers/scripts/fineweb_50b_extra.sh

@@ -0,0 +1,45 @@
+#! /bin/bash
+
+DOCUMENTS='s3://ai2-llm/pretraining-data/sources/dclm/v0/documents/20240909-50b/*zstd'
+NUM_NODES=1
+MODEL_NAME="HuggingFaceFW/fineweb-edu-classifier"
+CLUSTER="ai2/jupiter*"
+BATCH_SIZE=1024
+PRIORITY="high"
+
+
+# Generate a hash for the run name by combining model name and documents
+RUN_HASH=$(echo -n "${MODEL_NAME}${DOCUMENTS}" | md5sum | awk '{print $1}')
+RUN_NAME="fineweb_classifier_${RUN_HASH:0:8}"
+
+# Set the run name as an environment variable
+export BEAKER_EXPERIMENT_NAME="${RUN_NAME}"
+
+
+gantry run \
+    --task-name "${RUN_NAME}" \
+    --description "Score ${DOCUMENTS} with ${MODEL_NAME}" \
+    --allow-dirty \
+    --workspace ai2/davidw-oe-annealing \
+    --beaker-image 'petew/olmo-torch23-gantry' \
+    --timeout -1 \
+    --show-logs \
+    --host-networking \
+    --venv 'base' \
+    --priority "${PRIORITY}" \
+    --leader-selection \
+    --gpus 8 \
+    --replicas ${NUM_NODES} \
+    --preemptible \
+    --cluster "${CLUSTER}" \
+    --budget ai2/oe-data \
+    --env LOG_FILTER_TYPE=local_rank0_only \
+    --env OMP_NUM_THREADS=8 \
+    --env BEAKER_USER_ID=$(beaker account whoami --format json | jq '.[0].name' -cr) \
+    --env-secret AWS_ACCESS_KEY_ID=lucas-AWS_ACCESS_KEY_ID \
+    --env-secret AWS_SECRET_ACCESS_KEY=lucas-AWS_SECRET_ACCESS_KEY \
+    --env-secret WANDB_API_KEY=lucas-WANDB_API_KEY \
+    --shared-memory 10GiB \
+    --install "pip install -e classifiers/" \
+    --yes \
+    -- /bin/bash -c "huggingface-cli download ${MODEL_NAME} && torchrun --nnodes "${NUM_NODES}:${NUM_NODES}" --nproc-per-node 8 --rdzv_id 12347 --rdzv_backend static --rdzv_endpoint "\${BEAKER_LEADER_REPLICA_HOSTNAME}:29400" --node_rank "\${BEAKER_REPLICA_RANK}" --rdzv_conf 'read_timeout=3600' -m dolma_classifiers.inference --source-prefix ${DOCUMENTS} --batch-size ${BATCH_SIZE} --use-wandb --wandb-project 'dolma-classifiers' --wandb-entity ai2-llm --model-name ${MODEL_NAME} --num-workers 8 --prefetch-factor 8"

classifiers/scripts/fineweb_automath_arxiv.sh

@@ -0,0 +1,45 @@
+#! /bin/bash
+
+DOCUMENTS='s3://ai2-llm/pretraining-data/sources/math-ai_AutoMathText/v0/documents/arxiv/*/*.gz'
+
+NUM_NODES=1
+MODEL_NAME="HuggingFaceFW/fineweb-edu-classifier"
+CLUSTER="ai2/jupiter*"
+BATCH_SIZE=1024
+PRIORITY="urgent"
+
+# Generate a hash for the run name by combining model name and documents
+RUN_HASH=$(echo -n "${MODEL_NAME}${DOCUMENTS}" | md5sum | awk '{print $1}')
+RUN_NAME="fineweb_classifier_${RUN_HASH:0:8}"
+
+# Set the run name as an environment variable
+export BEAKER_EXPERIMENT_NAME="${RUN_NAME}"
+
+
+gantry run \
+    --task-name "${RUN_NAME}" \
+    --description "Score ${DOCUMENTS} with ${MODEL_NAME}" \
+    --allow-dirty \
+    --workspace ai2/davidw-oe-annealing \
+    --beaker-image 'petew/olmo-torch23-gantry' \
+    --timeout -1 \
+    --show-logs \
+    --host-networking \
+    --venv 'base' \
+    --priority "${PRIORITY}" \
+    --leader-selection \
+    --gpus 8 \
+    --replicas ${NUM_NODES} \
+    --preemptible \
+    --cluster "${CLUSTER}" \
+    --budget ai2/oe-data \
+    --env LOG_FILTER_TYPE=local_rank0_only \
+    --env OMP_NUM_THREADS=8 \
+    --env BEAKER_USER_ID=$(beaker account whoami --format json | jq '.[0].name' -cr) \
+    --env-secret AWS_ACCESS_KEY_ID=lucas-AWS_ACCESS_KEY_ID \
+    --env-secret AWS_SECRET_ACCESS_KEY=lucas-AWS_SECRET_ACCESS_KEY \
+    --env-secret WANDB_API_KEY=lucas-WANDB_API_KEY \
+    --shared-memory 10GiB \
+    --install "pip install -e classifiers/" \
+    --yes \
+    -- /bin/bash -c "huggingface-cli download ${MODEL_NAME} && torchrun --nnodes "${NUM_NODES}:${NUM_NODES}" --nproc-per-node 8 --rdzv_id 12347 --rdzv_backend static --rdzv_endpoint "\${BEAKER_LEADER_REPLICA_HOSTNAME}:29400" --node_rank "\${BEAKER_REPLICA_RANK}" --rdzv_conf 'read_timeout=3600' -m dolma_classifiers.inference --source-prefix ${DOCUMENTS} --batch-size ${BATCH_SIZE} --use-wandb --wandb-project 'dolma-classifiers' --wandb-entity ai2-llm --model-name ${MODEL_NAME} --num-workers 8 --prefetch-factor 8"

classifiers/scripts/fineweb_automath_code.sh

@@ -0,0 +1,45 @@
+#! /bin/bash
+
+DOCUMENTS='s3://ai2-llm/pretraining-data/sources/math-ai_AutoMathText/v0/documents/code/*/*.gz'
+
+NUM_NODES=1
+MODEL_NAME="HuggingFaceFW/fineweb-edu-classifier"
+CLUSTER="ai2/jupiter*"
+BATCH_SIZE=1024
+PRIORITY="urgent"
+
+# Generate a hash for the run name by combining model name and documents
+RUN_HASH=$(echo -n "${MODEL_NAME}${DOCUMENTS}" | md5sum | awk '{print $1}')
+RUN_NAME="fineweb_classifier_${RUN_HASH:0:8}"
+
+# Set the run name as an environment variable
+export BEAKER_EXPERIMENT_NAME="${RUN_NAME}"
+
+
+gantry run \
+    --task-name "${RUN_NAME}" \
+    --description "Score ${DOCUMENTS} with ${MODEL_NAME}" \
+    --allow-dirty \
+    --workspace ai2/davidw-oe-annealing \
+    --beaker-image 'petew/olmo-torch23-gantry' \
+    --timeout -1 \
+    --show-logs \
+    --host-networking \
+    --venv 'base' \
+    --priority "${PRIORITY}" \
+    --leader-selection \
+    --gpus 8 \
+    --replicas ${NUM_NODES} \
+    --preemptible \
+    --cluster "${CLUSTER}" \
+    --budget ai2/oe-data \
+    --env LOG_FILTER_TYPE=local_rank0_only \
+    --env OMP_NUM_THREADS=8 \
+    --env BEAKER_USER_ID=$(beaker account whoami --format json | jq '.[0].name' -cr) \
+    --env-secret AWS_ACCESS_KEY_ID=lucas-AWS_ACCESS_KEY_ID \
+    --env-secret AWS_SECRET_ACCESS_KEY=lucas-AWS_SECRET_ACCESS_KEY \
+    --env-secret WANDB_API_KEY=lucas-WANDB_API_KEY \
+    --shared-memory 10GiB \
+    --install "pip install -e classifiers/" \
+    --yes \
+    -- /bin/bash -c "huggingface-cli download ${MODEL_NAME} && torchrun --nnodes "${NUM_NODES}:${NUM_NODES}" --nproc-per-node 8 --rdzv_id 12347 --rdzv_backend static --rdzv_endpoint "\${BEAKER_LEADER_REPLICA_HOSTNAME}:29400" --node_rank "\${BEAKER_REPLICA_RANK}" --rdzv_conf 'read_timeout=3600' -m dolma_classifiers.inference --source-prefix ${DOCUMENTS} --batch-size ${BATCH_SIZE} --use-wandb --wandb-project 'dolma-classifiers' --wandb-entity ai2-llm --model-name ${MODEL_NAME} --num-workers 8 --prefetch-factor 8"