Warc Backoff

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

Makefile

@@ -1,5 +1,29 @@
+UNAME := $(shell uname)
+
+ifeq ($(UNAME), Darwin)
+    OS_MESSAGE := "MacOS detected"
+	CMAKE_SETUP := "which cmake || brew install cmake"
+	PROTOBUF_SETUP := "which protoc || brew install protobuf"
+	OPENSSL_SETUP := "which openssl || brew install openssl"
+else ifeq ($(UNAME), Linux)
+	OS_MESSAGE := "Linux detected"
+	CMAKE_SETUP := "which cmake || sudo apt-get install --yes build-essential cmake"
+	PROTOBUF_SETUP := "which protoc || sudo apt-get install --yes protobuf-compiler"
+	OPENSSL_SETUP := "which openssl || sudo apt-get install --yes libssl-dev"
+else
+	OS_MESSAGE := "Unsupported OS; please install rust, cmake, protobuf, and openssl manually"
+	CMAKE_SETUP := ""
+	PROTOBUF_SETUP := ""
+	OPENSSL_SETUP := ""
+endif
+
 setup:
-	@./setup.sh
+	@echo "${OS_MESSAGE}: installing..."
+	$(shell "${CMAKE_SETUP}")
+	$(shell "${PROTOBUF_SETUP}")
+	$(shell "${OPENSSL_SETUP}")
+	which cargo || curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
+	which maturin || pip install 'maturin>=1.5,<2.0'
 
 publish:
 	maturin publish

configs/cc-news/dedupe-month.sh

@@ -0,0 +1,73 @@
+#! /usr/bin/env bash
+
+# documents:
+#   - s3://ai2-llm/pretraining-data/sources/c4/v0/documents/train/*.gz
+
+# dedupe:
+#   name: dedupe_para_ngrams_13_1
+#   paragraphs:
+#     attribute_name: dedupe_para_ngrams_13_1
+#     by_ngram:
+#       ngram_length: 13
+#       stride: 1
+#       overlap_threshold: 0.5
+#   skip_empty: true
+
+# bloom_filter:
+#   file: ${oc.env:HOME}/c4_dedupe_para_ngrams_13_1.bin
+#   read_only: false
+#   # estimated doc count is obtained by counting number of words in paragraphs
+#   # then dividing by 13 (ngram_length) and multiplying by 2 (for each ngram)
+#   estimated_doc_count: 359_916_731_334
+#   desired_false_positive_rate: 0.1
+
+# processes: 188
+# work_dir:
+#   input: /tmp/c4_dedupe_para_ngrams_13_1/input
+#   output: /tmp/c4_dedupe_para_ngrams_13_1/output
+
+# run years between 2016 and 2024
+for year in {2016..2024}; do
+    # run months between 1 and 12
+    for month in {1..12}; do
+        # skip months after 7 if year is 2024
+        if [ $year -eq 2024 ] && [ $month -gt 7 ]; then
+            continue
+        fi
+
+        # skip months before 8 if year is 2016
+        if [ $year -eq 2016 ] && [ $month -lt 8 ]; then
+            continue
+        fi
+
+        # rename month to 2 digits
+        month=$(printf "%02d" $month)
+
+        documents="s3://ai2-llm/pretraining-data/sources/cc-news/v0-resiliparse/documents/${year}-${month}/*.zst"
+
+        size=$(aws s3api list-objects --bucket ai2-llm --prefix "pretraining-data/sources/cc-news/v0-resiliparse/documents/${year}-${month}/" --output json --query "[sum(Contents[].Size)]" | jq '.[0]' -rc)
+
+        # run deduplication
+        echo "Running fuzzy dedupe for ${year}-${month} with ${size} bytes Bloom filter"
+
+        set -ex
+
+        dolma dedupe \
+            --documents ${documents} \
+            --dedupe.name dedupe_ngrams_13_1 \
+            --dedupe.paragraphs.attribute_name dedupe_ngrams_13_1 \
+            --dedupe.paragraphs.by_ngram.ngram_length 13 \
+            --dedupe.paragraphs.by_ngram.stride 1 \
+            --dedupe.paragraphs.by_ngram.overlap_threshold 0.5 \
+            --dedupe.skip_empty \
+            --bloom_filter.file "${HOME}/cc-news/dedupe_ngrams_13_1-${year}-${month}.bin" \
+            --no-bloom_filter.read_only \
+            --bloom_filter.estimated_doc_count $size \
+            --bloom_filter.desired_false_positive_rate 0.01 \
+            --processes "$(expr $(nproc) - 4)" \
+            --work_dir.input /tmp/cc-news/dedupe_ngrams_13_1/${year}-${month}/input \
+            --work_dir.output /tmp/cc-news/dedupe_ngrams_13_1/${year}-${month}/output
+
+        set +ex
+    done
+done

