huggingface / huggingface/datatrove

Implementing DecileEstimator

Open
#336 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
3.3k
Forks
302
Avg merge
2h 18m
Merged PRs (30d)
2

Description

Hi,
I'm playing a lot with Datatrove and it's very handy to use, we can easily add a lot of Custom Step.

Recently I'm trying to create a Decile estimator to filter base on the distribution of the data.

To do that I created a `PipelineStep` which is very basic and it's working well.
My only issue is when launch with several tasks, at the end it sums the stats of the different run instead of making something more intelligent in this case.

I need to merge the different `TDigest` dictionnary, so I tried to make a `DigestStatsDict` which is working well but now the issue lays in the `MetricStats` where it's trying to compute the mean and standard deviation etc... which of course fail because it doesn't apply here

So I'm looking for suggestion on how to integrate properly this Step and I would be happy to open a MR on the subject.

Best regards,

My code :
```python
from tdigest import TDigest

class DigestStatsDict:
"""
Stores multiple stats
"""

def __init__(self, init=None, **kwargs):
self.tdigest = TDigest()
if init:
self.tdigest.update_from_dict(init)

def __add__(self, other):
self.tdigest = self.tdigest + other.tdigest
self.tdigest.compress()

def update(self, value, unit=None):
self.tdigest = self.tdigest + value
self.tdigest.compress()

def __repr__(self):
rep = ""
for i in range(10):
rep += f"decile_{i}: {self.tdigest.percentile(i * 10)}, "
rep += f"decile_95: {self.tdigest.percentile(95)}, "
rep += f"decile_99: {self.tdigest.percentile(99)}, "
rep += f"decile_99.5: {self.tdigest.percentile(99.5)}"
return rep

def to_dict(self):
return self.tdigest.to_dict()

@classmethod
def from_dict(cls, data):
return DigestStatsDict(init=data)

class DecileEstimator(PipelineStep):
name = "📒 DecileEstimator"
type = "📖 - COUNTER"

def __init__(self, batch_size: int = 10000):
super().__init__()
self.batch_size = batch_size
self.stats["digest"] = DigestStatsDict()

def set_stat(self, label, value):
self.stats[label] = value

def run(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1) -> DocumentsPipeline:
tdigest = TDigest()
for batch_doc in batched(data, self.batch_size):
tokens = [document.metadata["token_count"] for document in batch_doc]
tdigest.batch_update(tokens)
yield from batch_doc
tdigest.compress()
self.stat_update("digest", value=tdigest)

```

The error :
```bash
File "/data/home/j-g.barthelemy/the-stack/llm-data-preparation/pretraining/.venv/lib/python3.12/site-packages/datatrove/utils/stats.py", line 257, in __add__
mean = (self.n * self.mean + other.n * other.mean) / n
~~~~~~~~^~~~~~~~~~~~
TypeError: unsupported operand type(s) for *: 'int' and 'DigestStatsDict'
```

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.