Repository: allenai/dolma (language: python)

Issue:
ff



Discussion / hints:

Fix patch:
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index d2e67904..e139963b 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -20,6 +20,7 @@ permissions:
 env:
   DOLMA_TESTS_SKIP_AWS: ${{ secrets.AWS_ACCESS_KEY_ID == '' && 'true' || 'false' }}
   DOLMA_TEST_S3_PREFIX: s3://dolma-tests
+  RUST_CHANNEL: stable
 
 
 jobs:
@@ -38,9 +39,62 @@ jobs:
           echo "PR base repo:  ${{ github.event.pull_request.base.repo.full_name }}/tree/${{ github.event.pull_request.base.ref }}"
           echo "PR head repo:  ${{ github.event.pull_request.head.repo.full_name }}/tree/${{ github.event.pull_request.head.ref }}"
 
+  prepare-venv:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v3
+
+      - name: Cache Virtual Env
+        uses: actions/cache@v3
+        # name for referring later
+        id: cache-venv
+        with:
+          # what we cache: the virtualenv
+          path: ./.venv/
+          # The cache key depends on pyproject.toml and Cargo.toml
+          key: ${{ runner.os }}-venv-${{ hashFiles('**/pyproject.toml', '**/Cargo.toml, **/Cargo.lock') }}--${{ hashFiles('python/**', 'src/**') }}
+
+      - name: Setup system libraries
+        if: steps.cache-venv.outputs.cache-hit != 'true'
+        run:  |
+          sudo apt-get update
+          sudo apt-get install --yes --upgrade build-essential cmake protobuf-compiler libssl-dev glibc-source
+
+      - name: Install Rust toolchain
+        if: steps.cache-venv.outputs.cache-hit != 'true'
+        run: |
+          rustup update ${{ env.RUST_CHANNEL }}
+          rustup component add --toolchain ${{ env.RUST_CHANNEL }} rustfmt rust-src
+          rustup default ${{ env.RUST_CHANNEL }}
+
+      - name: Install Python
+        if: steps.cache-venv.outputs.cache-hit != 'true'
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.8'
+          architecture: "x64"
+          cache: 'pip'
+
+      - name: Create a new Python environment & install maturin
+        if: steps.cache-venv.outputs.cache-hit != 'true'
+        run: |
+          python -m venv .venv
+          source .venv/bin/activate
+          pip install -U pip
+          pip install maturin
+
+      - name: Install dolma wheels
+        if: steps.cache-venv.outputs.cache-hit != 'true'
+        run: |
+          source .venv/bin/activate
+          maturin build --release  -i $(which python) --out dist
+          wheel_path=$(ls dist/*.whl)
+          pip install "${wheel_path}[all]"
 
   tests:
     runs-on: ubuntu-latest
+    needs: prepare-venv
     env:
       AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
       AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
@@ -48,7 +102,6 @@ jobs:
     strategy:
       fail-fast: true
       matrix:
-        python: [3.8]
         task:
           - name: Check Python style
             run: |
@@ -73,50 +126,23 @@ jobs:
 
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v1
+        uses: actions/checkout@v3
 
-      - name: Setup system libraries
-        run:  |
-          sudo apt-get update
-          sudo apt-get install --yes --upgrade build-essential cmake protobuf-compiler libssl-dev glibc-source
-
-      - name: Install Rust
-        uses: actions-rs/toolchain@v1
+      - name: Cache Virtual Env
+        uses: actions/cache@v3
+        # name for referring later
+        id: cache-venv
         with:
-          toolchain: stable
-          components: rustfmt
-
-      - name: Install Python
-        uses: actions/setup-python@v2
-        with:
-          python-version: ${{ matrix.python }}
-          architecture: "x64"
-          sccache: true
-
-      - name: Create a new Python environment & install maturin
-        run: |
-          python -m venv .venv
-          source .venv/bin/activate
-          pip install -U pip
-          pip install maturin
-
-      - name: Install dolma wheels
-        run: |
-          source .venv/bin/activate
-          maturin develop --extras=dev
+          # what we cache: the virtualenv
+          path: ./.venv/
+          # The cache key depends on pyproject.toml and Cargo.toml
+          key: ${{ runner.os }}-venv-${{ hashFiles('**/pyproject.toml', '**/Cargo.toml, **/Cargo.lock') }}--${{ hashFiles('python/**', 'src/**') }}
 
       - name: ${{ matrix.task.name }}
         run: |
           source .venv/bin/activate
           ${{ matrix.task.run }}
 
-      - name: Clean up
-        if: always()
-        run: |
-          source .venv/bin/activate
-          pip uninstall -y dolma
-
-
 
   build-linux:
     if: ${{ github.ref == 'refs/heads/main' || github.ref == 'refs/heads/master' || startsWith(github.ref, 'refs/tags/') }}
@@ -132,6 +158,7 @@ jobs:
       - uses: actions/setup-python@v4
         with:
           python-version: '3.10'
+          cache: 'pip'
       - name: Setup environment
         run:  |
           sudo apt-get update
@@ -165,6 +192,7 @@ jobs:
         with:
           python-version: '3.10'
           architecture: ${{ matrix.target }}
+          cache: 'pip'
       - name: Build wheels
         uses: PyO3/maturin-action@v1
         with:
@@ -188,6 +216,7 @@ jobs:
       - uses: actions/setup-python@v4
         with:
           python-version: '3.10'
+          cache: 'pip'
       - name: Build wheels
         uses: PyO3/maturin-action@v1
         with:
diff --git a/.gitignore b/.gitignore
index 5edee38b..c85d3007 100644
--- a/.gitignore
+++ b/.gitignore
@@ -67,3 +67,7 @@ target/
 
 # ignore vscode directory
 .vscode
+
+# ignore temporary directories
+/tmp/
+/temp/
diff --git a/Makefile b/Makefile
index d80c4af4..d756f386 100644
--- a/Makefile
+++ b/Makefile
@@ -39,7 +39,7 @@ test-rust:
 	rm -rf tests/work/*
 
 develop:
-	maturin develop --extras=dev
+	maturin develop --extras=all
 
 style:
 	rustfmt --edition 2021 src/*.rs
diff --git a/pyproject.toml b/pyproject.toml
index 30248c87..9cfe5f97 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -8,22 +8,20 @@ requires-python = ">=3.8"
 dependencies = [
     "anyascii>=0.3.2",
     "blingfire==0.1.8",
-    "boto3",
+    "boto3>=1.28",
     "cached-path==1.3.4",
-    "detect-secrets==1.4.0",
     # "fasttext==0.9.2",    # broken with new version of setuptools; using fasttext-wheel instead
     "fasttext-wheel==0.9.2",
-    "fsspec",
+    "fsspec>=2023.6.0",
     "msgspec>=0.14.2",
     "nltk==3.8.1",
     "omegaconf>=2.3.0",
-    "presidio_analyzer==2.2.32",
     "pycld2==0.41",
     # "pycld3==0.22",       # does not install correctly
     "pyyaml",
     "requests",
     "rich",
-    "s3fs",
+    "s3fs>=2023.6.0",
     "smart-open",
     "tokenizers>=0.13.3,<1.0.0",
     "tqdm",
@@ -108,18 +106,39 @@ dev = [
     "flake8-pyi>=22.8.1",
     "Flake8-pyproject>=1.1.0",
 ]
-warc = [
-    "warcio>=1.7.4",
-    "trafilatura>=1.6.1",
-    "justext>=3.0.0",
-    "goose3>=3.1.17",
-
-    # following are all for speeding up trafilatura
-    "brotli",
-    "cchardet >= 2.1.7; python_version < '3.11'",  # build issue
-    "faust-cchardet >= 2.1.18; python_version >= '3.11'",  # fix for build
-    "htmldate[speed] >= 1.4.3",
-    "py3langid >= 0.2.2",
+# extension to process code
+code = [
+    "detect-secrets==1.4.0",
+    "beautifulsoup4>=4",
+    "pygments",
+    "regex"
+]
+# extension to detect PIIs using presidio
+pii = [
+    "presidio_analyzer==2.2.32",
+    "regex"
+]
+# # extension to parse warc files
+# warc = [
+#     "warcio>=1.7.4",
+#     "trafilatura>=1.6.1",
+#     "justext>=3.0.0",
+#     "goose3>=3.1.17",
+
+#     # following are all for speeding up trafilatura
+#     "brotli",
+#     "cchardet >= 2.1.7; python_version < '3.11'",  # build issue
+#     "faust-cchardet >= 2.1.18; python_version >= '3.11'",  # fix for build
+#     "htmldate[speed] >= 1.4.3",
+#     "py3langid >= 0.2.2",
+# ]
+
+# all extensions
+all = [
+    "dolma[dev]",
+    "dolma[code]",
+    "dolma[pii]",
+    # "dolma[warc]",
 ]
 
 [build-system]
diff --git a/python/dolma/core/data_types.py b/python/dolma/core/data_types.py
index 1ce2a0c5..9bd1c3ee 100644
--- a/python/dolma/core/data_types.py
+++ b/python/dolma/core/data_types.py
@@ -21,8 +21,10 @@ class InputSpec(Struct):
     text: str
     source: str = ""
     version: Optional[str] = None
-    # ignoring metadata for now; taggers run on text only
-    # metadata: Optional[Dict[str, Any]] = None
+
+
+class InputSpecWithMetadata(InputSpec):
+    metadata: Optional[Dict[str, Any]] = None
 
 
 class OutputSpec(Struct):
@@ -48,17 +50,67 @@ def to_spec(self) -> InputSpec:
         return InputSpec(source=self.source, version=self.version, id=self.id, text=self.text)
 
     @classmethod
-    def from_json(cls, d: Dict) -> "Document":
+    def from_json(cls, d: Dict[str, Any]) -> "Document":
         return Document(source=d["source"], version=d["version"], id=d["id"], text=d["text"])
 
-    def to_json(self) -> Dict:
+    def to_json(self) -> Dict[str, Any]:
         return {"source": self.source, "version": self.version, "id": self.id, "text": self.text}
 
     def __str__(self) -> str:
-        attributes_string = ",".join([f"{k}:{repr(v)}" for k, v in self.to_json()])
+        attributes_string = ",".join([f"{k}:{repr(v)}" for k, v in self.to_json().items()])
         return f"{self.__class__.__name__}({attributes_string})"
 
 
+class DocumentWithMetadata(Document):
+    __slots__ = ("metadata",)
+
+    def __init__(self, *args, metadata: Optional[Dict[str, Any]] = None, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.metadata = metadata or {}
+
+    @classmethod
+    def from_spec(cls, spec: InputSpecWithMetadata) -> "DocumentWithMetadata":
+        return DocumentWithMetadata(
+            source=spec.source,
+            version=spec.version,
+            id=spec.id,
+            text=spec.text,
+            metadata=spec.metadata,
+        )
+
+    def to_spec(self) -> InputSpecWithMetadata:
+        return InputSpecWithMetadata(
+            source=self.source,
+            version=self.version,
+            id=self.id,
+            text=self.text,
+            metadata=self.metadata,
+        )
+
+    @classmethod
+    def from_json(cls, d: Dict) -> "DocumentWithMetadata":
+        return DocumentWithMetadata(
+            source=d["source"],
+            version=d["version"],
+            id=d["id"],
+            text=d["text"],
+            metadata=d["metadata"],
+        )
+
+    def to_json(self) -> Dict:
+        return {
+            "source": self.source,
+            "version": self.version,
+            "id": self.id,
+            "text": self.text,
+            "metadata": self.metadata,
+        }
+
+    def __str__(self) -> str:
+        repr_ = super().__str__()
+        return repr_.rstrip(")") + f",metadata={'...' if self.metadata else 'none'})"
+
+
 class Span:
     __slots__ = "start", "end", "type", "score", "experiment", "tagger"
 
@@ -127,6 +179,19 @@ def __str__(self) -> str:
         cls_name = self.__class__.__name__
         return f"{cls_name}(start={self.start},end={self.end},type={repr(self.type)},score={self.score:.5f})"
 
+    def __repr__(self) -> str:
+        return str(self)
+
+    def __eq__(self, other: Any) -> bool:
+        if not isinstance(other, self.__class__):
+            return False
+        return (
+            self.start == other.start
+            and self.end == other.end
+            and self.type == other.type
+            and self.score == other.score
+        )
+
 
 class DocResult:
     __slots__ = "doc", "spans"
diff --git a/python/dolma/core/loggers.py b/python/dolma/core/loggers.py
index 7974d66a..c84bf66f 100644
--- a/python/dolma/core/loggers.py
+++ b/python/dolma/core/loggers.py
@@ -1,8 +1,13 @@
 import logging
+import multiprocessing
 
 
 def get_logger(name: str) -> logging.Logger:
-    name = f"dolma.{name}"
+    if (proc_name := multiprocessing.current_process().name) == "MainProcess":
+        proc_name = "main"
+    proc_name = proc_name.replace(" ", "_")
+
+    name = f"{proc_name}.dolma.{name}"
     logger = logging.getLogger(name)
     logger.setLevel(logging.WARN)
 
diff --git a/python/dolma/core/runtime.py b/python/dolma/core/runtime.py
index 62ea0f5e..a181e3e2 100644
--- a/python/dolma/core/runtime.py
+++ b/python/dolma/core/runtime.py
@@ -18,7 +18,14 @@
 import msgspec
 import smart_open
 
-from .data_types import InputSpec, OutputSpec, TaggerOutputDictType
+from dolma.core.taggers import BaseTaggerWithMetadata
+
+from .data_types import (
+    InputSpec,
+    InputSpecWithMetadata,
+    OutputSpec,
+    TaggerOutputDictType,
+)
 from .errors import DolmaFatalError, DolmaRetryableFailure, DolmaShardError
 from .parallel import BaseParallelProcessor, QueueType
 from .paths import delete_dir, join_path, make_relative, mkdir_p, split_glob, split_path
@@ -252,12 +259,19 @@ def process_single(
         # too full
         update_interval = 1
 
-        # running document count; gets reset every time we update the progress
-        # bar
+        # running document count; gets reset every time we update the progress bar
         docs_cnt = 0
 
+        # total number of documents processed
+        total_docs_cnt = 0
+
         # creating dedicated decoder speeds up the process
-        decoder = msgspec.json.Decoder(InputSpec)
+        # if any of the taggers require metadata, we use a decoder that can handle it
+        # otherwise, we use a decoder that does not parse metadata, which is faster
+        if any(isinstance(tagger, BaseTaggerWithMetadata) for tagger in taggers.values()):
+            decoder = msgspec.json.Decoder(InputSpecWithMetadata)
+        else:
+            decoder = msgspec.json.Decoder(InputSpec)
 
         with ExitStack() as stack:
             in_stream = stack.enter_context(smart_open.open(source_path, "rt", encoding="utf-8"))
@@ -279,8 +293,9 @@ def process_single(
 
                     # increment the number of documents processed so far
                     docs_cnt += 1
+                    total_docs_cnt += 1
 
-                    if steps is not None and docs_cnt >= steps:
+                    if steps is not None and total_docs_cnt >= steps:
                         # if we have reached the maximum number of steps, we break
                         break
 
diff --git a/python/dolma/core/taggers.py b/python/dolma/core/taggers.py
index 8d51d9aa..342985fd 100644
--- a/python/dolma/core/taggers.py
+++ b/python/dolma/core/taggers.py
@@ -8,7 +8,14 @@
 from abc import abstractmethod
 from typing import List
 
-from .data_types import DocResult, Document, InputSpec, TaggerOutputDictType
+from .data_types import (
+    DocResult,
+    Document,
+    DocumentWithMetadata,
+    InputSpec,
+    InputSpecWithMetadata,
+    TaggerOutputDictType,
+)
 
 # digits after the decimal point
 TAGGER_SCORE_PRECISION = 5
@@ -35,14 +42,27 @@ def defaults(self) -> List[str]:
     def predict(self, doc: Document) -> DocResult:
         raise NotImplementedError
 
-    def tag(self, row: InputSpec) -> TaggerOutputDictType:
-        """Internal function that is used by the tagger to get data"""
-        doc = Document(source=row.source, version=row.version, id=row.id, text=row.text)
-        doc_result = self.predict(doc)
-
+    def group_output(self, doc_result: DocResult) -> TaggerOutputDictType:
         tagger_output: TaggerOutputDictType = {field: [] for field in self.defaults}
         for span in doc_result.spans:
             output = (span.start, span.end, round(float(span.score), TAGGER_SCORE_PRECISION))
             tagger_output.setdefault(span.type, []).append(output)
-
         return tagger_output
+
+    def tag(self, row: InputSpec) -> TaggerOutputDictType:
+        """Internal function that is used by the tagger to get data"""
+        doc = Document.from_spec(row)
+        doc_result = self.predict(doc)
+        return self.group_output(doc_result)
+
+
+class BaseTaggerWithMetadata(BaseTagger):
+    @abstractmethod
+    def predict(self, doc: DocumentWithMetadata) -> DocResult:  # type: ignore
+        raise NotImplementedError
+
+    def tag(self, row: InputSpecWithMetadata) -> TaggerOutputDictType:
+        """Internal function that is used by the tagger to get data"""
+        doc = DocumentWithMetadata.from_spec(row)
+        doc_result = self.predict(doc)
+        return self.group_output(doc_result)
diff --git a/python/dolma/core/utils.py b/python/dolma/core/utils.py
index 6ec012a9..a11a3aa5 100644
--- a/python/dolma/core/utils.py
+++ b/python/dolma/core/utils.py
@@ -13,6 +13,7 @@
     BLINGFIRE_AVAILABLE = False
 
 import nltk
+import uniseg.wordbreak
 from nltk.tokenize.punkt import PunktSentenceTokenizer
 
 try:
@@ -42,6 +43,20 @@ def make_variable_name(name: str, remove_multiple_underscores: bool = False) ->
     return name
 
 
+def split_words(text: str, remove_empty: bool = True) -> List[TextSlice]:
+    """
+    Split a string into words, as defined by the unicode standard.
+    For more info, see https://www.unicode.org/reports/tr29/
+    """
+    text_slices: List[TextSlice] = []
+    offset = 0
+    for word in uniseg.wordbreak.words(text):
+        if word.strip() or not remove_empty:
+            text_slices.append(TextSlice(doc=text, start=offset, end=offset + len(word)))
+        offset += len(word)
+    return text_slices
+
+
 def split_paragraphs(text: str, remove_empty: bool = True) -> List[TextSlice]:
     """
     Split a string into paragraphs. A paragraph is defined as a sequence of zero or more characters, followed
diff --git a/python/dolma/data/ext_to_lang_mapping.json b/python/dolma/data/ext_to_lang_mapping.json
new file mode 100644
index 00000000..72240bcb
--- /dev/null
+++ b/python/dolma/data/ext_to_lang_mapping.json
@@ -0,0 +1,855 @@
+{
+    "abap": "abap",
+    "ash": "ags-script",
+    "ampl": "ampl",
+    "g4": "antlr",
+    "apib": "api-blueprint",
+    "apl": "apl",
+    "dyalog": "apl",
+    "asp": "asp",
+    "asax": "asp",
+    "ascx": "asp",
+    "ashx": "asp",
+    "asmx": "asp",
+    "aspx": "asp",
+    "axd": "asp",
+    "dats": "ats",
+    "hats": "ats",
+    "sats": "ats",
+    "as": "actionscript",
+    "adb": "ada",
+    "ada": "ada",
+    "ads": "ada",
+    "agda": "agda",
+    "als": "alloy",
+    "apacheconf": "apacheconf",
+    "vhost": "apacheconf",
+    "applescript": "applescript",
+    "scpt": "applescript",
+    "arc": "arc",
+    "ino": "arduino",
+    "asciidoc": "asciidoc",
+    "adoc": "asciidoc",
+    "aj": "aspectj",
+    "asm": "assembly",
+    "a51": "assembly",
+    "nasm": "assembly",
+    "aug": "augeas",
+    "ahk": "autohotkey",
+    "ahkl": "autohotkey",
+    "au3": "autoit",
+    "awk": "awk",
+    "auk": "awk",
+    "gawk": "awk",
+    "mawk": "awk",
+    "nawk": "awk",
+    "bat": "batchfile",
+    "cmd": "batchfile",
+    "befunge": "befunge",
+    "bison": "bison",
+    "bb": "bitbake",
+    "decls": "blitzbasic",
+    "bmx": "blitzmax",
+    "bsv": "bluespec",
+    "boo": "boo",
+    "bf": "brainfuck",
+    "brs": "brightscript",
+    "bro": "bro",
+    "c": "c",
+    "cats": "c",
+    "h": "c",
+    "idc": "c",
+    "w": "c",
+    "cs": "c-sharp",
+    "cake": "c-sharp",
+    "cshtml": "c-sharp",
+    "csx": "c-sharp",
+    "cpp": "c++",
+    "c++": "c++",
+    "cc": "c++",
+    "cp": "c++",
+    "cxx": "c++",
+    "h++": "c++",
+    "hh": "c++",
+    "hpp": "c++",
+    "hxx": "c++",
+    "inl": "c++",
+    "ipp": "c++",
+    "tcc": "c++",
+    "tpp": "c++",
+    "C": "c++",
+    "H": "c++",
+    "chs": "c2hs-haskell",
+    "clp": "clips",
+    "cmake": "cmake",
+    "cmakein": "cmake",
+    "cob": "cobol",
+    "cbl": "cobol",
+    "ccp": "cobol",
+    "cobol": "cobol",
+    "cpy": "cobol",
+    "css": "css",
+    "csv": "csv",
+    "capnp": "cap'n-proto",
+    "mss": "cartocss",
+    "ceylon": "ceylon",
+    "chpl": "chapel",
+    "ck": "chuck",
+    "cirru": "cirru",
+    "clw": "clarion",
+    "icl": "clean",
+    "dcl": "clean",
+    "click": "click",
+    "clj": "clojure",
+    "boot": "clojure",
+    "cl2": "clojure",
+    "cljc": "clojure",
+    "cljs": "clojure",
+    "cljshl": "clojure",
+    "cljscm": "clojure",
+    "cljx": "clojure",
+    "hic": "clojure",
+    "coffee": "coffeescript",
+    "_coffee": "coffeescript",
+    "cjsx": "coffeescript",
+    "cson": "coffeescript",
+    "iced": "coffeescript",
+    "cfm": "coldfusion",
+    "cfml": "coldfusion",
+    "cfc": "coldfusion-cfc",
+    "lisp": "common-lisp",
+    "asd": "common-lisp",
+    "lsp": "common-lisp",
+    "ny": "common-lisp",
+    "podsl": "common-lisp",
+    "sexp": "common-lisp",
+    "cps": "component-pascal",
+    "coq": "coq",
+    "creole": "creole",
+    "cr": "crystal",
+    "csd": "csound",
+    "feature": "cucumber",
+    "cu": "cuda",
+    "cuh": "cuda",
+    "cy": "cycript",
+    "pyx": "cython",
+    "pxd": "cython",
+    "pxi": "cython",
+    "di": "d",
+    "com": "digital-command-language",
+    "dm": "dm",
+    "zone": "dns-zone",
+    "arpa": "dns-zone",
+    "darcspatch": "darcs-patch",
+    "dpatch": "darcs-patch",
+    "dart": "dart",
+    "diff": "diff",
+    "patch": "diff",
+    "dockerfile": "dockerfile",
+    "Dockerfile": "dockerfile",
+    "djs": "dogescript",
+    "dylan": "dylan",
+    "dyl": "dylan",
+    "intr": "dylan",
+    "lid": "dylan",
+    "ecl": "ecl",
+    "eclxml": "ecl",
+    "sch": "eagle",
+    "brd": "eagle",
+    "epj": "ecere-projects",
+    "e": "eiffel",
+    "ex": "elixir",
+    "exs": "elixir",
+    "elm": "elm",
+    "el": "emacs-lisp",
+    "emacs": "emacs-lisp",
+    "emacsdesktop": "emacs-lisp",
+    "em": "emberscript",
+    "emberscript": "emberscript",
+    "erl": "erlang",
+    "escript": "erlang",
+    "hrl": "erlang",
+    "xrl": "erlang",
+    "yrl": "erlang",
+    "fs": "f-sharp",
+    "fsi": "f-sharp",
+    "fsx": "f-sharp",
+    "flux": "flux",
+    "f90": "fortran",
+    "f": "fortran",
+    "f03": "fortran",
+    "f08": "fortran",
+    "f77": "fortran",
+    "f95": "fortran",
+    "for": "fortran",
+    "fpp": "fortran",
+    "factor": "factor",
+    "fy": "fancy",
+    "fancypack": "fancy",
+    "fan": "fantom",
+    "fth": "forth",
+    "4th": "forth",
+    "forth": "forth",
+    "frt": "forth",
+    "ftl": "freemarker",
+    "g": "g-code",
+    "gco": "g-code",
+    "gcode": "g-code",
+    "gms": "gams",
+    "gap": "gap",
+    "gi": "gap",
+    "s": "gas",
+    "gd": "gdscript",
+    "glsl": "glsl",
+    "fp": "glsl",
+    "frag": "glsl",
+    "frg": "glsl",
+    "fsh": "glsl",
+    "fshader": "glsl",
+    "geo": "glsl",
+    "geom": "glsl",
+    "glslv": "glsl",
+    "gshader": "glsl",
+    "shader": "glsl",
+    "vert": "glsl",
+    "vrx": "glsl",
+    "vsh": "glsl",
+    "vshader": "glsl",
+    "kid": "genshi",
+    "ebuild": "gentoo-ebuild",
+    "eclass": "gentoo-eclass",
+    "po": "gettext-catalog",
+    "pot": "gettext-catalog",
+    "glf": "glyph",
+    "gp": "gnuplot",
+    "gnu": "gnuplot",
+    "gnuplot": "gnuplot",
+    "plot": "gnuplot",
+    "plt": "gnuplot",
+    "go": "go",
+    "golo": "golo",
+    "gst": "gosu",
+    "gsx": "gosu",
+    "vark": "gosu",
+    "grace": "grace",
+    "gf": "grammatical-framework",
+    "graphql": "graphql",
+    "dot": "graphviz-(dot)",
+    "gv": "graphviz-(dot)",
+    "man": "groff",
+    "1": "groff",
+    "1in": "groff",
+    "1m": "groff",
+    "1x": "groff",
+    "2": "groff",
+    "3": "groff",
+    "3in": "groff",
+    "3m": "groff",
+    "3qt": "groff",
+    "3x": "groff",
+    "4": "groff",
+    "5": "groff",
+    "6": "groff",
+    "7": "groff",
+    "8": "groff",
+    "9": "groff",
+    "me": "groff",
+    "rno": "groff",
+    "roff": "groff",
+    "groovy": "groovy",
+    "grt": "groovy",
+    "gtpl": "groovy",
+    "gvy": "groovy",
+    "gsp": "groovy-server-pages",
+    "hcl": "hcl",
+    "tf": "hcl",
+    "hlsl": "hlsl",
+    "fxh": "hlsl",
+    "hlsli": "hlsl",
+    "html": "html",
+    "htm": "html",
+    "htmlhl": "html",
+    "xht": "html",
+    "xhtml": "html",
+    "mustache": "html+django",
+    "jinja": "html+django",
+    "eex": "html+eex",
+    "erb": "html+erb",
+    "erbdeface": "html+erb",
+    "phtml": "html+php",
+    "http": "http",
+    "haml": "haml",
+    "hamldeface": "haml",
+    "handlebars": "handlebars",
+    "hbs": "handlebars",
+    "hb": "harbour",
+    "hs": "haskell",
+    "hsc": "haskell",
+    "hx": "haxe",
+    "hxsl": "haxe",
+    "hy": "hy",
+    "dlm": "idl",
+    "ipf": "igor-pro",
+    "ini": "ini",
+    "cfg": "ini",
+    "prefs": "ini",
+    "properties": "ini",
+    "irclog": "irc-log",
+    "weechatlog": "irc-log",
+    "idr": "idris",
+    "lidr": "idris",
+    "ni": "inform-7",
+    "i7x": "inform-7",
+    "iss": "inno-setup",
+    "io": "io",
+    "ik": "ioke",
+    "thy": "isabelle",
+    "ijs": "j",
+    "flex": "jflex",
+    "jflex": "jflex",
+    "json": "json",
+    "geojson": "json",
+    "lock": "json",
+    "topojson": "json",
+    "json5": "json5",
+    "jsonld": "jsonld",
+    "jq": "jsoniq",
+    "jsx": "jsx",
+    "jade": "jade",
+    "j": "jasmin",
+    "java": "java",
+    "jsp": "java-server-pages",
+    "js": "javascript",
+    "_js": "javascript",
+    "bones": "javascript",
+    "es6": "javascript",
+    "jake": "javascript",
+    "jsb": "javascript",
+    "jscad": "javascript",
+    "jsfl": "javascript",
+    "jsm": "javascript",
+    "jss": "javascript",
+    "njs": "javascript",
+    "pac": "javascript",
+    "sjs": "javascript",
+    "ssjs": "javascript",
+    "xsjs": "javascript",
+    "xsjslib": "javascript",
+    "jl": "julia",
+    "ipynb": "jupyter-notebook",
+    "krl": "krl",
+    "kicad_pcb": "kicad",
+    "kit": "kit",
+    "kt": "kotlin",
+    "ktm": "kotlin",
+    "kts": "kotlin",
+    "lfe": "lfe",
+    "ll": "llvm",
+    "lol": "lolcode",
+    "lsl": "lsl",
+    "lslp": "lsl",
+    "lvproj": "labview",
+    "lasso": "lasso",
+    "las": "lasso",
+    "lasso8": "lasso",
+    "lasso9": "lasso",
+    "ldml": "lasso",
+    "latte": "latte",
+    "lean": "lean",
+    "hlean": "lean",
+    "less": "less",
+    "lex": "lex",
+    "ly": "lilypond",
+    "ily": "lilypond",
+    "ld": "linker-script",
+    "lds": "linker-script",
+    "liquid": "liquid",
+    "lagda": "literate-agda",
+    "litcoffee": "literate-coffeescript",
+    "lhs": "literate-haskell",
+    "ls": "livescript",
+    "_ls": "livescript",
+    "xm": "logos",
+    "x": "logos",
+    "xi": "logos",
+    "lgt": "logtalk",
+    "logtalk": "logtalk",
+    "lookml": "lookml",
+    "lua": "lua",
+    "nse": "lua",
+    "pd_lua": "lua",
+    "rbxs": "lua",
+    "wlua": "lua",
+    "mumps": "m",
+    "m4": "m4",
+    "mcr": "maxscript",
+    "mtml": "mtml",
+    "muf": "muf",
+    "mak": "makefile",
+    "mk": "makefile",
+    "mkfile": "makefile",
+    "Makefile": "makefile",
+    "mako": "mako",
+    "mao": "mako",
+    "mpl": "maple",
+    "md": "markdown",
+    "markdown": "markdown",
+    "mkd": "markdown",
+    "mkdn": "markdown",
+    "mkdown": "markdown",
+    "ron": "markdown",
+    "mask": "mask",
+    "mathematica": "mathematica",
+    "cdf": "mathematica",
+    "ma": "mathematica",
+    "mt": "mathematica",
+    "nb": "mathematica",
+    "nbp": "mathematica",
+    "wl": "mathematica",
+    "wlt": "mathematica",
+    "matlab": "matlab",
+    "maxpat": "max",
+    "maxhelp": "max",
+    "maxproj": "max",
+    "mxt": "max",
+    "pat": "max",
+    "mediawiki": "mediawiki",
+    "wiki": "mediawiki",
+    "metal": "metal",
+    "druby": "mirah",
+    "duby": "mirah",
+    "mir": "mirah",
+    "mirah": "mirah",
+    "mo": "modelica",
+    "mms": "module-management-system",
+    "mmk": "module-management-system",
+    "monkey": "monkey",
+    "moon": "moonscript",
+    "myt": "myghty",
+    "nsi": "nsis",
+    "nsh": "nsis",
+    "axs": "netlinx",
+    "axi": "netlinx",
+    "nlogo": "netlogo",
+    "nginxconf": "nginx",
+    "nim": "nimrod",
+    "nimrod": "nimrod",
+    "ninja": "ninja",
+    "nit": "nit",
+    "nix": "nix",
+    "nu": "nu",
+    "numpy": "numpy",
+    "numpyw": "numpy",
+    "numsc": "numpy",
+    "ml": "ocaml",
+    "eliom": "ocaml",
+    "eliomi": "ocaml",
+    "ml4": "ocaml",
+    "mli": "ocaml",
+    "mll": "ocaml",
+    "mly": "ocaml",
+    "objdump": "objdump",
+    "mm": "objective-c++",
+    "sj": "objective-j",
+    "oct": "octave",
+    "omgrofl": "omgrofl",
+    "opa": "opa",
+    "opal": "opal",
+    "cl": "opencl",
+    "opencl": "opencl",
+    "scad": "openscad",
+    "org": "org",
+    "ox": "ox",
+    "oxh": "ox",
+    "oxo": "ox",
+    "oxygene": "oxygene",
+    "oz": "oz",
+    "pwn": "pawn",
+    "php": "php",
+    "aw": "php",
+    "ctp": "php",
+    "php3": "php",
+    "php4": "php",
+    "php5": "php",
+    "phps": "php",
+    "phpt": "php",
+    "pov": "pov-ray-sdl",
+    "pan": "pan",
+    "psc": "papyrus",
+    "parrot": "parrot",
+    "pasm": "parrot-assembly",
+    "pir": "parrot-internal-representation",
+    "pas": "pascal",
+    "dfm": "pascal",
+    "dpr": "pascal",
+    "lpr": "pascal",
+    "pl": "perl",
+    "al": "perl",
+    "perl": "perl",
+    "ph": "perl",
+    "plx": "perl",
+    "pm": "perl",
+    "psgi": "perl",
+    "t": "perl",
+    "6pl": "perl6",
+    "6pm": "perl6",
+    "nqp": "perl6",
+    "p6": "perl6",
+    "p6l": "perl6",
+    "p6m": "perl6",
+    "pl6": "perl6",
+    "pm6": "perl6",
+    "pig": "piglatin",
+    "pike": "pike",
+    "pmod": "pike",
+    "pod": "pod",
+    "pogo": "pogoscript",
+    "pony": "pony",
+    "ps": "postscript",
+    "eps": "postscript",
+    "ps1": "powershell",
+    "psd1": "powershell",
+    "psm1": "powershell",
+    "pde": "processing",
+    "prolog": "prolog",
+    "yap": "prolog",
+    "spin": "propeller-spin",
+    "proto": "protocol-buffer",
+    "pd": "pure-data",
+    "pb": "purebasic",
+    "pbi": "purebasic",
+    "purs": "purescript",
+    "py": "python",
+    "bzl": "python",
+    "gyp": "python",
+    "lmi": "python",
+    "pyde": "python",
+    "pyp": "python",
+    "pyt": "python",
+    "pyw": "python",
+    "tac": "python",
+    "wsgi": "python",
+    "xpy": "python",
+    "pytb": "python-traceback",
+    "qml": "qml",
+    "qbs": "qml",
+    "pri": "qmake",
+    "r": "r",
+    "rd": "r",
+    "rsx": "r",
+    "raml": "raml",
+    "rdoc": "rdoc",
+    "rbbas": "realbasic",
+    "rbfrm": "realbasic",
+    "rbmnu": "realbasic",
+    "rbres": "realbasic",
+    "rbtbar": "realbasic",
+    "rbuistate": "realbasic",
+    "rhtml": "rhtml",
+    "rmd": "rmarkdown",
+    "rkt": "racket",
+    "rktd": "racket",
+    "rktl": "racket",
+    "scrbl": "racket",
+    "rl": "ragel-in-ruby-host",
+    "reb": "rebol",
+    "r2": "rebol",
+    "r3": "rebol",
+    "rebol": "rebol",
+    "red": "red",
+    "reds": "red",
+    "cw": "redcode",
+    "rpy": "ren'py",
+    "rsh": "renderscript",
+    "robot": "robotframework",
+    "rg": "rouge",
+    "rb": "ruby",
+    "builder": "ruby",
+    "gemspec": "ruby",
+    "god": "ruby",
+    "irbrc": "ruby",
+    "jbuilder": "ruby",
+    "mspec": "ruby",
+    "podspec": "ruby",
+    "rabl": "ruby",
+    "rake": "ruby",
+    "rbuild": "ruby",
+    "rbw": "ruby",
+    "rbx": "ruby",
+    "ru": "ruby",
+    "ruby": "ruby",
+    "thor": "ruby",
+    "watchr": "ruby",
+    "rs": "rust",
+    "rsin": "rust",
+    "sas": "sas",
+    "scss": "scss",
+    "smt2": "smt",
+    "smt": "smt",
+    "sparql": "sparql",
+    "rq": "sparql",
+    "sqf": "sqf",
+    "hqf": "sqf",
+    "pls": "sql",
+    "pck": "sql",
+    "pkb": "sql",
+    "pks": "sql",
+    "plb": "sql",
+    "plsql": "sql",
+    "sql": "sql",
+    "cql": "sql",
+    "ddl": "sql",
+    "prc": "sql",
+    "tab": "sql",
+    "udf": "sql",
+    "viw": "sql",
+    "db2": "sql",
+    "ston": "ston",
+    "svg": "svg",
+    "sage": "sage",
+    "sagews": "sage",
+    "sls": "saltstack",
+    "sass": "sass",
+    "scala": "scala",
+    "sbt": "scala",
+    "scaml": "scaml",
+    "scm": "scheme",
+    "sld": "scheme",
+    "sps": "scheme",
+    "ss": "scheme",
+    "sci": "scilab",
+    "sce": "scilab",
+    "self": "self",
+    "sh": "shell",
+    "bash": "shell",
+    "bats": "shell",
+    "command": "shell",
+    "ksh": "shell",
+    "shin": "shell",
+    "tmux": "shell",
+    "tool": "shell",
+    "zsh": "shell",
+    "sh-session": "shellsession",
+    "shen": "shen",
+    "sl": "slash",
+    "slim": "slim",
+    "smali": "smali",
+    "st": "smalltalk",
+    "tpl": "smarty",
+    "sol": "solidity",
+    "sp": "sourcepawn",
+    "sma": "sourcepawn",
+    "nut": "squirrel",
+    "stan": "stan",
+    "ML": "standard-ml",
+    "fun": "standard-ml",
+    "sig": "standard-ml",
+    "sml": "standard-ml",
+    "do": "stata",
+    "ado": "stata",
+    "doh": "stata",
+    "ihlp": "stata",
+    "mata": "stata",
+    "matah": "stata",
+    "sthlp": "stata",
+    "styl": "stylus",
+    "scd": "supercollider",
+    "swift": "swift",
+    "sv": "systemverilog",
+    "svh": "systemverilog",
+    "vh": "systemverilog",
+    "toml": "toml",
+    "txl": "txl",
+    "tcl": "tcl",
+    "adp": "tcl",
+    "tm": "tcl",
+    "tcsh": "tcsh",
+    "csh": "tcsh",
+    "tex": "tex",
+    "aux": "tex",
+    "bbx": "tex",
+    "bib": "tex",
+    "cbx": "tex",
+    "dtx": "tex",
+    "ins": "tex",
+    "lbx": "tex",
+    "ltx": "tex",
+    "mkii": "tex",
+    "mkiv": "tex",
+    "mkvi": "tex",
+    "sty": "tex",
+    "toc": "tex",
+    "tea": "tea",
+    "txt": "text",
+    "no": "text",
+    "textile": "textile",
+    "thrift": "thrift",
+    "tu": "turing",
+    "ttl": "turtle",
+    "twig": "twig",
+    "ts": "typescript",
+    "tsx": "typescript",
+    "upc": "unified-parallel-c",
+    "anim": "unity3d-asset",
+    "asset": "unity3d-asset",
+    "mat": "unity3d-asset",
+    "meta": "unity3d-asset",
+    "prefab": "unity3d-asset",
+    "unity": "unity3d-asset",
+    "uno": "uno",
+    "uc": "unrealscript",
+    "ur": "urweb",
+    "urs": "urweb",
+    "vcl": "vcl",
+    "vhdl": "vhdl",
+    "vhd": "vhdl",
+    "vhf": "vhdl",
+    "vhi": "vhdl",
+    "vho": "vhdl",
+    "vhs": "vhdl",
+    "vht": "vhdl",
+    "vhw": "vhdl",
+    "vala": "vala",
+    "vapi": "vala",
+    "veo": "verilog",
+    "vim": "viml",
+    "vb": "visual-basic",
+    "bas": "visual-basic",
+    "frm": "visual-basic",
+    "frx": "visual-basic",
+    "vba": "visual-basic",
+    "vbhtml": "visual-basic",
+    "vbs": "visual-basic",
+    "volt": "volt",
+    "vue": "vue",
+    "owl": "web-ontology-language",
+    "wat": "webassembly",
+    "webidl": "webidl",
+    "x10": "x10",
+    "xc": "xc",
+    "xml": "xml",
+    "ant": "xml",
+    "axml": "xml",
+    "ccxml": "xml",
+    "clixml": "xml",
+    "cproject": "xml",
+    "csl": "xml",
+    "csproj": "xml",
+    "ct": "xml",
+    "dita": "xml",
+    "ditamap": "xml",
+    "ditaval": "xml",
+    "dllconfig": "xml",
+    "dotsettings": "xml",
+    "filters": "xml",
+    "fsproj": "xml",
+    "fxml": "xml",
+    "glade": "xml",
+    "grxml": "xml",
+    "iml": "xml",
+    "ivy": "xml",
+    "jelly": "xml",
+    "jsproj": "xml",
+    "kml": "xml",
+    "launch": "xml",
+    "mdpolicy": "xml",
+    "mxml": "xml",
+    "nproj": "xml",
+    "nuspec": "xml",
+    "odd": "xml",
+    "osm": "xml",
+    "plist": "xml",
+    "props": "xml",
+    "ps1xml": "xml",
+    "psc1": "xml",
+    "pt": "xml",
+    "rdf": "xml",
+    "rss": "xml",
+    "scxml": "xml",
+    "srdf": "xml",
+    "storyboard": "xml",
+    "stTheme": "xml",
+    "sublime-snippet": "xml",
+    "targets": "xml",
+    "tmCommand": "xml",
+    "tml": "xml",
+    "tmLanguage": "xml",
+    "tmPreferences": "xml",
+    "tmSnippet": "xml",
+    "tmTheme": "xml",
+    "ui": "xml",
+    "urdf": "xml",
+    "ux": "xml",
+    "vbproj": "xml",
+    "vcxproj": "xml",
+    "vssettings": "xml",
+    "vxml": "xml",
+    "wsdl": "xml",
+    "wsf": "xml",
+    "wxi": "xml",
+    "wxl": "xml",
+    "wxs": "xml",
+    "x3d": "xml",
+    "xacro": "xml",
+    "xaml": "xml",
+    "xib": "xml",
+    "xlf": "xml",
+    "xliff": "xml",
+    "xmi": "xml",
+    "xmldist": "xml",
+    "xproj": "xml",
+    "xsd": "xml",
+    "xul": "xml",
+    "zcml": "xml",
+    "xsp-config": "xpages",
+    "xspmetadata": "xpages",
+    "xpl": "xproc",
+    "xproc": "xproc",
+    "xquery": "xquery",
+    "xq": "xquery",
+    "xql": "xquery",
+    "xqm": "xquery",
+    "xqy": "xquery",
+    "xs": "xs",
+    "xslt": "xslt",
+    "xsl": "xslt",
+    "xojo_code": "xojo",
+    "xojo_menu": "xojo",
+    "xojo_report": "xojo",
+    "xojo_script": "xojo",
+    "xojo_toolbar": "xojo",
+    "xojo_window": "xojo",
+    "xtend": "xtend",
+    "yml": "yaml",
+    "reek": "yaml",
+    "rviz": "yaml",
+    "sublime-syntax": "yaml",
+    "syntax": "yaml",
+    "yaml": "yaml",
+    "yaml-tmlanguage": "yaml",
+    "yang": "yang",
+    "y": "yacc",
+    "yacc": "yacc",
+    "yy": "yacc",
+    "zep": "zephir",
+    "zig": "zig",
+    "zimpl": "zimpl",
+    "zmpl": "zimpl",
+    "zpl": "zimpl",
+    "desktop": "desktop",
+    "desktopin": "desktop",
+    "ec": "ec",
+    "eh": "ec",
+    "edn": "edn",
+    "fish": "fish",
+    "mu": "mupad",
+    "nc": "nesc",
+    "ooc": "ooc",
+    "rst": "restructuredtext",
+    "rest": "restructuredtext",
+    "resttxt": "restructuredtext",
+    "rsttxt": "restructuredtext",
+    "wisp": "wisp",
+    "prg": "xbase",
+    "prw": "xbase"
+}
\ No newline at end of file
diff --git a/python/dolma/taggers/__init__.py b/python/dolma/taggers/__init__.py
index 00b3599d..ab99a2cc 100644
--- a/python/dolma/taggers/__init__.py
+++ b/python/dolma/taggers/__init__.py
@@ -7,6 +7,7 @@
     length,
     pii,
     punctuation,
+    repetitions,
     sampling,
     tokenizers,
 )
diff --git a/python/dolma/taggers/code.py b/python/dolma/taggers/code.py
deleted file mode 100644
index 714ca33a..00000000
--- a/python/dolma/taggers/code.py
+++ /dev/null
@@ -1,199 +0,0 @@
-"""
-
-Code-related taggers.
-
-@akshitab
-
-"""
-import logging
-import re
-from typing import Generator, List
-
-import numpy as np
-import regex
-from detect_secrets import SecretsCollection
-from detect_secrets.core.scan import (
-    PotentialSecret,
-    _process_line_based_plugins,
-    get_plugins,
-)
-from detect_secrets.settings import default_settings
-
-from ..core.data_types import DocResult, Document, Span
-from ..core.registry import TaggerRegistry
-from ..core.taggers import BaseTagger
-
-logger = logging.getLogger(__name__)
-
-
-def scan_code(code: str) -> Generator[PotentialSecret, None, None]:
-    if not get_plugins():
-        logger.error("No plugins to scan with!")
-        return
-
-    has_secret = False
-    for lines in [code.splitlines()]:
-        for secret in _process_line_based_plugins(
-            lines=list(enumerate(lines, start=1)),
-            filename="code_str.yml",
-        ):
-            has_secret = True
-            yield secret
-
-        if has_secret:
-            break
-
-
-class SecretsCollectionForStringInput(SecretsCollection):
-    def scan_str(self, code_str: str):
-        for secret in scan_code(code_str):
-            self["code_str.yml"].add(secret)
-
-
-def get_secrets(code: str):
-    secrets = SecretsCollectionForStringInput()
-    with default_settings():
-        secrets.scan_str(code)
-
-    return secrets
-
-
-@TaggerRegistry.add("code_secrets_v1")
-class CodeSecretsTagger(BaseTagger):
-    @classmethod
-    def _extract_code_secrets(cls, text: str) -> List[Span]:
-        secrets_spans: List[Span] = []
-
-        text_lines = text.splitlines()
-        secrets = get_secrets(text)
-        for _, secret in secrets:
-            line_number = secret.line_number - 1
-            span = secret.secret_value
-            span_line = text_lines[line_number]
-            line_start = text.find(span_line)
-            start = line_start + span_line.find(span)
-            end = start + len(span)
-            assert text[start:end] == span
-            secret_type = secret.type.replace(" ", "_")
-            secrets_spans.append(Span(start=start, end=end, type=f"SECRET_{secret_type}"))  # , span])
-
-        return secrets_spans
-
-    def predict(self, doc: Document) -> DocResult:
-        """Main runner."""
-        spans = self._extract_code_secrets(doc.text)
-
-        # document-level score
-        score = self._score(text=doc.text, secrets_spans=spans)
-        spans.append(Span(start=0, end=len(doc.text), type="doc", score=score))
-        return DocResult(doc=doc, spans=spans)
-
-    def _score(self, text: str, secrets_spans: List[Span]) -> float:
-        try:
-            score = len(secrets_spans) * 1.0 / len(text.split())
-        except ZeroDivisionError:
-            score = -1.0
-        return score
-
-
-@TaggerRegistry.add("code_copyright_comments_v1")
-class CodeCopyrightTagger(BaseTagger):
-    """
-    Based on RedPajama code filtering.
-    """
-
-    def __init__(self):
-        self.cpat = re.compile("copyright", re.IGNORECASE)
-        self.pat = re.compile("/\\*[^*]*\\*+(?:[^/*][^*]*\\*+)*/")
-
-    def _extract_copyright_spans(self, text: str) -> List[Span]:
-        copyright_spans: List[Span] = []
-
-        reg = self.pat.search(text)
-
-        if reg:
-            # found one, now see if it contains "copyright", if so strip it
-            span = reg.span()
-            sub = text[span[0] : span[1]]
-            if self.cpat.search(sub):
-                copyright_spans.append(Span(start=span[0], end=span[1], type="copyright_notice", score=1.0))
-            return copyright_spans
-
-        lines = text.split("\n")
-        skip = 0
-        # Greedy replace any file that begins with comment block, most
-        # are copyright headers
-        end = 0
-        for k in range(len(lines)):
-            if lines[k].startswith("//") or lines[k].startswith("#") or lines[k].startswith("--") or not lines[k]:
-                skip = skip + 1
-                if not lines[k]:
-                    end += 1
-                else:
-                    end += len(lines[k])
-            else:
-                break
-
-        if skip:
-            copyright_spans.append(Span(start=0, end=end, type="comment_block", score=1.0))
-        return copyright_spans
-
-    def predict(self, doc: Document) -> DocResult:
-        """Main runner."""
-        spans = self._extract_copyright_spans(doc.text)
-
-        # document-level score
-        score = self._score(text=doc.text, copyright_spans=spans)
-        spans.append(Span(start=0, end=len(doc.text), type="doc", score=score))
-        return DocResult(doc=doc, spans=spans)
-
-    def _score(self, text: str, copyright_spans: List[Span]) -> float:
-        try:
-            if len(copyright_spans) == 0:
-                score = 0.0
-            else:
-                span = copyright_spans[0]
-                # percentage of content affected
-                score = (span.end - span.start + 1) * 1.0 / len(text)
-        except ZeroDivisionError:
-            score = -1.0
-        return score
-
-
-@TaggerRegistry.add("code_redpajama_taggers_v1")
-class CodeRedPajamaTaggers(BaseTagger):
-    """
-    Based on RedPajama code filtering.
-    """
-
-    WHITESPACE_REGEX = regex.compile(r"\w+|[^\w\s]+")
-
-    def _get_num_tokens(self, text: str):
-        return len(self.WHITESPACE_REGEX.split(text))
-
-    def predict(self, doc: Document) -> DocResult:
-        """Main runner."""
-
-        spans: List[Span] = []
-
-        doc_length = len(doc.text)
-
-        line_lengths = list(map(len, doc.text.splitlines()))
-
-        max_line_length = max(line_lengths, default=0.0)
-        avg_line_length = np.mean(line_lengths) if len(line_lengths) > 0 else 0.0
-
-        alnum_count = sum(map(lambda char: 1 if char.isalnum() else 0, doc.text))
-        alnum_prop = (alnum_count / doc_length) if doc_length > 0 else 0.0
-
-        num_tokens = self._get_num_tokens(doc.text)
-        num_alpha = len([c for c in doc.text if c.isalpha()])
-        alpha_token_prop = (num_alpha / num_tokens) if num_tokens > 0 else 0.0
-
-        # document-level scores
-        spans.append(Span(start=0, end=doc_length, type="max_line_length_doc", score=max_line_length))
-        spans.append(Span(start=0, end=doc_length, type="avg_line_length_doc", score=avg_line_length))
-        spans.append(Span(start=0, end=doc_length, type="alnum_prop_doc", score=alnum_prop))
-        spans.append(Span(start=0, end=doc_length, type="alpha_token_prop_doc", score=alpha_token_prop))
-
-        return DocResult(doc=doc, spans=spans)
diff --git a/python/dolma/taggers/code/__init__.py b/python/dolma/taggers/code/__init__.py
new file mode 100644
index 00000000..13117d6c
--- /dev/null
+++ b/python/dolma/taggers/code/__init__.py
@@ -0,0 +1,15 @@
+from .code_taggers import (
+    CodeCopyrightTagger,
+    CodeRedPajamaTaggers,
+    CodeSecretsTagger,
+    CodeStarCoderTaggers,
+    CodeStarCoderTaggers2,
+)
+
+__all__ = [
+    "CodeSecretsTagger",
+    "CodeCopyrightTagger",
+    "CodeRedPajamaTaggers",
+    "CodeStarCoderTaggers",
+    "CodeStarCoderTaggers2",
+]
diff --git a/python/dolma/taggers/code/code_taggers.py b/python/dolma/taggers/code/code_taggers.py
new file mode 100644
index 00000000..feefbefd
--- /dev/null
+++ b/python/dolma/taggers/code/code_taggers.py
@@ -0,0 +1,270 @@
+"""
+
+Code-related taggers.
+
+@akshitab
+
+"""
+import logging
+import re
+from typing import List
+
+import numpy as np
+from necessary import necessary
+
+from ...core.data_types import DocResult, Document, DocumentWithMetadata, Span
+from ...core.registry import TaggerRegistry
+from ...core.taggers import BaseTagger, BaseTaggerWithMetadata
+
+with necessary(["detect_secrets", "bs4", "regex", "pygments"], soft=True) as CODE_DEPENDENCIES_AVAILABLE:
+    if CODE_DEPENDENCIES_AVAILABLE:
+        from .starcoder import get_nl_ratio
+        from .utils import (
+            filter_html,
+            get_ext_to_lang_mapping,
+            get_secrets,
+            get_whitespace_regex,
+        )
+
+
+logger = logging.getLogger(__name__)
+
+
+def check_code_dependencies() -> None:
+    """Check if code dependencies are available."""
+    if not CODE_DEPENDENCIES_AVAILABLE:
+        raise RuntimeError("Code dependencies are not available; please run `pip install dolma[code]`")
+
+
+@TaggerRegistry.add("code_secrets_v1")
+class CodeSecretsTagger(BaseTagger):
+    def __init__(self) -> None:
+        check_code_dependencies()
+        super().__init__()
+
+    @classmethod
+    def _extract_code_secrets(cls, text: str) -> List[Span]:
+        secrets_spans: List[Span] = []
+
+        text_lines = text.splitlines()
+        secrets = get_secrets(text)
+        for _, secret in secrets:
+            line_number = secret.line_number - 1
+            span = secret.secret_value
+            span_line = text_lines[line_number]
+            line_start = text.find(span_line)
+            start = line_start + span_line.find(span or "")
+            end = start + len(span or "")
+            assert text[start:end] == span
+            secret_type = secret.type.replace(" ", "_")
+            secrets_spans.append(Span(start=start, end=end, type=f"SECRET_{secret_type}"))  # , span])
+
+        return secrets_spans
+
+    def predict(self, doc: Document) -> DocResult:
+        """Main runner."""
+        spans = self._extract_code_secrets(doc.text)
+
+        # document-level score
+        score = self._score(text=doc.text, secrets_spans=spans)
+        spans.append(Span(start=0, end=len(doc.text), type="doc", score=score))
+        return DocResult(doc=doc, spans=spans)
+
+    def _score(self, text: str, secrets_spans: List[Span]) -> float:
+        try:
+            score = len(secrets_spans) * 1.0 / len(text.split())
+        except ZeroDivisionError:
+            score = -1.0
+        return score
+
+
+@TaggerRegistry.add("code_copyright_comments_v1")
+class CodeCopyrightTagger(BaseTagger):
+    """
+    Based on RedPajama code filtering.
+    """
+
+    def __init__(self):
+        check_code_dependencies()
+        self.cpat = re.compile("copyright", re.IGNORECASE)
+        self.pat = re.compile("/\\*[^*]*\\*+(?:[^/*][^*]*\\*+)*/")
+
+    def _extract_copyright_spans(self, text: str) -> List[Span]:
+        copyright_spans: List[Span] = []
+
+        reg = self.pat.search(text)
+
+        if reg:
+            # found one, now see if it contains "copyright", if so strip it
+            span = reg.span()
+            sub = text[span[0] : span[1]]
+            if self.cpat.search(sub):
+                copyright_spans.append(Span(start=span[0], end=span[1], type="copyright_notice", score=1.0))
+            return copyright_spans
+
+        lines = text.split("\n")
+        skip = 0
+        # Greedy replace any file that begins with comment block, most
+        # are copyright headers
+        end = 0
+        for line in lines:
+            if line.startswith("//") or line.startswith("#") or line.startswith("--") or not line:
+                skip = skip + 1
+                if not line:
+                    end += 1
+                else:
+                    end += len(line)
+            else:
+                break
+
+        if skip:
+            copyright_spans.append(Span(start=0, end=end, type="comment_block", score=1.0))
+        return copyright_spans
+
+    def predict(self, doc: Document) -> DocResult:
+        """Main runner."""
+        spans = self._extract_copyright_spans(doc.text)
+
+        # document-level score
+        score = self._score(text=doc.text, copyright_spans=spans)
+        spans.append(Span(start=0, end=len(doc.text), type="doc", score=score))
+        return DocResult(doc=doc, spans=spans)
+
+    def _score(self, text: str, copyright_spans: List[Span]) -> float:
+        try:
+            if len(copyright_spans) == 0:
+                score = 0.0
+            else:
+                span = copyright_spans[0]
+                # percentage of content affected
+                score = (span.end - span.start + 1) * 1.0 / len(text)
+        except ZeroDivisionError:
+            score = -1.0
+        return score
+
+
+@TaggerRegistry.add("code_redpajama_taggers_v1")
+class CodeRedPajamaTaggers(BaseTagger):
+    """
+    Based on RedPajama code filtering.
+    """
+
+    def __init__(self):
+        check_code_dependencies()
+        self.whitespace_regex = get_whitespace_regex()
+        super().__init__()
+
+    def _get_num_tokens(self, text: str) -> int:
+        return len(self.whitespace_regex.split(text))
+
+    def predict(self, doc: Document) -> DocResult:
+        """Main runner."""
+
+        spans: List[Span] = []
+
+        doc_length = len(doc.text)
+
+        line_lengths = list(map(len, doc.text.splitlines()))
+
+        max_line_length = max(line_lengths, default=0.0)
+        avg_line_length = float(np.mean(line_lengths) if len(line_lengths) > 0 else 0.0)
+
+        alnum_count = sum(map(lambda char: 1 if char.isalnum() else 0, doc.text))
+        alnum_prop = (alnum_count / doc_length) if doc_length > 0 else 0.0
+
+        num_tokens = self._get_num_tokens(doc.text)
+        num_alpha = len([c for c in doc.text if c.isalpha()])
+        alpha_token_prop = float((num_alpha / num_tokens) if num_tokens > 0 else 0.0)
+
+        # document-level scores
+        spans.append(Span(start=0, end=doc_length, type="max_line_length_doc", score=max_line_length))
+        spans.append(Span(start=0, end=doc_length, type="avg_line_length_doc", score=avg_line_length))
+        spans.append(Span(start=0, end=doc_length, type="alnum_prop_doc", score=alnum_prop))
+        spans.append(Span(start=0, end=doc_length, type="alpha_token_prop_doc", score=alpha_token_prop))
+
+        return DocResult(doc=doc, spans=spans)
+
+
+@TaggerRegistry.add("code_starcoder_taggers_v1")
+class CodeStarCoderTaggers(BaseTaggerWithMetadata):
+    """
+    Based on StarCoder filtering.
+    """
+
+    def __init__(self) -> None:
+        check_code_dependencies()
+        self.ext_to_lang_mapping = get_ext_to_lang_mapping()
+        super().__init__()
+
+    def predict(self, doc: DocumentWithMetadata) -> DocResult:  # type: ignore
+        spans: List[Span] = []
+        doc_length = len(doc.text)
+
+        has_xml_template = 1.0 if "<?xml version=" in doc.text[:100] else 0.0
+        num_github_stars = doc.metadata.get("max_stars_count", 0) or 0
+
+        try:
+            lang = self.ext_to_lang_mapping[doc.metadata.get("ext", "-no-lang")]
+            nl_ratio = get_nl_ratio(doc.text, lang)
+
+            if lang == "html":
+                code_to_text_ratio = filter_html(doc.text)
+            else:
+                # Not relevant for non-html code
+                code_to_text_ratio = 1.0
+        except:  # pylint: disable=bare-except   # noqa: E722
+            nl_ratio = -1.0
+            code_to_text_ratio = -1.0
+
+        # document-level scores
+        spans.append(Span(start=0, end=doc_length, type="has_xml_template_doc", score=has_xml_template))
+        spans.append(Span(start=0, end=doc_length, type="num_github_stars_doc", score=float(num_github_stars)))
+        spans.append(Span(start=0, end=doc_length, type="nl_ratio_doc", score=nl_ratio))
+        spans.append(Span(start=0, end=doc_length, type="code_to_text_ratio_html_doc", score=code_to_text_ratio))
+
+        return DocResult(doc=doc, spans=spans)
+
+
+@TaggerRegistry.add("code_starcoder_taggers_v2")
+class CodeStarCoderTaggers2(BaseTaggerWithMetadata):
+    """
+    Based on StarCoder filtering.
+    """
+
+    def __init__(self) -> None:
+        check_code_dependencies()
+        self.ext_to_lang_mapping = get_ext_to_lang_mapping()
+        super().__init__()
+
+    def predict(self, doc: DocumentWithMetadata) -> DocResult:  # type: ignore
+        spans: List[Span] = []
+        doc_length = len(doc.text)
+
+        has_xml_template = 1.0 if "<?xml version=" in doc.text[:100] else 0.0
+        num_github_stars = doc.metadata.get("max_stars_count", 0) or 0
+
+        try:
+            lang = self.ext_to_lang_mapping[doc.metadata.get("ext", "-no-lang")]
+        except KeyError:
+            lang = "-no-lang"
+
+        if lang in ["python", "java", "javascript"]:
+            code_to_comment_ratio = get_nl_ratio(doc.text, lang)
+        else:
+            code_to_comment_ratio = 0.5  # We use an upper and lower bound of filters; this is in the middle.
+
+        if lang == "html":
+            try:
+                code_to_text_ratio = filter_html(doc.text)
+            except:  # pylint: disable=bare-except   # noqa: E722
+                code_to_text_ratio = -1.0
+        else:
+            code_to_text_ratio = 1.0
+
+        # document-level scores
+        spans.append(Span(start=0, end=doc_length, type="has_xml_template_doc", score=has_xml_template))
+        spans.append(Span(start=0, end=doc_length, type="num_github_stars_doc", score=float(num_github_stars)))
+        spans.append(Span(start=0, end=doc_length, type="code_to_comment_ratio_doc", score=code_to_comment_ratio))
+        spans.append(Span(start=0, end=doc_length, type="code_to_text_ratio_html_doc", score=code_to_text_ratio))
+
+        return DocResult(doc=doc, spans=spans)
diff --git a/python/dolma/taggers/code/starcoder.py b/python/dolma/taggers/code/starcoder.py
new file mode 100644
index 00000000..8a58d40a
--- /dev/null
+++ b/python/dolma/taggers/code/starcoder.py
@@ -0,0 +1,129 @@
+"""
+Copied from https://github.com/bigcode-project/bigcode-dataset/blob/main/preprocessing/utils/text_extraction.py
+
+Extract Python comments (using Python tokenizer) and
+docstrings (using AST parsing).
+Extract Java and JavaScript comments (using Pygments)
+"""
+
+import ast
+import io
+import tokenize
+import warnings
+from itertools import groupby
+from os.path import basename, splitext
+
+import pygments
+from pygments.lexers import get_lexer_by_name
+from pygments.token import Comment
+
+StringIO = io.StringIO
+
+NODE_TYPES = {
+    ast.ClassDef: "Class",
+    ast.FunctionDef: "Function/Method",
+    ast.Module: "Module",
+}
+
+
+# comment extraction
+def get_comments(s, clean=False):
+    "Returns a string including all comments in python code"
+    coments = []
+    g = tokenize.generate_tokens(StringIO(s).readline)
+    for toknum, tokval, _, _, _ in g:
+        # print(toknum,tokval)
+        if toknum == tokenize.COMMENT:
+            coments.append((toknum, tokval))
+    result = tokenize.untokenize(coments)
+    if clean:
+        result = result.replace("#", "")
+    return result
+
+
+# Note: sometimes this can miss examples with decorators over classes
+# ast parsing, source: https://gist.github.com/SpotlightKid/1548cb6c97f2a844f72d
+def parse_docstrings(source):
+    """Parse Python source code and yield a tuple of ast node instance, name,
+    and docstring for each function/method, class and module."""
+    tree = ast.parse(source)
+
+    for node in ast.walk(tree):
+        if isinstance(node, tuple(NODE_TYPES)):
+            docstring = ast.get_docstring(node)  # pyright: ignore
+
+            yield (node, getattr(node, "name", None), docstring)
+
+
+def get_docstrings(source, module="<string>"):
+    """Parse Python source code from file or string and print docstrings."""
+    if hasattr(source, "read"):
+        filename = getattr(source, "name", module)
+        module = splitext(basename(filename))[0]
+        source = source.read()
+
+    docstrings = sorted(
+        parse_docstrings(source), key=lambda x: (NODE_TYPES.get(type(x[0])), x[1])  # pyright: ignore
+    )
+
+    grouped = groupby(docstrings, key=lambda x: NODE_TYPES.get(type(x[0])))  # pyright: ignore
+    results = []
+    for _, group in grouped:
+        for _, name, docstring in group:
+            name = name if name else module
+            # print(docstring or '')
+            if docstring:
+                results.append(docstring)
+    return results
+
+
+def get_text_python(source, comments=True, clean_comments=True):
+    """Extract all natural text in source: comments + docstrings
+    the extraction fails in case of syntax errors in the file
+    Args:
+        source: the code to parse
+        comments: if True extract comments two
+        clean_comment: if True remove # from extracted comments
+    Returns:
+        a string with concatenated docstrings and comments"""
+
+    try:
+        docstrings = "\n".join(get_docstrings(source))
+    except:  # pylint: disable=bare-except   # noqa: E722
+        docstrings = ""
+        warnings.warn("code couldn't be parsed due to compilation failure, no docstring is extracted")
+
+    if comments:
+        try:
+            comments = get_comments(source, clean=clean_comments)
+        except:  # pylint: disable=bare-except    # noqa: E722
+            comments = ""
+            warnings.warn("tokenization error, no comment is extracted")
+    else:
+        comments = ""
+
+    output = docstrings + "\n" + comments
+    return output.strip()
+
+
+def comment_size(text, language):
+    """
+    Compute the size of comments in a program (not necessarily python).
+    """
+    lexer = get_lexer_by_name(language)
+    tokens = pygments.lex(text, lexer)
+    comment_len = 0
+    for token_type, token in tokens:
+        if token_type == Comment.Multiline or token_type == Comment.Single:
+            comment_len += len(token)  # token is a string with the comment contents
+    return comment_len
+
+
+def get_nl_ratio(text, language):
+    """get the ratio of comments to code in a program"""
+    if language == "python":
+        comments = get_text_python(text)
+        ratio = len(comments) / len(text)
+    else:
+        ratio = comment_size(text, language) / len(text)
+    return ratio
diff --git a/python/dolma/taggers/code/utils.py b/python/dolma/taggers/code/utils.py
new file mode 100644
index 00000000..e7c2c474
--- /dev/null
+++ b/python/dolma/taggers/code/utils.py
@@ -0,0 +1,81 @@
+"""
+
+Utilities for code-related taggers.
+
+@akshitab, @soldni
+
+"""
+import json
+import logging
+from pathlib import Path
+from typing import Dict, Generator
+
+import regex
+import smart_open
+from bs4 import BeautifulSoup  # pylint: disable=import-error
+from detect_secrets.core.potential_secret import PotentialSecret
+from detect_secrets.core.scan import _process_line_based_plugins
+from detect_secrets.core.secrets_collection import SecretsCollection
+from detect_secrets.settings import default_settings, get_plugins
+
+logger = logging.getLogger(__name__)
+
+
+def scan_code(code: str) -> Generator["PotentialSecret", None, None]:
+    if not get_plugins():
+        logger.error("No plugins to scan with!")
+        return
+
+    has_secret = False
+    for lines in [code.splitlines()]:
+        for secret in _process_line_based_plugins(
+            lines=list(enumerate(lines, start=1)),
+            filename="code_str.yml",
+        ):
+            has_secret = True
+            yield secret
+
+        if has_secret:
+            break
+
+
+class SecretsCollectionForStringInput(SecretsCollection):
+    def scan_str(self, code_str: str):
+        for secret in scan_code(code_str):
+            self["code_str.yml"].add(secret)
+
+
+def get_secrets(code: str):
+    secrets = SecretsCollectionForStringInput()
+    with default_settings():
+        secrets.scan_str(code)
+
+    return secrets
+
+
+def filter_html(html: str) -> float:
+    """Filter HTML files based on displayed text VS code ratio"""
+    try:
+        soup = BeautifulSoup(html, features="html.parser")
+    except (TypeError, UnboundLocalError):
+        return False
+
+    # kill all script and style elements
+    for script in soup(["script", "style"]):
+        script.extract()  # rip it out
+
+    # get text
+    text = soup.get_text()
+    ratio = len(text) / len(html)
+
+    return (ratio) * (len(text) > 100)
+
+
+def get_whitespace_regex() -> regex.Pattern:
+    return regex.compile(r"\w+|[^\w\s]+")
+
+
+def get_ext_to_lang_mapping() -> Dict[str, str]:
+    path = Path(__file__).parent / "../../data/ext_to_lang_mapping.json"
+    with smart_open.open(path, "r") as f:
+        return json.load(f)
diff --git a/python/dolma/taggers/length.py b/python/dolma/taggers/length.py
index 995d8f00..4dbc649b 100644
--- a/python/dolma/taggers/length.py
+++ b/python/dolma/taggers/length.py
@@ -25,6 +25,13 @@ def predict(self, doc: Document) -> DocResult:
         return DocResult(doc=doc, spans=[Span(start=0, end=len(doc.text), type="length", score=score)])
 
 
+@TaggerRegistry.add("char_length_strip_ws_v1")
+class CharLengthStripWsV1(BaseTagger):
+    def predict(self, doc: Document) -> DocResult:
+        score = len(doc.text.strip())
+        return DocResult(doc=doc, spans=[Span(start=0, end=len(doc.text), type="length_no_ws", score=score)])
+
+
 @TaggerRegistry.add("char_length_with_paragraphs_v1")
 class CharLengthWithParagraphsV1(BaseTagger):
     def predict(self, doc: Document) -> DocResult:
diff --git a/python/dolma/taggers/pii.py b/python/dolma/taggers/pii.py
index 34f47ae4..afa7c8df 100644
--- a/python/dolma/taggers/pii.py
+++ b/python/dolma/taggers/pii.py
@@ -6,18 +6,16 @@
 
 """
 
-try:
-    import re2 as re
-except ImportError:
-    import re
-else:
-    re.set_fallback_notification(re.FALLBACK_WARNING)
-
 
+import re
 from typing import List
 from warnings import warn
 
-from presidio_analyzer import AnalyzerEngine
+from necessary import necessary
+
+with necessary("presidio-analyzer", soft=True) as PRESIDIO_AVAILABLE:
+    if PRESIDIO_AVAILABLE:
+        from presidio_analyzer import AnalyzerEngine  # pylint: disable=import-error # pyright: ignore
 
 from ..core.data_types import DocResult, Document, Span, TextSlice
 from ..core.registry import TaggerRegistry
@@ -70,6 +68,8 @@ def __init__(
 
         # presidio
         if self.method == self.PRESIDIO:
+            if not PRESIDIO_AVAILABLE:
+                raise RuntimeError("Presidio is not available; please run `pip install dolma[pii]`")
             self.analyzer = AnalyzerEngine()
 
     def predict(self, doc: Document) -> DocResult:
diff --git a/python/dolma/taggers/repetitions/__init__.py b/python/dolma/taggers/repetitions/__init__.py
new file mode 100644
index 00000000..3d941851
--- /dev/null
+++ b/python/dolma/taggers/repetitions/__init__.py
@@ -0,0 +1,7 @@
+from .repetitions_taggers import (
+    ParagraphRepetitionsTagger,
+    RepetitionsTagger,
+    TokenizerRepetitionsTagger,
+)
+
+__all__ = ["RepetitionsTagger", "ParagraphRepetitionsTagger", "TokenizerRepetitionsTagger"]
diff --git a/python/dolma/taggers/repetitions/repetitions_taggers.py b/python/dolma/taggers/repetitions/repetitions_taggers.py
new file mode 100644
index 00000000..44bc2bbe
--- /dev/null
+++ b/python/dolma/taggers/repetitions/repetitions_taggers.py
@@ -0,0 +1,136 @@
+"""
+Taggers to detect repetitions in the text.
+
+@soldni
+"""
+
+import re
+from abc import abstractmethod
+from typing import Generator, List
+
+import numpy as np
+from tokenizers import Tokenizer
+
+from ...core.data_types import DocResult, Document, Span
+from ...core.registry import TaggerRegistry
+from ...core.taggers import BaseTagger
+from ...core.utils import split_paragraphs
+from .utils import find_periodic_sequences
+
+
+class BaseRepetitionsTagger(BaseTagger):
+    @abstractmethod
+    def _extract_from_text(self, text: str) -> Generator[Span, None, None]:
+        raise NotImplementedError()
+
+    def _extract_from_doc(self, doc: Document) -> Generator[Span, None, None]:
+        yield from self._extract_from_text(doc.text)
+
+    def _compute_document_stats(self, spans: List[Span], doc: Document) -> List[Span]:
+        doc_max_span = Span(
+            start=0,
+            end=len(doc.text),
+            type="doc_max_repetition",
+            score=max(spans, key=lambda s: s.score).score if spans else 0.0,
+        )
+        doc_mean_reps_span = Span(
+            start=0,
+            end=len(doc.text),
+            type="doc_mean_repetition",
+            score=float(np.mean([s.score for s in spans]) if spans else 0),
+        )
+        doc_frac_reps_span = Span(
+            start=0,
+            end=len(doc.text),
+            type="doc_frac_repetition",
+            score=float(sum([s.score for s in spans]) / len(doc.text) if spans else 0),
+        )
+        return [doc_max_span, doc_mean_reps_span, doc_frac_reps_span]
+
+    def predict(self, doc: Document) -> DocResult:
+        """Predict method for the tagger."""
+        reps_spans = list(self._extract_from_doc(doc))
+        document_stats_spans = self._compute_document_stats(spans=reps_spans, doc=doc)
+        return DocResult(doc=doc, spans=reps_spans + document_stats_spans)
+
+
+@TaggerRegistry.add("repetitions_v1")
+class RepetitionsTagger(BaseRepetitionsTagger):
+    """Tagger to detect repetitions of of groups of characters.
+    Only repetitions that occur at least 4 times are detected."""
+
+    def __init__(self) -> None:
+        self.re_char_repetitions = re.compile(r"(.+?)(\s?\1){3,}")
+        super().__init__()
+
+    def _extract_from_text(self, text: str) -> Generator[Span, None, None]:
+        """Extract repetitions of characters in the text."""
+        for match in self.re_char_repetitions.finditer(text):
+            yield Span(
+                start=(start := match.start()),
+                end=(end := match.end()),
+                type="repetition",
+                score=float(end - start),
+            )
+
+
+@TaggerRegistry.add("paragraph_repetitions_v1")
+class ParagraphRepetitionsTagger(RepetitionsTagger):
+    """Tagger to detect repetitions of paragraphs.
+    It's faster than the char repetition tagger, but it does not account for
+    repetitions of characters that span multiple paragraphs."""
+
+    def _extract_from_doc(self, doc: Document) -> Generator[Span, None, None]:
+        offset = 0
+        for paragraph in split_paragraphs(doc.text, remove_empty=False):
+            for span in self._extract_from_text(paragraph.text):
+                span.start += offset
+                span.end += offset
+                yield span
+            offset += len(paragraph.text)
+
+
+@TaggerRegistry.add("tokenizer_repetitions_v1")
+class TokenizerRepetitionsTagger(BaseRepetitionsTagger):
+    """Tagger to detect repetitions of tokens.
+
+    It uses a tokenizer to split the text into tokens, and then identifies
+    sequences of tokens that repeat at least 3 times."""
+
+    TOKENIZER_IDENTIFIER = "allenai/eleuther-ai-gpt-neox-20b-pii-special"
+    MIN_PERIOD = 1
+    MAX_PERIOD = 13
+
+    def __init__(self) -> None:
+        self.tokenizer = Tokenizer.from_pretrained(self.TOKENIZER_IDENTIFIER)
+
+    def _extract_from_text(self, text: str) -> Generator[Span, None, None]:
+        tokens = self.tokenizer.encode(text, add_special_tokens=False)
+        sequences_iter = find_periodic_sequences(
+            arr=np.array(tokens.ids), min_period=self.MIN_PERIOD, max_period=self.MAX_PERIOD
+        )
+        for seq in sequences_iter:
+            yield Span(
+                start=(s := tokens.offsets[seq.start][0]),
+                end=(e := tokens.offsets[seq.end - 1][1]),
+                type="repetition",
+                score=float(e - s),
+            )
+
+
+@TaggerRegistry.add("paragraph_tokenizer_repetitions_v1")
+class ParagraphTokenizerRepetitionsTagger(TokenizerRepetitionsTagger):
+    """Tagger to detect repetitions of tokens in paragraphs.
+    It's faster than the tokenizer repetition tagger, but it does not account for
+    repetitions of tokens that span multiple paragraphs."""
+
+    def _extract_from_doc(self, doc: Document) -> Generator[Span, None, None]:
+        offset = 0
+        for paragraph in split_paragraphs(doc.text, remove_empty=False):
+            # space is required to avoid first symbol in the paragraph to be
+            # tokenized as a different token.
+            for span in self._extract_from_text(" " + paragraph.text):
+                span.start += offset - 1
+                span.end += offset - 1
+                yield span
+            offset += len(paragraph.text)
diff --git a/python/dolma/taggers/repetitions/utils.py b/python/dolma/taggers/repetitions/utils.py
new file mode 100644
index 00000000..e2fb9973
--- /dev/null
+++ b/python/dolma/taggers/repetitions/utils.py
@@ -0,0 +1,118 @@
+from typing import Generator, List, NamedTuple
+
+import numpy as np
+
+
+def find_end_first_consecutive_true(arr: np.ndarray) -> int:
+    """Function to find the end position of the first consecutive sequence of True in an array."""
+    if not arr[0]:
+        return 0
+
+    prog = np.cumsum(arr)
+    if prog[-1] == len(arr):
+        return len(arr)
+
+    true_locs = np.where(prog[:-1:] == prog[1::])[0]
+
+    return true_locs[0] + 1
+
+
+def find_start_last_consecutive_true(arr: np.ndarray) -> int:
+    """Function to find the start position of the last consecutive sequence of True in an array."""
+    reverse = find_end_first_consecutive_true(arr[::-1])
+    return len(arr) - reverse if reverse > 0 else -1
+
+
+def group_consecutive_values(arr: np.ndarray, stepsize: int = 1) -> List[np.ndarray]:
+    """Function to group consecutive values in an array."""
+    return np.split(arr, np.where(np.diff(arr) != stepsize)[0] + 1)
+
+
+class RepetitionTuple(NamedTuple):
+    """Tuple to store information about a periodic sequence."""
+
+    start: int
+    end: int
+    period: int
+    times: int
+
+
+def find_periodic_sequences(
+    arr: np.ndarray, max_period: int, min_period: int = 1, mask_value: int = -1
+) -> Generator[RepetitionTuple, None, None]:
+    """Function to find periodic sequences in an array.
+
+    This function sweeps through the array and checks for sequences of length
+    [min_period, max_period] that repeat at least 3 times. To do so, it
+    reshape the array into a matrix with `period` columns and checks if each
+    row is equal to the previous row. Blocks of repeating rows indicates repeating
+    sequences.
+
+    Because there's no guarantee that the sequences start at the beginning of each
+    row, it can only detect sequences that repeat at least 3 times. To account
+    for the fact that sequences may not start at the beginning of each row (or
+    end at the end of each row), we check the end of the previous row and the
+    start of the next row to determine the actual start and end positions of the
+    sequence.
+
+    Args:
+        arr (np.ndarray): The array to search for periodic sequences.
+        max_period (int): The maximum period to check for.
+        min_period (int, optional): The minimum period to check for. Defaults to 1.
+        mask_value (int, optional): The value to use to pad the array. Defaults to -1.
+    """
+    # make sure the mask_value is not in the array
+    if (arr == mask_value).sum() > 0:
+        raise ValueError("`mask_value` is in the array")
+
+    # no since we can only detect sequences that repeat at least 3 times,
+    # there is no point in checking for periods greater than 1/3 of the length
+    max_period = min(max_period, len(arr) // 3)
+
+    for period in range(min_period, max_period + 1):
+        # pad the array so that it can be reshaped into a matrix matching the period
+        padded_arr = np.pad(arr, (0, period - (len(arr) % period)), constant_values=mask_value)
+        shaped_arr = padded_arr.reshape(-1, period)
+
+        # find rows that are equal to the previous  row; these are the possibly-periodic sequences
+        is_equal_to_prev_row = shaped_arr == np.roll(shaped_arr, shift=1, axis=0)
+        rows_with_period, *_ = np.where(is_equal_to_prev_row.all(axis=1))
+
+        # no sequences found with this period
+        if len(rows_with_period) == 0:
+            continue
+
+        # this finds the start and end positions of the sequences with period `period`
+        where_true_consecutive = group_consecutive_values(rows_with_period)
+
+        for sequence in where_true_consecutive:
+            start_row = sequence[0]
+            end_row = sequence[-1]
+
+            # we check if any value at the end of the previous row is True, e.g.:
+            #     [[False, False, True, True]
+            #      [True, True, True, True]]
+            # (in the case above, start offset is 2). If so, we subtract that from the
+            # period to get the actual start offset.
+            start_offset = find_start_last_consecutive_true(is_equal_to_prev_row[start_row - 1])
+            start_offset = period - start_offset if start_offset > 0 else 0
+
+            # same idea as above, we want to compute offset. Only difference is that
+            # `find_end_first_consecutive_true` already returns the offset, so we don't
+            # need to subtract from the period.
+            end_offset = find_end_first_consecutive_true(is_equal_to_prev_row[end_row + 1])
+
+            # because we are always comparing with preceding row in
+            # `is_equal_to_prev_row`, we need to subtract 1 from the row number
+            start_pos = (start_row - 1) * period - start_offset
+
+            # note that the end position is exclusive
+            end_pos = ((end_row + 1) * period) + end_offset
+
+            out = RepetitionTuple(
+                start=start_pos, end=end_pos, period=period, times=(end_pos - start_pos) // period
+            )
+            if out.times > 2:
+                # cannot accurately determine the period of a sequence that repeats
+                # less than 3 times with this algorithm
+                yield out
diff --git a/scripts/dolma_paper_plots.sh b/scripts/dolma_paper_plots.sh
new file mode 100644
index 00000000..fdeb777c
--- /dev/null
+++ b/scripts/dolma_paper_plots.sh
@@ -0,0 +1,363 @@
+#!/usr/bin/env bash
+
+# get script directory
+SOURCE="${BASH_SOURCE[0]}"
+while [ -h "$SOURCE" ]; do # resolve $SOURCE until the file is no longer a symlink
+  SCRIPT_DIR="$( cd -P "$( dirname "$SOURCE" )" >/dev/null 2>&1 && pwd )"
+  SOURCE="$(readlink "$SOURCE")"
+  # if $SOURCE was a relative symlink, we need to resolve it
+  # relative to the path where the symlink file was located
+  [[ $SOURCE != /* ]] && SOURCE="$SCRIPT_DIR/$SOURCE"
+done
+SCRIPT_DIR="$( cd -P "$( dirname "$SOURCE" )" >/dev/null 2>&1 && pwd )"
+
+train_metrics=(
+    'train/Perplexity'
+    'train/CrossEntropyLoss'
+)
+TRAIN_METRICS="$(printf "%s " "${train_metrics[@]}" | sed 's/ $//')"
+
+eval_metrics=(
+    'eval/downstream/hellaswag_len_norm'
+    'eval/downstream/piqa_len_norm'
+    'eval/downstream/arc_easy_acc'
+    'eval/downstream/sciq_acc'
+    'eval/downstream/mrpc_f1'
+    'eval/downstream/commitment_bank_acc'
+    'eval/downstream/copa_acc'
+    'eval/downstream/openbook_qa_len_norm'
+    'eval/downstream/winogrande_acc'
+    'eval/downstream/rte_len_norm'
+)
+EVAL_METRICS="$(printf "%s " "${eval_metrics[@]}" | sed 's/ $//')"
+
+code_perplexity_suite=(
+    'eval/openai_humaneval_test/Perplexity'
+    'eval/mbpp_valid/Perplexity'
+    'eval/stack_v2_held_out/Perplexity'
+)
+CODE_PERPLEXITY_SUITE="$(printf "%s " "${code_perplexity_suite[@]}" | sed 's/ $//')"
+
+v1_perplexity_suite=(
+    'eval/4chan-validation/Perplexity'
+    'eval/c4_100_domains-validation/Perplexity'
+    'eval/c4_en-validation/Perplexity'
+    'eval/gab-validation/Perplexity'
+    'eval/ice-validation/Perplexity'
+    'eval/m2d2_s2orc-validation/Perplexity'
+    'eval/m2d2_wiki-validation/Perplexity'
+    'eval/manosphere-validation/Perplexity'
+    'eval/mc4_en-validation/Perplexity'
+    'eval/pile-validation/Perplexity'
+    'eval/twitterAEE-validation/Perplexity'
+)
+V1_PERPLEXITY_SUITE="$(printf "%s " "${v1_perplexity_suite[@]}" | sed 's/ $//')"
+
+og_perplexity_suite=(
+    'eval/gab-validation/Perplexity'
+    'eval/ice-validation/Perplexity'
+    'eval/ptb-validation/Perplexity'
+    'eval/pile-validation/Perplexity'
+    'eval/4chan-validation/Perplexity'
+    'eval/c4_en-validation/Perplexity'
+    'eval/mc4_en-validation/Perplexity'
+    'eval/m2d2_wiki-validation/Perplexity'
+    'eval/m2d2_s2orc-validation/Perplexity'
+    'eval/manosphere-validation/Perplexity'
+    'eval/twitterAEE-validation/Perplexity'
+    'eval/wikitext_103-validation/Perplexity'
+    'eval/c4_100_domains-validation/Perplexity'
+)
+OG_PERPLEXITY_SUITE="$(printf "%s " "${og_perplexity_suite[@]}" | sed 's/ $//')"
+
+v2_v3_perplexity_suite=(
+    'eval/v3-small-ice-validation/Perplexity'
+    'eval/v3-small-dolma_wiki-validation/Perplexity'
+    'eval/v3-small-dolma_books-validation/Perplexity'
+    'eval/v3-small-dolma_pes2o-validation/Perplexity'
+    'eval/v3-small-dolma_stack-validation/Perplexity'
+    'eval/v3-small-dolma_reddit-validation/Perplexity'
+    'eval/v3-small-dolma_common-crawl-validation/Perplexity'
+    'eval/v2-small-gab-validation/Perplexity'
+    'eval/v2-small-ptb-validation/Perplexity'
+    'eval/v2-small-pile-validation/Perplexity'
+    'eval/v2-small-4chan-validation/Perplexity'
+    'eval/v2-small-c4_en-validation/Perplexity'
+    'eval/v2-small-mc4_en-validation/Perplexity'
+    'eval/v2-small-m2d2_wiki-validation/Perplexity'
+    'eval/v2-small-m2d2_s2orc-validation/Perplexity'
+    'eval/v2-small-manosphere-validation/Perplexity'
+    'eval/v2-small-twitterAEE-validation/Perplexity'
+    'eval/v2-small-wikitext_103-validation/Perplexity'
+    'eval/v2-small-c4_100_domains-validation/Perplexity'
+)
+V2_V3_PERPLEXITY_SUITE="$(printf "%s " "${v2_v3_perplexity_suite[@]}" | sed 's/ $//')"
+
+runs_up_to_150b=(
+    'olmo-small-rpj-*'
+    'olmo-small-pile-fixed-*'
+    'olmo-small-dolma-*'
+    'olmo-small-falcon-*'
+)
+RUNS_UP_TO_150B="$(printf "%s " "${runs_up_to_150b[@]}" | sed 's/ $//')"
+
+# base directory for all plots and points to sample
+SAMPLES=1000
+BASE_DIR="${1}"
+
+###############################################################################
+
+# Figure 1: comparison of training curves between diff datasets.
+FIGURE_1_DIR="${BASE_DIR}/150b_runs"
+
+if [ ! -d "${FIGURE_1_DIR}/train" ]; then
+    # only plot if the directory doesn't exist
+    set -ex
+
+    python ${SCRIPT_DIR}/wandb_to_plot.py \
+        -t ai2-llm \
+        -p olmo-small \
+        -n $RUNS_UP_TO_150B \
+        -y $TRAIN_METRICS \
+        -s $SAMPLES \
+        -d ${FIGURE_1_DIR}/train \
+        --max-x-axis '150e9' \
+        --max-y-axis 100 \
+        --y-log-scale \
+        -v ${SCRIPT_DIR}/wandb_run_vocab.yaml
+
+    set +ex
+fi
+
+if [ ! -d "${FIGURE_1_DIR}/downstream" ]; then
+    # only plot if the directory doesn't exist
+    set -ex
+
+    python ${SCRIPT_DIR}/wandb_to_plot.py \
+        -t ai2-llm \
+        -p olmo-small \
+        -n $RUNS_UP_TO_150B \
+        -y $EVAL_METRICS \
+        -s $SAMPLES \
+        -d "${FIGURE_1_DIR}/downstream" \
+        --max-x-axis '150e9' \
+        -v ${SCRIPT_DIR}/wandb_run_vocab.yaml
+
+    set +ex
+
+fi
+
+if [ ! -d "${FIGURE_1_DIR}/ppl" ]; then
+    # only plot if the directory doesn't exist
+    set -ex
+
+    python ${SCRIPT_DIR}/wandb_to_plot.py \
+        -t ai2-llm \
+        -p olmo-small \
+        -n $RUNS_UP_TO_150B \
+        -y $OG_PERPLEXITY_SUITE \
+        -s $SAMPLES \
+        -d "${FIGURE_1_DIR}/ppl" \
+        --max-x-axis '150e9' \
+        --max-y-axis 100 \
+        --y-log-scale \
+        -v ${SCRIPT_DIR}/wandb_run_vocab.yaml
+
+    set +ex
+fi
+
+###############################################################################
+
+# Figure 2: long 1b run (up to 3T tokens if possible)
+FIGURE_2_DIR="${BASE_DIR}/long_1b_run"
+
+if [ ! -d "${FIGURE_2_DIR}/train" ]; then
+    # only plot if the directory doesn't exist
+    set -ex
+
+    python ${SCRIPT_DIR}/wandb_to_plot.py \
+        -t ai2-llm \
+        -p olmo-small \
+        -n 'olmo-small-3T-lower-lr-tie_*' \
+        -y $TRAIN_METRICS \
+        -s $SAMPLES \
+        -d ${FIGURE_2_DIR}/train \
+        --max-y-axis 100 \
+        --y-log-scale \
+        -v ${SCRIPT_DIR}/wandb_run_vocab.yaml
+
+    set +ex
+fi
+
+if [ ! -d "${FIGURE_2_DIR}/downstream" ]; then
+    # only plot if the directory doesn't exist
+    set -ex
+
+    python ${SCRIPT_DIR}/wandb_to_plot.py \
+        -t ai2-llm \
+        -p olmo-small \
+        -n 'olmo-small-3T-lower-lr-tie_*' \
+        -y $EVAL_METRICS \
+        -s $SAMPLES \
+        -d "${FIGURE_2_DIR}/downstream" \
+        -v ${SCRIPT_DIR}/wandb_run_vocab.yaml
+
+    set +ex
+
+fi
+
+if [ ! -d "${FIGURE_2_DIR}/ppl" ]; then
+    # only plot if the directory doesn't exist
+    set -ex
+
+    python ${SCRIPT_DIR}/wandb_to_plot.py \
+        -t ai2-llm \
+        -p olmo-small \
+        -n 'olmo-small-3T-lower-lr-tie_*' \
+        -y $V2_V3_PERPLEXITY_SUITE \
+        -s $SAMPLES \
+        -d "${FIGURE_2_DIR}/ppl" \
+        --max-y-axis 100 \
+        --y-log-scale \
+        -v ${SCRIPT_DIR}/wandb_run_vocab.yaml
+
+    set +ex
+fi
+
+
+###############################################################################
+
+
+
+
+ablations_runs=(
+    'v1-small-pi-less-than-5-anonymize_* v1-small-all-pi-removed_* abl-cc-v1-small-dedup_*'
+    'reddit-v1-ablation-base_* reddit-v1-ablation-pii-nsfw-toxic_filtered_* reddit-v1-ablation-toxic-filtered_*'
+    'olmo-mix-v1-sample_* olmo-mix-v1-sample-all-cc* olmo-mix-v1-sample-mix2_* olmo-mix-v1-gopher-like_*'
+    'stack-v2* stack-v4*'
+    'c4-stack-15p* c4_p85-stack_v4_p15* c4_p85-starcoder_p15*'
+    'v1-small-hatespeech-filtered-low* v1-small-nsfw-filtered-low* v1-small-hatespeech-filtered-high* v1-small-nsfw-filtered-high* abl-cc-v1-small-dedup_*'
+    'abl-cc-v1-small-dedup_* abl-cc-v2-small-dedup*'
+    'abl-cc-v1-small-dedup_* v1-small-c4-cleaned_\d+ v1-small-c4-filtered_\d+ v1-small-gopher-filtered_\d+ v1-small-c4-cleaned-gopher-filtered_\d+ v1-small-c4-cleaned-gopher-filtered-deduped_\d+ olmo-mix-v1-sample-all-cc*'
+)
+ablations_names=(
+    'cc_pii_filtering'
+    'reddit_toxic_filtering'
+    'dolma_mix'
+    'code_stack_v2_vs_v4'
+    'code_15p_stack_v2_v4_starcoder'
+    'cc_toxic_filtering'
+    'cc_dedupe'
+    'cc_quality'
+)
+
+limits=(
+    '150e9'
+    '150e9'
+    '150e9'
+    '50e9'
+    '50e9'
+    '150e9'
+    '150e9'
+    '150e9'
+)
+
+# Loop through the indices of the array.
+for index in "${!ablations_names[@]}"; do
+    # Access the element by its index.
+    ABLATION_NAME="${ablations_names[$index]}"
+    ABLATION_DIR="${BASE_DIR}/ablations_${ABLATION_NAME}"
+    RUNS_ABLATION="${ablations_runs[$index]}"
+    X_AXIS_LIMIT="${limits[$index]}"
+
+    if [ ! -d "${ABLATION_DIR}/train" ]; then
+        # only plot if the directory doesn't exist
+        set -ex
+
+        python ${SCRIPT_DIR}/wandb_to_plot.py \
+            -t ai2-llm \
+            -p c4-small \
+            -n $RUNS_ABLATION \
+            -N $ABLATION_NAME \
+            -y $TRAIN_METRICS \
+            -s $SAMPLES \
+            -d ${ABLATION_DIR}/train \
+            --max-x-axis $X_AXIS_LIMIT \
+            --max-y-axis 100 \
+            --y-log-scale \
+            -v ${SCRIPT_DIR}/wandb_run_vocab.yaml \
+            --plotly-font-size 9 \
+            --plotly-figure-width 400 \
+            --plotly-figure-height 400
+
+        set +ex
+    fi
+
+    if [ ! -d "${ABLATION_DIR}/downstream" ]; then
+        # only plot if the directory doesn't exist
+        set -ex
+
+        python ${SCRIPT_DIR}/wandb_to_plot.py \
+            -t ai2-llm \
+            -p c4-small \
+            -n $RUNS_ABLATION \
+            -N $ABLATION_NAME \
+            -y $EVAL_METRICS \
+            -s $SAMPLES \
+            -d "${ABLATION_DIR}/downstream" \
+            --max-x-axis $X_AXIS_LIMIT \
+            -v ${SCRIPT_DIR}/wandb_run_vocab.yaml \
+            --plotly-font-size 9 \
+            --plotly-figure-width 400 \
+            --plotly-figure-height 400
+
+        set +ex
+
+    fi
+
+    if [ ! -d "${ABLATION_DIR}/ppl" ]; then
+        # only plot if the directory doesn't exist
+        set -ex
+
+        python ${SCRIPT_DIR}/wandb_to_plot.py \
+            -t ai2-llm \
+            -p c4-small \
+            -n $RUNS_ABLATION \
+            -N $ABLATION_NAME \
+            -y $V1_PERPLEXITY_SUITE \
+            -s $SAMPLES \
+            -d "${ABLATION_DIR}/ppl" \
+            --max-x-axis $X_AXIS_LIMIT \
+            --max-y-axis 100 \
+            --y-log-scale \
+            -v ${SCRIPT_DIR}/wandb_run_vocab.yaml \
+            --plotly-font-size 9 \
+            --plotly-figure-width 400 \
+            --plotly-figure-height 400
+
+        set +ex
+    fi
+
+    if ([[ $ABLATION_NAME == *"code"* ]] || [[ $ABLATION_NAME == *"dolma"* ]]) && [[ ! -d "$ABLATION_DIR/code" ]]; then
+        # only plot if the directory doesn't exist
+        set -ex
+
+        python ${SCRIPT_DIR}/wandb_to_plot.py \
+            -t ai2-llm \
+            -p c4-small \
+            -n $RUNS_ABLATION \
+            -N $ABLATION_NAME \
+            -y $CODE_PERPLEXITY_SUITE \
+            -s $SAMPLES \
+            -d "${ABLATION_DIR}/code" \
+            --max-x-axis $X_AXIS_LIMIT \
+            --max-y-axis 100 \
+            --y-log-scale \
+            -v ${SCRIPT_DIR}/wandb_run_vocab.yaml \
+            --plotly-font-size 9 \
+            --plotly-figure-width 400 \
+            --plotly-figure-height 400
+
+        set +ex
+    fi
+done
diff --git a/scripts/dolma_stats.py b/scripts/dolma_stats.py
index d27a2dcf..c45a1f5e 100644
--- a/scripts/dolma_stats.py
+++ b/scripts/dolma_stats.py
@@ -209,11 +209,7 @@ class BaseStatsProcessor(BaseParallelProcessor):
 
     @classmethod
     def increment_progressbar(
-        cls,
-        queue: "Queue[Union[Tuple[int, ...], None]]",
-        /,
-        files: int = 0,
-        documents: int = 0
+        cls, queue: "Queue[Union[Tuple[int, ...], None]]", /, files: int = 0, documents: int = 0
     ) -> Dict[str, int]:
         return super().increment_progressbar(queue, files=files, documents=documents)
 
@@ -237,7 +233,7 @@ def _merge_dicts(cls, d1, d2):
     def _run_parallel_processor(cls, stats_root: str, num_workers: int, debug: bool, **process_single_kwargs: Any):
         with TemporaryDirectory() as tempdir:
             h = hashlib.md5()
-            for path in (cls.documents if isinstance(cls.documents, list) else [cls.documents]):
+            for path in cls.documents if isinstance(cls.documents, list) else [cls.documents]:
                 h.update(path.encode())
 
             metadata = os.path.join(tempdir, h.hexdigest())
@@ -419,7 +415,7 @@ def process_single(
 class cc_v1_c4_cleaned(BaseStatsProcessor):
     documents = "s3://ai2-llm/pretraining-data/sources/common-crawl/v1-c4-cleaned/documents/cc_en_*/*.gz"
     stats = "s3://ai2-llm/stats/olmo-mix/v1/cc/v1_c4_cleaned/**/*.gz"
-    decontamination_key: str = 'decontamination'
+    decontamination_key: str = "decontamination"
 
     @classmethod
     def gopher_rules(cls, attrs: Dict[str, List[Tuple[int, int, float]]]) -> List[Tuple[int, int, float]]:
@@ -637,7 +633,7 @@ def process_single(
 class v15_cc_c4_cleaned(cc_v1_c4_cleaned):
     documents = "s3://ai2-llm/pretraining-data/sources/common-crawl/v1-c4-cleaned/documents/cc_en_*/*.gz"
     stats = "s3://ai2-llm/stats/olmo-mix/v15/cc/v1_c4_cleaned/**/*.gz"
-    decontamination_key: str = 'perplexity_suite_v3_option2'
+    decontamination_key: str = "perplexity_suite_v3_option2"
 
 @Registry.add
 class LineStatsCC(cc_v1_c4_cleaned):
@@ -1153,13 +1149,11 @@ class decon_ppl_v3(BaseStatsProcessor):
         "s3://ai2-llm/pretraining-data/sources/wikipedia/v0/documents/lang=en/*.gz",
         "s3://ai2-llm/pretraining-data/sources/wikipedia/v0/documents/lang=simple/*.gz",
         "s3://ai2-llm/pretraining-data/sources/wikibooks/v0/documents/lang=en/*.gz",
-        "s3://ai2-llm/pretraining-data/sources/wikibooks/v0/documents/lang=simple/*.gz"
+        "s3://ai2-llm/pretraining-data/sources/wikibooks/v0/documents/lang=simple/*.gz",
     ]
     stats = "s3://ai2-llm/stats/olmo-mix/v1_5/decontamination_ppl_v3_option2"
 
 
-
-
 # # # BELOW HERE: AGGREGATION # # #
 
 
diff --git a/scripts/find_offset.py b/scripts/find_offset.py
new file mode 100644
index 00000000..eda8a96f
--- /dev/null
+++ b/scripts/find_offset.py
@@ -0,0 +1,63 @@
+"""
+Find source document given `dolma tokens` metadata and offset.
+
+@soldni
+"""
+
+
+import argparse
+import json
+from bisect import bisect
+from typing import Any, Dict, List
+
+import pandas as pd
+import smart_open
+import tqdm
+
+
+def parse_args() -> argparse.Namespace:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("-f", "--file", type=str, required=True)
+    ap.add_argument("-o", "--offsets", type=int, nargs="+", required=True)
+    ap.add_argument("-d", "--destination", type=str, required=True)
+    return ap.parse_args()
+
+
+def main() -> None:
+    opts = parse_args()
+
+    print("Loading metadata...", end=" ", flush=True)
+    data = pd.read_csv(opts.file, header=0, names=["start", "end", "sha", "path", "pos"])
+
+    locs: List[int] = []
+    for offset in opts.offsets:
+        loc = bisect(data["start"].values, offset)
+        if data["start"][loc] > offset:
+            # in case of not exact match
+            loc -= 1
+
+        if data.iloc[loc].start >= offset or data.iloc[loc].end <= offset:
+            raise ValueError(f"Offset {offset} not found in {opts.file}.")
+        locs.append(loc)
+
+    print(f"offset(s) found at {locs} row(s)", flush=True)
+
+    rows = data.iloc[locs]
+
+    documents: List[Dict[str, Any]] = []
+    for i, row in rows.iterrows():
+        with smart_open.open(row.path) as f:
+            for i, ln in tqdm.tqdm(enumerate(f, start=1), desc="Searching doc...", unit="lines", total=row.pos):
+                if i == row.pos:
+                    documents.append(json.loads(ln))
+                    break
+                elif i > row.pos:
+                    raise ValueError("location in document not found")
+
+    with smart_open.open(opts.destination, "w") as f:
+        for doc in documents:
+            f.write(json.dumps(doc) + "\n")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/hash_sample.py b/scripts/hash_sample.py
index 22024983..50e13db9 100644
--- a/scripts/hash_sample.py
+++ b/scripts/hash_sample.py
@@ -2,22 +2,21 @@
 import hashlib
 import json
 import multiprocessing
+import re
 from contextlib import ExitStack
 from itertools import product
 from queue import Queue
-import re
 from tempfile import TemporaryDirectory
-from typing import Any, Dict, Tuple, Union, List
-from dolma.core.paths import join_path
-from dolma.core.runtime import _make_paths_from_prefix
+from typing import Any, Dict, List, Tuple, Union
 
 import msgspec
-import uniseg.wordbreak
 import smart_open
+import uniseg.wordbreak
 
 from dolma.core.data_types import InputSpec
 from dolma.core.parallel import BaseParallelProcessor
-
+from dolma.core.paths import join_path
+from dolma.core.runtime import _make_paths_from_prefix
 
 # we like higher digits first for consistency with previous sampling strategies
 HEX_DIGITS = list("0123456789abcdef")[::-1]
@@ -30,12 +29,12 @@ def calculate_md5_suffix(prob: float) -> List[str]:
     n = 16 * prob
 
     if n > 1:
-        suffix = HEX_DIGITS[:int(round(n))]
+        suffix = HEX_DIGITS[: int(round(n))]
     else:
         # not enough hex positions to represent the probability
         # so we use more digits
         more_suffix = calculate_md5_suffix(prob * 16)
-        suffix = [''.join(e) for e in product(HEX_DIGITS[:1], more_suffix)]
+        suffix = ["".join(e) for e in product(HEX_DIGITS[:1], more_suffix)]
 
     assert all(len(s) < 32 for s in suffix), "Suffixes must be less than 32 characters"
 
@@ -54,11 +53,7 @@ def increment_progressbar(  # type: ignore
         extracted: int = 0,
     ) -> Dict[str, int]:
         return super().increment_progressbar(
-            queue,
-            files=files,
-            documents=documents,
-            extracted=extracted,
-            uniseg_words=words
+            queue, files=files, documents=documents, extracted=extracted, uniseg_words=words
         )
 
     @classmethod
@@ -74,7 +69,6 @@ def process_single(
 
         complement = kwargs.get("complement", False)
 
-
         suffixes = calculate_md5_suffix(probability)
         re_suffixes = re.compile(rf'({"|".join(suffixes)})$')
         extracted_count = documents_count = words_count = 0
@@ -103,10 +97,7 @@ def process_single(
                     # update the progress bar every 1000 documents to prevent
                     # buffering
                     cls.increment_progressbar(
-                        queue,
-                        documents=documents_count,
-                        extracted=extracted_count,
-                        words=words_count
+                        queue, documents=documents_count, extracted=extracted_count, words=words_count
                     )
                     extracted_count = documents_count = words_count = 0
 
@@ -115,11 +106,7 @@ def process_single(
                         update_interval *= 2
 
         cls.increment_progressbar(
-            queue,
-            files=1,
-            documents=documents_count,
-            extracted=extracted_count,
-            words=words_count
+            queue, files=1, documents=documents_count, extracted=extracted_count, words=words_count
         )
 
 
@@ -130,9 +117,8 @@ def main(
     complement: bool = False,
     num_workers: int = 1,
     debug: bool = False,
-    dryrun: bool = False
+    dryrun: bool = False,
 ) -> None:
-
     # make source always a list
     source = [source] if isinstance(source, str) else source
 
@@ -165,10 +151,10 @@ def main(
 
 def parse_args():
     ap = argparse.ArgumentParser()
-    ap.add_argument("-s", '--source', type=str, required=True, help="glob pattern for source files", nargs="+")
-    ap.add_argument("-d", '--destination', type=str, required=True, help="destination prefix")
-    ap.add_argument("-p", '--probability', type=float, required=True, help="sampling probability")
-    ap.add_argument("-n", '--num-workers', type=int, default=1, help="number of workers")
+    ap.add_argument("-s", "--source", type=str, required=True, help="glob pattern for source files", nargs="+")
+    ap.add_argument("-d", "--destination", type=str, required=True, help="destination prefix")
+    ap.add_argument("-p", "--probability", type=float, required=True, help="sampling probability")
+    ap.add_argument("-n", "--num-workers", type=int, default=1, help="number of workers")
     ap.add_argument("--debug", action="store_true", help="debug mode")
     ap.add_argument("--dryrun", action="store_true", help="dry run")
     ap.add_argument("--complement", action="store_true", help="get the complement of the sample")
diff --git a/scripts/make_wikipedia.py b/scripts/make_wikipedia.py
index 31058d31..6001d138 100644
--- a/scripts/make_wikipedia.py
+++ b/scripts/make_wikipedia.py
@@ -1,30 +1,31 @@
-'''
+"""
 Minimal example to download a Wikipedia dump, process via WikiExtractor, and transform to Dolma format.
 
 Author: Luca Soldaini (@soldni)
-'''
+"""
 
 
 import argparse
-from contextlib import ExitStack
 import datetime
-from io import TextIOWrapper
+import json
 import logging
 import multiprocessing
-from pathlib import Path
 import re
-from tempfile import TemporaryDirectory
-from typing import Any, Dict, Union, List
-import smart_open
-import requests
 import string
 import sys
-import json
-from dolma.core.parallel import BaseParallelProcessor, QueueType
+from contextlib import ExitStack
+from io import TextIOWrapper
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Any, Dict, List, Union
+
+import requests
+import smart_open
 from uniseg.wordbreak import words as uniseg_get_words
 
+from dolma.core.parallel import BaseParallelProcessor, QueueType
 
-CMD_INSTALL = 'pip install git+https://github.com/santhoshtr/wikiextractor.git requests smart_open tqdm'
+CMD_INSTALL = "pip install git+https://github.com/santhoshtr/wikiextractor.git requests smart_open tqdm"
 
 try:
     from wikiextractor import WikiExtractor
@@ -33,19 +34,19 @@
     sys.exit(1)
 
 try:
-    import requests     # noqa
+    import requests  # noqa
 except ImportError:
     print(f"Please install requests with `{CMD_INSTALL}`")
     sys.exit(1)
 
 try:
-    import smart_open   # noqa
+    import smart_open  # noqa
 except ImportError:
     print(f"Please install smart_open with `{CMD_INSTALL}`")
     sys.exit(1)
 
 try:
-    import tqdm     # noqa
+    import tqdm  # noqa
 except ImportError:
     print(f"Please install tqdm with `{CMD_INSTALL}`")
     sys.exit(1)
@@ -70,13 +71,13 @@ def download_file(url, filename):
     if response.status_code != 200:
         raise ValueError(f"File {url} not found; try a more recent date")
 
-    file_size = int(response.headers['Content-Length'])
+    file_size = int(response.headers["Content-Length"])
 
-    with open(filename, 'wb') as f, tqdm.tqdm(
-        desc="Downloading", total=file_size, unit='B', unit_scale=True, unit_divisor=1024
+    with open(filename, "wb") as f, tqdm.tqdm(
+        desc="Downloading", total=file_size, unit="B", unit_scale=True, unit_divisor=1024
     ) as pbar:
         for chunk in response.iter_content(chunk_size=1024 * 1024):
-            if chunk:   # filter out keep-alive new chunks
+            if chunk:  # filter out keep-alive new chunks
                 f.write(chunk)
                 pbar.update(len(chunk))
 
@@ -128,7 +129,7 @@ def __init__(self, nextFile, max_file_size=0, compress=True):
         self.compress = compress
         self.max_file_size = max_file_size
         self.stack = ExitStack()
-        self.ext = '.gz' if compress else ''
+        self.ext = ".gz" if compress else ""
         self.written = 0
         self.open()
 
@@ -149,22 +150,17 @@ def close(self):
         self.stack.pop_all().close()
 
     def open(self):
-        self.file = self.stack.enter_context(smart_open.open(self.nextFile.next() + self.ext, 'wt'))
+        self.file = self.stack.enter_context(smart_open.open(self.nextFile.next() + self.ext, "wt"))
 
 
-def wiki_extract(
-    output_gzip: Path,
-    output_json: Path,
-    processes: int = 1,
-    overwrite: bool = False
-):
+def wiki_extract(output_gzip: Path, output_json: Path, processes: int = 1, overwrite: bool = False):
     if output_json.exists() and not overwrite:
         LOGGER.warning(f"Output file {output_json} exists. Use --overwrite to overwrite it.")
         return
 
     WikiExtractor.Extractor.keepLinks = False
     WikiExtractor.Extractor.HtmlFormatting = False
-    WikiExtractor.ignoreTag('a')
+    WikiExtractor.ignoreTag("a")
     WikiExtractor.Extractor.to_json = True  # type: ignore
     WikiExtractor.acceptedNamespaces = ["Article"]
     # input_file = output_gzip
@@ -187,7 +183,7 @@ def wiki_extract(
         file_size=file_size,
         file_compress=True,
         process_count=processes,
-        html_safe=False
+        html_safe=False,
     )
 
 
@@ -216,7 +212,7 @@ def process_single(cls, source_path: str, destination_path: str, queue: QueueTyp
         documents_count = 0
         update_interval = 1
 
-        with smart_open.open(source_path) as f, smart_open.open(destination_path, 'w') as g:
+        with smart_open.open(source_path) as f, smart_open.open(destination_path, "w") as g:
             try:
                 for i, line in enumerate(f):
                     try:
@@ -268,9 +264,9 @@ def process_single(cls, source_path: str, destination_path: str, queue: QueueTyp
 
 def main():
     args = get_arguments()
-    output_gzip = Path(args.output) / f'wiki_{args.date}_{args.lang}.xml.bz2'
-    output_json = Path(args.output) / f'wiki_{args.date}_{args.lang}'
-    output_final = Path(args.output) / 'v0/documents'
+    output_gzip = Path(args.output) / f"wiki_{args.date}_{args.lang}.xml.bz2"
+    output_json = Path(args.output) / f"wiki_{args.date}_{args.lang}"
+    output_final = Path(args.output) / "v0/documents"
 
     download_wiki(date=args.date, lang=args.lang, output_path=output_gzip, overwrite=args.overwrite)
     wiki_extract(output_gzip=output_gzip, output_json=output_json, processes=args.processes)
@@ -280,7 +276,7 @@ def main():
             source_prefix=f"{output_json}/*/*.gz",
             destination_prefix=str(output_final),
             metadata_prefix=tempdir,
-            num_processes=args.processes
+            num_processes=args.processes,
         )
         processor(date=args.date, lang=args.lang)
 
diff --git a/scripts/remove_empty_docs.py b/scripts/remove_empty_docs.py
index 3a0c4ece..3e12b786 100644
--- a/scripts/remove_empty_docs.py
+++ b/scripts/remove_empty_docs.py
@@ -1,7 +1,7 @@
 import json
+import multiprocessing
 from argparse import ArgumentParser
 from contextlib import ExitStack
-import multiprocessing
 from queue import Queue
 from tempfile import TemporaryDirectory
 from typing import Any, Tuple, Union
diff --git a/scripts/wandb_run_vocab.yaml b/scripts/wandb_run_vocab.yaml
new file mode 100644
index 00000000..666e2c8d
--- /dev/null
+++ b/scripts/wandb_run_vocab.yaml
@@ -0,0 +1,123 @@
+# Run Names
+olmo-small-pile-fixed-*: The Pile
+olmo-small-rpj-*: Red Pajama
+olmo-small-dolma-*: Dolma
+olmo-small-falcon-*: RefinedWeb
+olmo-small-3T-lower-lr-tie_*: 1B Dolma Model
+
+# Ablation for PII masking
+v1-small-pi-less-than-5-anonymize_*: PII Remove (>=5) + Mask (<5)
+v1-small-all-pi-removed_*: PII Remove All
+abl-cc-v1-small-dedup_*: Baseline
+
+# Ablation for Reddit Toxicity Filtering
+reddit-v1-ablation-base_*: No Filtering
+reddit-v1-ablation-pii-nsfw-toxic_filtered_*: PII + NSFW + Hate Filter
+reddit-v1-ablation-toxic-filtered_*: NSFW + Hate Filter
+
+# Data mixture
+olmo-mix-v1-sample_*: Equal Data Mix
+olmo-mix-v1-sample-mix2_*: "200% Wiki, 200% peS2o"
+olmo-mix-v1-sample-all-cc*: "Only Common Crawl (CC)"
+olmo-mix-v1-gopher-like_*: "20% Code, 60% CC, 200% peS2o"
+
+# Dedupe
+abl-cc-v2-small-dedup*: Paragraph Deduplication
+
+# Path to final cleanup
+v1-small-c4-cleaned_\d+: C4 Rules (EoS)
+v1-small-c4-filtered_\d+: C4 Rules (All)
+v1-small-gopher-filtered_\d+: Gopher Rules (All)
+v1-small-c4-cleaned-gopher-filtered_\d+: C4 (EoS), Gopher (All)
+v1-small-c4-cleaned-gopher-filtered-deduped_\d+: C4 (EoS), Gopher (All), Dedup
+
+cc_quality:
+  olmo-mix-v1-sample-all-cc*: C4 (EoS), Gopher (All), Dedup, PII, Toxic
+
+# Code Experiments
+stack-v2*: Dolma (Stack v2)
+stack-v4*: Dolma (Stack v5)
+stack-v5*: Dolma (Stack v5)
+c4_p85-starcoder_p15*: 85% Dolma (C4), 15% StarCoder
+c4_p85-stack_v4_p15*: 85% Dolma (C4), 15% Dolma (Stack v5)
+c4-stack-15p*: 85% Dolma (C4), 15% Dolma (Stack v2)
+
+# Hate Speech / NSFW Experiments
+v1-small-hatespeech-filtered-low*: Hate Filter (Low Threshold)
+# uses FastText classifier with a threshold of 0.0004 (drops 34.9% of tokens)
+v1-small-nsfw-filtered-low*: NSFW Filter (Low Threshold)
+# uses FastText classifier with a threshold of 0.00017 (drops 29.1% of tokens)
+v1-small-hatespeech-filtered-high*: Hate Filter (High Threshold)
+# uses FastText classifier with a threshold of 0.4 (drops 7.3% of tokens)
+v1-small-nsfw-filtered-high*: NSFW Filter (High Threshold)
+# uses FastText classifier with a threshold of 0.4 (drops 5.5% of tokens)
+
+
+# Training Metrics
+throughput/total_tokens: Total Tokens
+train/Perplexity: Perplexity (Train)
+train/CrossEntropyLoss: CE Loss (Train)
+
+
+# Downstream Metrics
+eval/downstream/hellaswag_len_norm: HellaSwag
+eval/downstream/piqa_len_norm: PIQA
+eval/downstream/arc_easy_acc: ARC-Easy
+eval/downstream/sciq_acc: SciQ
+eval/downstream/mrpc_f1: MRPC
+eval/downstream/commitment_bank_acc: CommitmentBank
+eval/downstream/copa_acc: COPA
+eval/downstream/openbook_qa_len_norm: OpenBookQA
+eval/downstream/winogrande_acc: Winogrande
+eval/downstream/rte_len_norm: RTE
+
+
+# Perplexity Eval (V3)
+eval/v3-small-ice-validation/Perplexity: ICE
+eval/v3-small-c4_en-validation/Perplexity: C4
+eval/v3-small-dolma_wiki-validation/Perplexity: Dolma (Wiki)
+eval/v3-small-m2d2_s2orc-validation/Perplexity: M2D2 (S2ORC)
+eval/v3-small-dolma_books-validation/Perplexity: Dolma (Books)
+eval/v3-small-dolma_pes2o-validation/Perplexity: Dolma (peS2o v2)
+eval/v3-small-dolma_stack-validation/Perplexity: Dolma (Stack v5)
+eval/v3-small-dolma_reddit-validation/Perplexity: Dolma (Reddit)
+eval/v3-small-wikitext_103-validation/Perplexity: WikiText-103
+eval/v3-small-dolma_common-crawl-validation/Perplexity: Dolma (Common Crawl)
+
+
+# Perplexity Eval (V2)
+eval/v2-small-gab-validation/Perplexity: GAB
+eval/v2-small-ice-validation/Perplexity: ICE
+eval/v2-small-ptb-validation/Perplexity: Penn Treebank
+eval/v2-small-pile-validation/Perplexity: The Pile
+eval/v2-small-4chan-validation/Perplexity: 4chan
+eval/v2-small-c4_en-validation/Perplexity: C4
+eval/v2-small-mc4_en-validation/Perplexity: mC4 (English)
+eval/v2-small-m2d2_wiki-validation/Perplexity: M2D2 (Wiki)
+eval/v2-small-m2d2_s2orc-validation/Perplexity: M2D2 (S2ORC)
+eval/v2-small-manosphere-validation/Perplexity: Manosphere
+eval/v2-small-twitterAEE-validation/Perplexity: TwitterAEE
+eval/v2-small-wikitext_103-validation/Perplexity: WikiText-103
+eval/v2-small-c4_100_domains-validation/Perplexity: C4 (100 Domains)
+
+
+# Perplexity Eval (Unnamed, equivalent to V2)
+eval/gab-validation/Perplexity: GAB
+eval/ice-validation/Perplexity: ICE
+eval/ptb-validation/Perplexity: Penn Treebank
+eval/pile-validation/Perplexity: The Pile
+eval/4chan-validation/Perplexity: 4chan
+eval/c4_en-validation/Perplexity: C4
+eval/mc4_en-validation/Perplexity: mC4 (English)
+eval/m2d2_wiki-validation/Perplexity: M2D2 (Wiki)
+eval/m2d2_s2orc-validation/Perplexity: M2D2 (S2ORC)
+eval/manosphere-validation/Perplexity: Manosphere
+eval/twitterAEE-validation/Perplexity: TwitterAEE
+eval/wikitext_103-validation/Perplexity: WikiText-103
+eval/c4_100_domains-validation/Perplexity: C4 (100 Domains)
+
+
+# Perplexity Eval; Code
+eval/mbpp_valid/Perplexity: MBPP
+eval/openai_humaneval_test/Perplexity: HumanEval
+eval/stack_v2_held_out/Perplexity: Dolma (Stack v2)
diff --git a/scripts/wandb_to_plot.py b/scripts/wandb_to_plot.py
new file mode 100644
index 00000000..e0248772
--- /dev/null
+++ b/scripts/wandb_to_plot.py
@@ -0,0 +1,211 @@
+import argparse
+import bisect
+import fnmatch
+import re
+from collections import defaultdict
+from functools import partial
+from pathlib import Path
+from statistics import stdev
+from typing import List, Optional, Sequence, Tuple
+
+from necessary import necessary
+
+with necessary(
+    modules=["plotly", "wandb", "kaleido", "yaml"],
+    message="Please run `pip install plotly wandb kaleido pyyaml`",
+):
+    import plotly.graph_objs as go
+    import plotly.io as pio
+    import yaml
+
+    import wandb
+
+    pio.kaleido.scope.mathjax = None
+
+
+def parse_args():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("-t", "--wandb-team", type=str, required=True, help="Name of the W&B team")
+    ap.add_argument("-p", "--wandb-project", type=str, required=True, help="Name of the W&B project")
+    ap.add_argument("-n", "--wandb-names", type=str, nargs="+", required=True, help="Run name or regex")
+    ap.add_argument("-x", "--x-axis", type=str, default="throughput/total_tokens", help="X axis")
+    ap.add_argument("-y", "--y-axis", nargs="+", type=str, default=["train/Perplexity"], help="Y axis")
+    ap.add_argument("-s", "--samples", default=0, type=int, help="Samples to get; 0 for no sampling")
+    ap.add_argument("-d", "--destination", type=Path, required=True, help="Directory to save plots to")
+    ap.add_argument("-e", "--plot-extension", type=str, default="pdf", help="Plot extension to use")
+    ap.add_argument("--max-x-axis", type=float, default=None, help="Maximum x value")
+    ap.add_argument("--max-y-axis", type=float, default=None, help="Maximum y value")
+    ap.add_argument("--remove-outliers", action="store_true", help="Remove outliers")
+    ap.add_argument("--legend-title", type=str, default=None, help="Legend title")
+    ap.add_argument("--y-log-scale", action="store_true", help="Use log scale for y axis")
+    ap.add_argument("-v", "--vocabulary", type=Path, default=None, help="Path json file with pretty names")
+    ap.add_argument("-N", "--experiment-nickname", type=str, default=None, help="Experiment nickname")
+    ap.add_argument("--plotly-theme", type=str, default="none", help="Plotly theme to use")
+    ap.add_argument("--plotly-font-size", type=int, default=10, help="Plotly font size")
+    ap.add_argument("--plotly-show-title", action="store_true", help="Show plot title")
+    ap.add_argument("--plotly-figure-width", type=int, default=800, help="Plotly figure width")
+    ap.add_argument("--plotly-figure-height", type=int, default=500, help="Plotly figure height")
+    return ap.parse_args()
+
+
+def match_run_name(name: str, run_names: List[str]) -> Optional[str]:
+    for run_name in run_names:
+        if fnmatch.filter([name], run_name) or re.search(run_name, name):
+            return run_name
+    return None
+
+
+def remove_outliers(
+    x: Sequence[float], y: Sequence[float], z: float = 10.0
+) -> Tuple[Sequence[float], Sequence[float]]:
+    if min(y) >= 0 and max(y) <= 1:
+        # do not crop values that are in [0, 1]
+        return x, y
+
+    std = stdev(y)
+    x, y = zip(*[(x, y) for x, y in zip(x, y) if abs(y - std) < z * std])
+    return x, y
+
+
+def main():
+    opts = parse_args()
+
+    # make sure we're logged in
+    wandb.login()
+
+    # get all the runs matching name filters
+    api = wandb.Api()
+    wb_path = f"{opts.wandb_team}/{opts.wandb_project}"
+    # wb_filters = {"$or": [{"config.run_name": {"$regex": n}} for n in opts.wandb_names]}
+    wb_filters = {"$or": [{"display_name": {"$regex": n}} for n in opts.wandb_names]}
+    wb_runs = api.runs(path=wb_path, filters=wb_filters)
+
+    vocabulary = {}
+    if opts.vocabulary is not None and opts.vocabulary.exists():
+        with opts.vocabulary.open("r") as f:
+            vocabulary = yaml.safe_load(f.read())
+
+    # prepare the vocabulary by (a) overriding with experiment specific
+    # names that are not in the vocabulary file and (b) experiment config
+    # (nested vocabulary)
+    vocabulary = {
+        **{k: v for k, v in vocabulary.items() if isinstance(v, str)},
+        **{k: v for k, v in vocabulary.get(opts.experiment_nickname, {}).items() if isinstance(v, str)},
+    }
+
+    metrics = defaultdict(lambda: {n: {"x": [], "y": []} for n in opts.wandb_names})
+    run_name_matcher = partial(match_run_name, run_names=opts.wandb_names)
+
+    print(f"Found {len(wb_runs)} matching runs in {wb_path}")
+
+    for wb_run in wb_runs:
+        plot_group_name = run_name_matcher(wb_run.name)
+        if plot_group_name is None:
+            raise ValueError(f"Could not find a name match for {wb_run.name}")
+
+        print(f"Processing run {wb_run.name} into group {plot_group_name}")
+
+        if opts.samples > 0:
+            x_axis_history = wb_run.history(samples=opts.samples, keys=[opts.x_axis], pandas=False)
+        else:
+            x_axis_history = list(wb_run.scan_history(keys=[opts.x_axis]))
+
+        if len(x_axis_history) == 0:
+            # this run has crashed and has no history
+            print(f"WARNING: skipping {wb_run.name} because it has no history. Crashed early?")
+            continue
+
+        steps, x_axis = zip(*sorted([(wb_step["_step"], wb_step[opts.x_axis]) for wb_step in x_axis_history]))
+
+        if opts.samples > 0:
+            history = wb_run.history(samples=opts.samples, keys=opts.y_axis, pandas=False)
+        else:
+            history = wb_run.scan_history(keys=opts.y_axis)
+
+        for y_axis in opts.y_axis:
+            yaxis_pretty_name = vocabulary.get(y_axis, y_axis)
+            for wb_step in history:
+                loc = min(bisect.bisect_left(steps, wb_step["_step"]), len(x_axis) - 1)
+                metrics[yaxis_pretty_name][plot_group_name]["x"].append(x_axis[loc])
+                metrics[yaxis_pretty_name][plot_group_name]["y"].append(wb_step[y_axis])
+
+    xaxis_pretty_name = vocabulary.get(opts.x_axis, opts.x_axis)
+
+    for y_axis, plot_groups in metrics.items():
+        fig = go.Figure()
+
+        # these we figure out as we go
+        use_y_log = opts.y_log_scale
+        top_right_legend = False
+
+        for run_name, run_data in plot_groups.items():
+            if len(run_data["y"]) == 0:
+                print(f"WARNING: skipping {run_name} because it has no data for {y_axis}")
+                continue
+
+            # start by sorting the data by x axis
+            x, y = zip(*sorted(zip(run_data["x"], run_data["y"])))
+
+            if opts.remove_outliers:
+                x, y = remove_outliers(x=x, y=y)
+
+            if opts.max_x_axis and min(x) < opts.max_x_axis:
+                x, y = zip(*[(x, y) for x, y in zip(x, y) if x <= opts.max_x_axis])
+
+            if opts.max_y_axis and min(y) < opts.max_y_axis:
+                x, y = zip(*[(x, y) for x, y in zip(x, y) if y <= opts.max_y_axis])
+
+            if y[-1] < y[0]:
+                # decreasing y values, so we want the legend on the top right
+                top_right_legend = True
+
+            # only use log scale if we have values that are not in [0, 1]
+            if max(y) > 1 or min(y) < 0:
+                min_y = min([y for y in y if y > 0] or [1e-3])  # avoid diving by zero
+                use_y_log = use_y_log or (max(y) / min_y > 100)
+
+            figure_run_name = vocabulary.get(run_name, run_name)
+            fig.add_trace(go.Scatter(name=figure_run_name, x=x, y=y, mode="lines"))
+
+        legend_config = {
+            "yanchor": "top" if top_right_legend else "bottom",
+            "y": 0.99 if top_right_legend else 0.01,
+            "xanchor": "right",
+            "x": 0.99,
+            "font": {"size": opts.plotly_font_size},
+            "borderwidth": 1,
+            "bordercolor": "Gray",
+        }
+        fig.update_layout(legend=legend_config)
+
+        title_text = vocabulary.get(opts.experiment_nickname, opts.experiment_nickname)
+        fig.update_layout(
+            template=opts.plotly_theme,
+            xaxis_title=xaxis_pretty_name,
+            yaxis_title=y_axis,
+            legend_title=opts.legend_title,
+            title_text=title_text if opts.plotly_show_title else None,
+            font=dict(size=opts.plotly_font_size),
+            width=opts.plotly_figure_width,
+            height=opts.plotly_figure_height,
+            margin=dict(
+                l=4 * opts.plotly_font_size,
+                r=opts.plotly_font_size,
+                b=4 * opts.plotly_font_size,
+                t=opts.plotly_font_size,
+            ),
+        )
+        fig.update_yaxes(type="log" if use_y_log else "linear")
+        fig.update_xaxes(range=([0, opts.max_x_axis] if opts.max_x_axis is not None else None))
+
+        file_name = re.sub(r"\W+", "_", y_axis).lower().strip("_")
+        file_path = opts.destination / f"{file_name}.{opts.plot_extension}"
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+        fig.write_image(str(file_path))
+        print(f"Saved {y_axis} plot to {file_path}")
+
+
+NORMALIZE_PLOTLY_NAME = re.compile(r"([a-z])([A-Z])")
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/wimbd_to_dolma.py b/scripts/wimbd_to_dolma.py
index 22473a04..6bd0d553 100644
--- a/scripts/wimbd_to_dolma.py
+++ b/scripts/wimbd_to_dolma.py
@@ -1,8 +1,8 @@
-'''
+"""
 Convert a dataset that is in the format used by the WIMBD project to the format used in Dolma.
 
 Author: Luca Soldaini (@soldni)
-'''
+"""
 
 import argparse
 import datetime
diff --git a/src/s3_util.rs b/src/s3_util.rs
index 75ba1036..6f97b9d1 100644
--- a/src/s3_util.rs
+++ b/src/s3_util.rs
@@ -39,7 +39,8 @@ pub async fn download_to_file(
     let result = s3_client
         .get_object()
         .bucket(bucket)
-        .key(key.clone())
+        // the type `str` does not implement `Clone`, so calling `clone` on `&str` copies the reference, which does not do anything and can be removed
+        .key(key)
         .send()
         .await
         .map_err(|e| {
@@ -70,7 +71,8 @@ pub async fn upload_file(
     s3_client
         .put_object()
         .bucket(bucket)
-        .key(key.clone())
+        // note: the type `str` does not implement `Clone`, so calling `clone` on `&str` copies the reference, which does not do anything and can be removed
+        .key(key)
         .body(ByteStream::from_path(path).await?)
         .send()
         .await


Test patch:
diff --git a/tests/python/test_code.py b/tests/python/test_code.py
new file mode 100644
index 00000000..6af2738f
--- /dev/null
+++ b/tests/python/test_code.py
@@ -0,0 +1,238 @@
+"""
+
+Unit tests for code taggers.
+
+@soldni
+
+"""
+
+import re
+import unittest
+
+from bs4 import BeautifulSoup
+
+from dolma.core.data_types import Document, DocumentWithMetadata
+from dolma.taggers.code import (
+    CodeCopyrightTagger,
+    CodeRedPajamaTaggers,
+    CodeSecretsTagger,
+    CodeStarCoderTaggers2,
+)
+
+DOC_WITH_SECRETS_AND_COPYRIGHT = """
+/* copyright: Test 2023 **/
+
+This is a document.
+
+This line contains a secret: https://username:password@dolma.allen.ai
+
+This is a line with just text.
+"""
+
+
+class TestCodeTaggers(unittest.TestCase):
+    def setUp(self) -> None:
+        self.doc = Document(id="0", text=DOC_WITH_SECRETS_AND_COPYRIGHT.strip(), source=__file__)
+        return super().setUp()
+
+    def test_code_secrets_tagger(self) -> None:
+        tagger = CodeSecretsTagger()
+        result = tagger.predict(self.doc)
+
+        self.assertEqual(len(result.spans), 3)
+
+        self.assertEqual(result.spans[0].type, "SECRET_Secret_Keyword")
+        self.assertEqual(result.spans[0].select(self.doc), "https://username:password@dolma.allen.ai")
+
+        self.assertEqual(result.spans[1].type, "SECRET_Basic_Auth_Credentials")
+        self.assertEqual(result.spans[1].select(self.doc), "password")
+
+        self.assertEqual(result.spans[2].type, "doc")
+        self.assertEqual(result.spans[2].select(self.doc), self.doc.text)
+
+    def test_copyright_notice(self):
+        tagger = CodeCopyrightTagger()
+        result = tagger.predict(self.doc)
+
+        self.assertEqual(len(result.spans), 2)
+
+        self.assertEqual(result.spans[0].type, "copyright_notice")
+        self.assertEqual(result.spans[0].select(self.doc), "/* copyright: Test 2023 **/")
+
+        self.assertEqual(result.spans[1].type, "doc")
+        self.assertEqual(result.spans[1].select(self.doc), self.doc.text)
+
+
+class TestRedPajamaTaggers(unittest.TestCase):
+    def setUp(self) -> None:
+        self.doc = Document(id="0", text=DOC_WITH_SECRETS_AND_COPYRIGHT.strip(), source=__file__)
+        self.whitespace_regex = re.compile(r"\w+|[^\w\s]+")
+        return super().setUp()
+
+    def test_code_red_pajama_taggers(self) -> None:
+        tagger = CodeRedPajamaTaggers()
+        result = tagger.predict(self.doc)
+
+        # handy precomputed values
+        line_lengths = list(map(len, self.doc.text.splitlines()))
+        words = self.whitespace_regex.findall(self.doc.text)
+        self.assertGreater(len(line_lengths), 0)
+        self.assertGreater(len(words), 0)
+
+        self.assertEqual(len(result.spans), 4)
+        self.assertEqual(result.spans[0].type, "max_line_length_doc")
+        self.assertEqual(result.spans[0].score, max(line_lengths))
+
+        self.assertEqual(result.spans[1].type, "avg_line_length_doc")
+        self.assertEqual(result.spans[1].score, sum(line_lengths) / len(line_lengths))
+
+        self.assertEqual(result.spans[2].type, "alnum_prop_doc")
+        self.assertEqual(result.spans[2].score, len(list(filter(str.isalnum, self.doc.text))) / len(self.doc.text))
+
+        # TODO: This test fail; check with Akshita if this is expected
+        # self.assertEqual(result.spans[3].type, "alpha_token_prop_doc")
+        # self.assertEqual(result.spans[3].score, len(list(filter(str.isalpha, words))) / len(words))
+
+
+DOC_WITH_METADATA = """
+An XML file begins as follows:
+
+```
+<?xml version="1.0" encoding="UTF-8"?>
+...
+</xml>
+```
+
+An HTML file begins as follows:
+
+```
+<!DOCTYPE html>
+<html>
+...
+</html>
+```
+
+These are different.
+"""
+
+DOC_WITH_PYTHON_CONTENT = """
+def foo():
+    # prints hello world
+    print("Hello, World!")
+"""
+
+DOC_WITH_HTML_CONTENT = """
+<!DOCTYPE html>
+<html>
+<head>
+    <title>Page Title</title>
+</head>
+<body>
+    <h1>This is a Heading</h1>
+    <p>This is a paragraph.</p>
+</body>
+<javascript>
+    console.log("Hello, World!")
+    for (let i = 0; i < 10; i++) {
+        console.log(i)
+    }
+</javascript>
+</html>
+"""
+
+
+class TestStarCoderTaggers(unittest.TestCase):
+    def setUp(self) -> None:
+        self.md_doc = DocumentWithMetadata(
+            id="0",
+            text=DOC_WITH_METADATA.strip(),
+            source=__file__,
+            metadata={"ext": "md", "max_stars_count": 10},
+        )
+        self.python_doc = DocumentWithMetadata(
+            id="1",
+            text=DOC_WITH_PYTHON_CONTENT.strip(),
+            source=__file__,
+            metadata={"ext": "py", "max_stars_count": 1},
+        )
+        self.html_doc = DocumentWithMetadata(
+            id="2",
+            text=DOC_WITH_HTML_CONTENT.strip(),
+            source=__file__,
+            metadata={"ext": "html", "max_stars_count": 5},
+        )
+        self.tagger = CodeStarCoderTaggers2()
+        return super().setUp()
+
+    def test_metadata_tagger(self):
+        result = self.tagger.predict(self.md_doc)
+        self.assertEqual(len(result.spans), 4)
+
+        self.assertEqual(result.spans[0].type, "has_xml_template_doc")
+        self.assertEqual(result.spans[0].score, 1.0)
+
+        self.assertEqual(result.spans[1].type, "num_github_stars_doc")
+        self.assertEqual(result.spans[1].score, 10.0)
+
+        # not a python, js, or java, so this is pinned to 0.5
+        self.assertEqual(result.spans[2].type, "code_to_comment_ratio_doc")
+        self.assertEqual(result.spans[2].score, 0.5)
+
+        # not html, so this is pinned to 1.0
+        self.assertEqual(result.spans[3].type, "code_to_text_ratio_html_doc")
+        self.assertEqual(result.spans[3].score, 1.0)
+
+    def test_python_tagger(self):
+        result = self.tagger.predict(self.python_doc)
+
+        comment_lines = [
+            lns.split("#")[1].strip()
+            for ln in self.python_doc.text.split("\n")
+            if (lns := ln.strip()).startswith("#")
+        ]
+
+        self.assertEqual(len(result.spans), 4)
+
+        self.assertEqual(result.spans[0].type, "has_xml_template_doc")
+        self.assertEqual(result.spans[0].score, 0.0)
+
+        self.assertEqual(result.spans[1].type, "num_github_stars_doc")
+        self.assertEqual(result.spans[1].score, 1.0)
+
+        self.assertEqual(result.spans[2].type, "code_to_comment_ratio_doc")
+        self.assertEqual(result.spans[2].score, sum(map(len, comment_lines)) / len(self.python_doc.text))
+
+        self.assertEqual(result.spans[3].type, "code_to_text_ratio_html_doc")
+        self.assertEqual(result.spans[3].score, 1.0)
+
+    def test_html_tagger(self):
+        result = self.tagger.predict(self.html_doc)
+
+        soup = BeautifulSoup(self.html_doc.text, features="html.parser")
+
+        self.assertEqual(len(result.spans), 4)
+
+        self.assertEqual(result.spans[0].type, "has_xml_template_doc")
+        self.assertEqual(result.spans[0].score, 0.0)
+
+        self.assertEqual(result.spans[1].type, "num_github_stars_doc")
+        self.assertEqual(result.spans[1].score, 5.0)
+
+        self.assertEqual(result.spans[2].type, "code_to_comment_ratio_doc")
+        self.assertEqual(result.spans[2].score, 0.5)
+
+        self.assertEqual(result.spans[3].type, "code_to_text_ratio_html_doc")
+        self.assertEqual(result.spans[3].score, len(soup.get_text()) / len(self.html_doc.text))
+
+    def test_html_tagger_doc_too_short(self):
+        doc = DocumentWithMetadata(
+            id="3",
+            text="<html><head></head><body></body></html>",
+            source=__file__,
+            metadata={"ext": "html", "max_stars_count": 5},
+        )
+        doc.text = doc.text[:100]
+        result = self.tagger.predict(doc)
+
+        self.assertEqual(result.spans[3].type, "code_to_text_ratio_html_doc")
+        self.assertEqual(result.spans[3].score, 0.0)
diff --git a/tests/python/test_length.py b/tests/python/test_length.py
new file mode 100644
index 00000000..381d83fd
--- /dev/null
+++ b/tests/python/test_length.py
@@ -0,0 +1,51 @@
+import unittest
+
+from dolma.core.data_types import Document
+from dolma.taggers.length import CharLengthStripWsV1, CharLengthV1
+
+
+class TestLengthTaggers(unittest.TestCase):
+    def setUp(self) -> None:
+        self.doc_with_ws = Document(text="\n\tThis is a test.  ", id="0", source=__file__)
+        self.doc_no_ws = Document(text="This is a test.", id="1", source=__file__)
+        self.doc_has_inter_ws = Document(text="This is\n a test.", id="2", source=__file__)
+        self.empty_doc = Document(text="", id="3", source=__file__)
+        self.empty_doc_with_ws = Document(text="\n\t  \n", id="4", source=__file__)
+
+        self.char_length_v1 = CharLengthV1()
+        self.char_length_strip_ws_v1 = CharLengthStripWsV1()
+
+    def test_doc_with_no_ws(self):
+        result = self.char_length_v1.predict(self.doc_no_ws)
+        self.assertEqual(result.spans[0].score, 15)
+
+        result = self.char_length_strip_ws_v1.predict(self.doc_no_ws)
+        self.assertEqual(result.spans[0].score, 15)
+
+    def test_doc_with_ws(self):
+        result = self.char_length_v1.predict(self.doc_with_ws)
+        self.assertEqual(result.spans[0].score, 19)
+
+        result = self.char_length_strip_ws_v1.predict(self.doc_with_ws)
+        self.assertEqual(result.spans[0].score, 15)
+
+    def test_doc_with_inter_ws(self):
+        result = self.char_length_v1.predict(self.doc_has_inter_ws)
+        self.assertEqual(result.spans[0].score, 16)
+
+        result = self.char_length_strip_ws_v1.predict(self.doc_has_inter_ws)
+        self.assertEqual(result.spans[0].score, 16)
+
+    def test_empty_doc(self):
+        result = self.char_length_v1.predict(self.empty_doc)
+        self.assertEqual(result.spans[0].score, 0)
+
+        result = self.char_length_strip_ws_v1.predict(self.empty_doc)
+        self.assertEqual(result.spans[0].score, 0)
+
+    def test_empty_doc_with_ws(self):
+        result = self.char_length_v1.predict(self.empty_doc_with_ws)
+        self.assertEqual(result.spans[0].score, 5)
+
+        result = self.char_length_strip_ws_v1.predict(self.empty_doc_with_ws)
+        self.assertEqual(result.spans[0].score, 0)
diff --git a/tests/python/test_repetitions.py b/tests/python/test_repetitions.py
new file mode 100644
index 00000000..03077a1f
--- /dev/null
+++ b/tests/python/test_repetitions.py
@@ -0,0 +1,142 @@
+import unittest
+
+from dolma.core.data_types import Document
+from dolma.taggers.repetitions import (
+    ParagraphRepetitionsTagger,
+    RepetitionsTagger,
+    TokenizerRepetitionsTagger,
+)
+
+DOCUMENT_WITH_REPETITIONS = """
+This is a text with repetitions.
+This is a text with repetitions repetitions.
+This is a text with repetitions repetitions repetitions.
+This is a text with repetitions repetitions repetitions repetitions.
+
+blah blah blah blah blah
+
+No reps at the beginning of this sentence but MMMMMMMMMM
+
+Seeing doubles: bass banana bass banana bass banana bass banana
+"""
+
+D0M0 = " repetitions repetitions repetitions repetitions"
+D0M1 = "blah blah blah blah blah"
+D0M2 = "MMMMMMMMMM"
+D0M3 = " bass banana bass banana bass banana bass banana"
+
+DOCUMENT_WITHOUT_REPETITIONS = """
+This is a text without repetitions.
+
+Absolutely no repetitions here.
+
+There are double letters here and there but those don't count.
+"""
+
+
+class TestRepetitionsTagger(unittest.TestCase):
+    def setUp(self) -> None:
+        self.doc_with_reps = Document(source=__file__, id="0", text=DOCUMENT_WITH_REPETITIONS)
+        self.doc_without_reps = Document(source=__file__, id="1", text=DOCUMENT_WITHOUT_REPETITIONS)
+
+        self.repetitions_tagger = RepetitionsTagger()
+        self.para_reps_tagger = ParagraphRepetitionsTagger()
+
+        return super().setUp()
+
+    def test_doc_with_repetitions(self):
+        all_result = self.repetitions_tagger.predict(self.doc_with_reps)
+        par_result = self.para_reps_tagger.predict(self.doc_with_reps)
+        self.assertEqual(len(all_result.spans), 7)
+        self.assertEqual(len(par_result.spans), 7)
+
+        self.assertEqual(all_result.spans[0].type, "repetition")
+        self.assertEqual(all_result.spans[0].select(self.doc_with_reps), D0M0)
+        self.assertEqual(all_result.spans[0].score, len(D0M0))
+        self.assertEqual(all_result.spans[0], par_result.spans[0])
+
+        self.assertEqual(all_result.spans[1].type, "repetition")
+        self.assertEqual(all_result.spans[1].select(self.doc_with_reps), D0M1)
+        self.assertEqual(all_result.spans[1].score, len(D0M1))
+        self.assertEqual(all_result.spans[1], par_result.spans[1])
+
+        self.assertEqual(all_result.spans[2].type, "repetition")
+        self.assertEqual(all_result.spans[2].select(self.doc_with_reps), D0M2)
+        self.assertEqual(all_result.spans[2].score, len(D0M2))
+        self.assertEqual(all_result.spans[2], par_result.spans[2])
+
+        self.assertEqual(all_result.spans[3].type, "repetition")
+        self.assertEqual(all_result.spans[3].select(self.doc_with_reps), D0M3)
+        self.assertEqual(all_result.spans[3].score, len(D0M3))
+        self.assertEqual(all_result.spans[3], par_result.spans[3])
+
+        self.assertEqual(all_result.spans[4].type, "doc_max_repetition")
+        self.assertEqual(all_result.spans[4].score, len(D0M0))
+        self.assertEqual(all_result.spans[4], par_result.spans[4])
+
+        matches_length = len(D0M0) + len(D0M1) + len(D0M2) + len(D0M3)
+        self.assertEqual(all_result.spans[5].type, "doc_mean_repetition")
+        self.assertEqual(all_result.spans[5].score, matches_length / 4)
+        self.assertEqual(all_result.spans[5], par_result.spans[5])
+
+        self.assertEqual(all_result.spans[6].type, "doc_frac_repetition")
+        self.assertEqual(all_result.spans[6].score, matches_length / len(self.doc_with_reps.text))
+        self.assertEqual(all_result.spans[6], par_result.spans[6])
+
+    def test_doc_without_repetitions(self):
+        all_result = self.repetitions_tagger.predict(self.doc_without_reps)
+        par_result = self.para_reps_tagger.predict(self.doc_without_reps)
+        self.assertEqual(len(all_result.spans), 3)
+        self.assertEqual(len(par_result.spans), 3)
+
+        self.assertEqual(all_result.spans[0].type, "doc_max_repetition")
+        self.assertEqual(all_result.spans[0].score, 0)
+        self.assertEqual(all_result.spans[0], par_result.spans[0])
+
+        self.assertEqual(all_result.spans[1].type, "doc_mean_repetition")
+        self.assertEqual(all_result.spans[1].score, 0)
+        self.assertEqual(all_result.spans[1], par_result.spans[1])
+
+        self.assertEqual(all_result.spans[2].type, "doc_frac_repetition")
+        self.assertEqual(all_result.spans[2].score, 0)
+        self.assertEqual(all_result.spans[2], par_result.spans[2])
+
+
+class TestTokenizerRepetitionsTagger(unittest.TestCase):
+    def setUp(self) -> None:
+        self.doc_with_reps = Document(source=__file__, id="0", text=DOCUMENT_WITH_REPETITIONS)
+        self.doc_without_reps = Document(source=__file__, id="1", text=DOCUMENT_WITHOUT_REPETITIONS)
+        self.repetitions_tagger = TokenizerRepetitionsTagger()
+
+        return super().setUp()
+
+    def test_doc_with_repetitions(self):
+        repeated_strings = [
+            "repetitions repetitions repetitions",
+            "repetitions repetitions repetitions repetitions",
+            "blah blah blah blah",  # missing a blah bc the first element in this seq has diff token id
+            "MMMMMMMM",  # shorter bc sequence is tokenized as 'ĠM', 'MM', 'MM', 'MM', 'MM', 'M'
+            "bass banana bass banana bass banana bass banana",
+        ]
+
+        all_results = self.repetitions_tagger.predict(self.doc_with_reps)
+        self.assertEqual(len(all_results.spans), len(repeated_strings) + 3)
+
+        i = 0
+        for string in repeated_strings:
+            self.assertEqual(all_results.spans[i].type, "repetition")
+            self.assertEqual(string, all_results.spans[i].select(self.doc_with_reps))
+            self.assertEqual(len(string), all_results.spans[i].score)
+            i += 1
+
+        self.assertEqual(all_results.spans[i].type, "doc_max_repetition")
+        self.assertEqual(all_results.spans[i].score, max(len(s) for s in repeated_strings))
+        i += 1
+
+        matches_length = sum(map(len, repeated_strings))
+        self.assertEqual(all_results.spans[i].type, "doc_mean_repetition")
+        self.assertEqual(all_results.spans[i].score, matches_length / len(repeated_strings))
+        i += 1
+
+        self.assertEqual(all_results.spans[i].type, "doc_frac_repetition")
+        self.assertEqual(all_results.spans[i].score, matches_length / len(self.doc_with_reps.text))
diff --git a/tests/python/test_repetitions_utils.py b/tests/python/test_repetitions_utils.py
new file mode 100644
index 00000000..81fe811a
--- /dev/null
+++ b/tests/python/test_repetitions_utils.py
@@ -0,0 +1,126 @@
+from unittest import TestCase
+
+import numpy as np
+
+from dolma.taggers.repetitions.utils import (
+    find_end_first_consecutive_true,
+    find_periodic_sequences,
+    find_start_last_consecutive_true,
+    group_consecutive_values,
+)
+
+
+class TestTrueLocsDetection(TestCase):
+    def test_find_end_first_consecutive_true(self):
+        arr = np.array([True, True, False, True])
+        self.assertEqual(find_end_first_consecutive_true(arr), 2)
+
+        arr = np.array([False, True])
+        self.assertEqual(find_end_first_consecutive_true(arr), 0)
+
+        arr = np.array([True, False])
+        self.assertEqual(find_end_first_consecutive_true(arr), 1)
+
+        arr = np.array([True])
+        self.assertEqual(find_end_first_consecutive_true(arr), 1)
+
+        arr = np.array([False])
+        self.assertEqual(find_end_first_consecutive_true(arr), 0)
+
+    def test_find_start_last_consecutive_true(self):
+        arr = np.array([False, False, True, True])
+        self.assertEqual(find_start_last_consecutive_true(arr), 2)
+
+        arr = np.array([True, False, False, True])
+        self.assertEqual(find_start_last_consecutive_true(arr), 3)
+
+        arr = np.array([True, True, True, True])
+        self.assertEqual(find_start_last_consecutive_true(arr), 0)
+
+        arr = np.array([False, False, False, False])
+        self.assertEqual(find_start_last_consecutive_true(arr), -1)
+
+        arr = np.array([True, True, True, False])
+        self.assertEqual(find_start_last_consecutive_true(arr), -1)
+
+    def test_group_consecutive_values(self):
+        arr = np.array([1, 2, 3, 4, 5])
+        grouped = group_consecutive_values(arr)
+        self.assertEqual(len(grouped), 1)
+        self.assertTrue((grouped[0] == arr).all())
+
+        arr = np.array([1, 2, 3, 5, 6, 7, 9, 10, 11])
+        grouped = group_consecutive_values(arr)
+        self.assertEqual(len(grouped), 3)
+        self.assertTrue((grouped[0] == np.array([1, 2, 3])).all())
+        self.assertTrue((grouped[1] == np.array([5, 6, 7])).all())
+        self.assertTrue((grouped[2] == np.array([9, 10, 11])).all())
+
+        arr = np.array([1, 3, 5, 7, 9])
+        grouped = group_consecutive_values(arr)
+        self.assertEqual(len(grouped), 5)
+        for i in range(5):
+            self.assertEqual(len(grouped[i]), 1)
+            self.assertEqual(grouped[i][0], arr[i])
+
+    def _to_array(self, s: str) -> np.ndarray:
+        return np.array(list(map(int, s)))
+
+    def test_find_periodic_sequences(self):
+        # should be 0 repeated 5 times
+        arr = self._to_array("5000007")
+        sequences = list(find_periodic_sequences(arr, max_period=1))
+        self.assertEqual(len(sequences), 1)
+        self.assertEqual(sequences[0].start, 1)
+        self.assertEqual(sequences[0].end, 6)
+        self.assertEqual(sequences[0].period, 1)
+        self.assertEqual(sequences[0].times, 5)
+
+        # should be 01 repeated 3 times
+        arr = self._to_array("501010127")
+        sequences = list(find_periodic_sequences(arr, min_period=2, max_period=2))
+        self.assertEqual(len(sequences), 1)
+        self.assertEqual(sequences[0].start, 1)
+        self.assertEqual(sequences[0].end, 7)
+        self.assertEqual(sequences[0].period, 2)
+        self.assertEqual(sequences[0].times, 3)
+
+        # should be 0104 repeated 3 times
+        arr = self._to_array("5301040104010401")
+        sequences = list(find_periodic_sequences(arr, min_period=4, max_period=4))
+        self.assertEqual(len(sequences), 1)
+        self.assertEqual(sequences[0].start, 2)
+        self.assertEqual(sequences[0].end, 16)
+        self.assertEqual(sequences[0].period, 4)
+        self.assertEqual(sequences[0].times, 3)
+
+        # should be 040 repeated 4 times
+        arr = np.array(list(map(int, "04004004004030")))
+        sequences = list(find_periodic_sequences(arr, min_period=3, max_period=3))
+        self.assertEqual(len(sequences), 1)
+        self.assertEqual(sequences[0].start, 0)
+        self.assertEqual(sequences[0].end, 12)
+        self.assertEqual(sequences[0].period, 3)
+        self.assertEqual(sequences[0].times, 4)
+
+        # should have two repetitions: 46 repeated 4 times, and 550 repeated 3 times
+        arr = np.array(list(map(int, "004646464639955055055046550")))
+        sequences = list(find_periodic_sequences(arr, min_period=2, max_period=3))
+        self.assertEqual(len(sequences), 2)
+        self.assertEqual(sequences[0].start, 2)
+        self.assertEqual(sequences[0].end, 10)
+        self.assertEqual(sequences[0].period, 2)
+        self.assertEqual(sequences[0].times, 4)
+        self.assertEqual(sequences[1].start, 13)
+        self.assertEqual(sequences[1].end, 22)
+        self.assertEqual(sequences[1].period, 3)
+        self.assertEqual(sequences[1].times, 3)
+
+    def test_find_no_periodic_sequences(self):
+        arr = np.array(list(map(int, "123456789")))
+        sequences = list(find_periodic_sequences(arr, max_period=10))
+        self.assertEqual(len(sequences), 0)
+
+        arr = np.array(list(map(int, "112233445566778899")))
+        sequences = list(find_periodic_sequences(arr, max_period=10))
+        self.assertEqual(len(sequences), 0)