configs/cc-news/dedupe-year.sh

@@ -0,0 +1,114 @@
+#! /usr/bin/env bash
+
+# documents:
+#   - s3://ai2-llm/pretraining-data/sources/c4/v0/documents/train/*.gz
+
+# dedupe:
+#   name: dedupe_para_ngrams_13_1
+#   paragraphs:
+#     attribute_name: dedupe_para_ngrams_13_1
+#     by_ngram:
+#       ngram_length: 13
+#       stride: 1
+#       overlap_threshold: 0.5
+#   skip_empty: true
+
+# bloom_filter:
+#   file: ${oc.env:HOME}/c4_dedupe_para_ngrams_13_1.bin
+#   read_only: false
+#   # estimated doc count is obtained by counting number of words in paragraphs
+#   # then dividing by 13 (ngram_length) and multiplying by 2 (for each ngram)
+#   estimated_doc_count: 359_916_731_334
+#   desired_false_positive_rate: 0.1
+
+# processes: 188
+# work_dir:
+#   input: /tmp/c4_dedupe_para_ngrams_13_1/input
+#   output: /tmp/c4_dedupe_para_ngrams_13_1/output
+
+# run years between 2016 and 2024
+for year in {2016..2024}; do
+    # run months between 1 and 12
+
+    # Initialize an empty array to store document paths and a variable for total size
+    documents=()
+    size=0
+
+    # Collect all month document paths into the array and accumulate size
+    for month in {1..12}; do
+        # Skip months after 7 if year is 2024
+        if [ $year -eq 2024 ] && [ $month -gt 7 ]; then
+            continue
+        fi
+
+        # Skip months before 8 if year is 2016
+        if [ $year -eq 2016 ] && [ $month -lt 8 ]; then
+            continue
+        fi
+
+        # Format month as 2 digits
+        month=$(printf "%02d" $month)
+
+        # Add the document path for this month to the array
+        documents+=("s3://ai2-llm/pretraining-data/sources/cc-news/v0-resiliparse/documents/${year}-${month}/*.zst")
+
+        # Get the size for this month and add it to the total size
+        month_size=$(aws s3api list-objects --bucket ai2-llm --prefix "pretraining-data/sources/cc-news/v0-resiliparse/documents/${year}-${month}/" --output json --query "[sum(Contents[].Size)]" | jq '.[0]' -rc)
+        size=$((size + month_size))
+    done
+
+
+    # run deduplication
+    echo "Running fuzzy dedupe for ${year} with ${size} bytes Bloom filter"
+
+    # Start the output
+    document_linearized="documents:\n"
+
+    # Loop through the array and append each element
+    for doc in "${documents[@]}"; do
+        document_linearized+="  - $doc\n"
+    done
+
+    config_yaml=$(cat <<EOF
+${document_linearized}
+dedupe:
+  name: dedupe_by_year
+  paragraphs:
+    attribute_name: dedupe_ngrams_13_1
+    by_ngram:
+      ngram_length: 13
+      stride: 1
+      overlap_threshold: 0.5
+      skip_short_paragraphs: true
+  skip_empty: true
+
+bloom_filter:
+  file: /tmp/cc_news_${year}_dedupe_ngram.bin
+  read_only: false
+  estimated_doc_count: ${size}
+  desired_false_positive_rate: 0.1
+
+work_dir:
+  input: /tmp/cc_news_${year}_dedupe_para_ngrams_13_1/input
+  output: /tmp/cc_news_${year}_dedupe_para_ngrams_13_1/output
+EOF
+)
+
+
+    # Create a temporary file for the YAML config
+    temp_config_file=$(mktemp)
+
+    # Write the YAML config to the temporary file
+    printf "$config_yaml" > "$temp_config_file"
+
+
+    set -ex
+    # Run dolma with the temporary config file
+    dolma -c "$temp_config_file" dedupe --processes $(expr $(nproc) - 4)
+    set +ex
+
+    # Remove the temporary file
+    rm "$temp_config_file"
+
+    done
+done

