3x Faster Text Preprocessing
Opened by ashvardanian (open)

I was preparing some datasets for AI training and noticed that `datasets` by HuggingFace uses the conventional `open` mechanism to read the file and split it into chunks. I thought it can be significantly accelerated, and [started with a benchmark](https://gist.github.com/ashvardanian/55c2052e9f78b05b8d614aa90cb12347):

```sh
$ pip install --upgrade --force-reinstall datasets
$ python benchmark_huggingface_datasets.py xlsum.csv 
Generating train split: 1004598 examples [00:47, 21116.16 examples/s]
Time taken to load the dataset: 48.66838526725769 seconds
Time taken to chunk the dataset into parts of size 10000: 0.11466407775878906 seconds
Total time taken: 48.78304934501648 seconds
```

For benchmarks I've used a [large CSV file with mixed UTF-8 content](https://github.com/ashvardanian/StringZilla/blob/main/CONTRIBUTING.md#benchmarking-datasets), most common in modern large-scale pre-training pipelines. I've later patched the `datasets` library to use `stringzilla`, which resulted in significantly lower memory consumption and in 2.9x  throughput improvement on the AWS `r7iz` instances. That's using slow SSDs mounted over the network. Performance on local SSDs on something like a DGX-H100 should be even higher:

```sh
$ pip install -e .
$ python benchmark_huggingface_datasets.py xlsum.csv 
Generating train split: 1004598 examples [00:15, 64529.90 examples/s]
Time taken to load the dataset: 16.45028805732727 seconds
Time taken to chunk the dataset into parts of size 10000: 0.1291060447692871 seconds
Total time taken: 16.579394102096558 seconds
```

I've already [pushed the patches to my fork](https://github.com/ashvardanian/datasets/tree/faster-text-parsers), and would love to contribute them to the upstream repository.

---

All the tests pass, but they leave a couple of important questions open. The default Python `open(..., newline=None)` uses universal newlines, where `\n`, `\r`, and `\r\n` are all converted to `\n` on the fly. I am not sure if its a good idea for a general purpose dataset preparation pipeline?

I can simulate the same behavior (which I don't yet do) for `"line"` splitter. Adjusting it for `"paragraph"`-splitter would be harder. Should we stick exactly to the old Pythonic behavior or stay closer to how C and other programming languages do that?

Comments:
Unfortunately, that won't improve the performance. StringZilla repository has extensive benchmarks comparing against different built-in functionality of several programming languages. Using `re.finditer` for tokenization is practically the slowest anti-pattern I've encountered in any language. The gap between that and a SIMD-accelerated kernel can be as big as 10 MB/s vs 10 GB/s.

I understand the need to keep the dependencies minimal. It helps the package remain small and portable. At this point, StringZilla provides 105 binaries for different OS and hardware versions (more portable than NumPy) and the [binary size generally ranges from 50 KB to 250 KB](https://pypi.org/project/stringzilla/), smaller than a single JPEG. 


The `text` builder is not very popular, so I'm also not a fan of introducing a dependency for it.

Moreover, I couldn't find any projects of this size/usage depending on StringZilla (with GitHub search), so we should at least wait for its greater adoption to merge this PR.


> Moreover, I couldn't find any projects of this size/usage depending on StringZilla (with GitHub search), so we should at least wait for its greater adoption to merge this PR.

Meanwhile I understand that you want to wait for a greater adoption - if things change in the future you would be stuck with an unsupported dependency (although I think, that really applies to everything) - the performance improvement is really significant!
I wonder if it's worth, perhaps, to provide an additional 'datasets.extras' library by huggingface which support these 3rd party improvements.
It would reduce the risk on the core components and, at the same time, it would definitely help on the performance side!

Hi @lhoestq & @mariosasko! Hope you both are doing well!

I've just pushed another year's worth of work on string algorithms, including both CPU- and GPU-accelerated processing, and I'm curious if they can be of bigger help this time 🤗 

## Tokenization

> Maybe we can improve speed and memory usage using built-in tools though, like regex `finditer`?

Revisiting this question, I've just benchmarked that exact use case, byte-set search for line breaks and whitespaces between words. The results on Intel Ice Lake+ CPUs installed in most DGX servers look like this:

| Library                     |  Shorter Words |   Longer Lines |
| --------------------------- | -------------: | -------------: |
| `re.finditer` 🐍             |     0.04 GiB/s |     0.19 GiB/s |
| `sz.Str.find_first_of` 🐍    |  __0.11 GiB/s__ | __8.79 GiB/s__ |

Moreover, on tokenization, with new Apache Arrow-compatible containers, one may easily end up with a 10x lower memory consumption, making much larger datasets feasible for in-memory processing.

## Hashing

I've also noticed that `xxhash` is used for non-cryptographic hashes. It's indeed a very good choice, but I believe it may be further improved:

| Library              | Languages         |  Shorter Words |    Longer Lines |
| -------------------- | ----------------- | -------------: | --------------: |
| `hash` 🐍             | Py                |     0.13 GiB/s |      4.27 GiB/s |
| `xxhash.xxh3_64` 🐍   | C, C++, Rs, Py... |     0.04 GiB/s |      6.38 GiB/s |
| `sz.hash` 🐍          | C, C++, Rs, Py... | __0.14 GiB/s__ |  __9.19 GiB/s__ |

---

There is a lot of other functionality for bulk-computing of MinHashes or Levenshtein distances on both CPUs and GPUs, but I'm not sure if it might be of interest to the `datasets` team. Happy to help if it is 🤗 

> All the benchmarks can be found in the [StringWa.rs repo](https://github.com/ashvardanian/StringWa.rs/).

Could you also compare with PyArrow string functions ? They offer good speed iirc and return arrow arrays directly (hence no need for an extra data copy)

Hi @lhoestq! I've been researching the PyArrow functionality over the last couple of days and not entirely sure if there are many obvious functionality overlaps. The functionality that I've already compared in [StringWa.rs, is sorting](https://github.com/ashvardanian/StringWa.rs/?tab=readme-ov-file#sequence-operations):

| Library                                     |                  Short Words |                 Long Lines |
| ------------------------------------------- | ---------------------------: | -------------------------: |
| `list.sort` on 1x CPU                       |        47.06 M comparisons/s |      22.36 M comparisons/s |
| `pandas.Series.sort_values` on 1x CPU       |         9.39 M comparisons/s |      11.93 M comparisons/s |
| `pyarrow.compute.sort_indices` on 1x CPU    |        62.17 M comparisons/s |       5.53 M comparisons/s |
| `polars.Series.sort` on 1x CPU              |       223.38 M comparisons/s | __181.60 M comparisons/s__ |
| `cudf.Series.sort_values` on 1x GPU         | __9'463.59 M comparisons/s__ |      66.44 M comparisons/s |
| `stringzilla.Strs.sorted` on 1x CPU         |       171.13 M comparisons/s |      77.88 M comparisons/s |

Please let me know if there's other functionality I should compare or add to the library 🤗 

---

PS: StringZilla's `Strs` class and PyArrow arrays are directly convertible without extra memory allocations:

```py
from pyarrow import foreign_buffer
from stringzilla import Strs

strs = Strs(["alpha", "beta", "gamma"])
arrow = foreign_buffer(strs.address, strs.nbytes, strs)
```

And only slightly harder to convert in reverse direction:

```py
arr = pa.Array.from_buffers(
    pa.large_string() if strs.offsets_are_large else pa.string(),
    len(strs),
    [None,
     pa.foreign_buffer(strs.offsets_address, strs.offsets_nbytes, strs),
     pa.foreign_buffer(strs.tape_address, strs.tape_nbytes, strs)],
)
```

Following up, here's a [success story](https://ashvardanian.com/posts/image-processing-with-strings) from integrating StringZilla into a package with 100M+ PyPI downloads lat year, replacing OpenCV for a 4x images processing speedup.

Interesting comparisons for sorting, what about the string functions that would be useful here ?