huggingface / huggingface/datatrove
Implementing DecileEstimator
- 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.