configs/cc-news/extract.sh

@@ -0,0 +1,43 @@
+#! /usr/bin/env bash
+
+# run years between 2016 and 2024
+for year in {2016..2024}; do
+    # run months between 1 and 12
+    for month in {1..12}; do
+        # skip months after 7 if year is 2024
+        if [ $year -eq 2024 ] && [ $month -gt 7 ]; then
+            continue
+        fi
+
+        # skip months before 8 if year is 2016
+        if [ $year -eq 2016 ] && [ $month -lt 8 ]; then
+            continue
+        fi
+
+        # rename month to 2 digits
+        month=$(printf "%02d" $month)
+
+        documents="s3://ai2-russella/crawl-data/CC-NEWS/${year}/${month}/*.warc.gz"
+
+        # run the extraction
+        echo "Running extraction for ${year}-${month}"
+
+        set -ex
+
+        dolma warc \
+            --documents ${documents} \
+            --destination s3://ai2-llm/pretraining-data/sources/cc-news/v0-resiliparse/documents/${year}-${month} \
+            --processes "$(expr $(nproc) - 4)" \
+            --source_name cc-news_${year}-${month} \
+            --linearizer resiliparse \
+            --pre.taggers cc_re \
+            --no-pre.skip \
+            --no-store.html \
+            --store.attr_spans 500 \
+            --skip_duplicate_urls \
+            --work_dir.input /tmp/cc-news/${year}-${month}/input \
+            --work_dir.output /tmp/cc-news/${year}-${month}/output
+
+        set +ex
+    done
+done

configs/cc-news/make_lang_partition.py

@@ -0,0 +1,59 @@
+import json
+from typing import List
+import smart_open
+
+
+SRC_BASE = "s3://ai2-llm/pretraining-data/sources/cc-news"
+SRC_PRFX = "v1-resiliparse"
+LANG_THR = 100_000
+DST_BASE = "s3://ai2-llm/pretraining-data/sources/cc-news"
+DST_PRFX = f"v2-resiliparse-l{LANG_THR // 1000}k"
+
+
+def base_stream_config(lang: str, year: int, months: List[int]):
+    return {
+        "name": f"cc-news_{year:04d}_{lang}",
+        "documents": [f"{SRC_BASE}/{SRC_PRFX}/documents/{year:04d}-{month:02d}/*.zst" for month in months],
+        "compression": {"input": "zst", "output": "zst"},
+        "output": {
+            "path": f"{DST_BASE}/{DST_PRFX}/documents/{lang}/{year:04d}",
+            "max_size_in_bytes": 10_000_000_000,
+        },
+        "attributes": ["ft_lang_id_1e2", "dolma_v2_tokenizer"],
+        "filter": {
+            "include": [],
+            "exclude": [
+                # at least 100 tokens
+                ".attributes.dolma_v2_tokenizer__dolma_v2_tokenizer__length[0][-1] <= 100",
+                # no language detected or low confidence
+                f"(.attributes.ft_lang_id_1e2__ft_lang_id_1e2__{lang} == null) or (.attributes.ft_lang_id_1e2__ft_lang_id_1e2__{lang}[0][-1] < 0.5)",
+
+            ],
+            "syntax": "jq",
+        },
+    }
+
+
+def main():
+    with smart_open.open("s3://ai2-llm/stats/cc-news/v1-resiliparse/attributes/ft_lang_id_1e2_summary.json") as f:
+        lang_counts = json.load(f)
+
+    languages = {k: v for k, v in lang_counts.items() if v >= LANG_THR}
+
+    streams = []
+    for year in range(2016, 2025):
+        if year == 2016:
+            months = list(range(8, 13))
+        elif year == 2024:
+            months = list(range(1, 8))
+        else:
+            months = list(range(1, 13))
+
+        streams.extend([base_stream_config(lang, year, months) for lang in languages])
+
+    with smart_open.open("configs/cc-news/mix_v2.json", "wt") as f:
+        json.dump({"processes": 1, "streams": streams}, f, indent=2)
+
+
+if __name__ == "__main__":
+    main()