From edbea5f1362ade9ffb235f0e78ccc3c60c977b45 Mon Sep 17 00:00:00 2001
From: guipenedo <nostrumg@gmail.com>
Date: Thu, 2 Nov 2023 15:43:51 +0000
Subject: [PATCH] refactoring minhash + 64bit hash support

---
 src/datatrove/pipeline/dedup/minhash.py | 109 ++++++-----
 tests/pipeline/test_minhash.py          | 250 +++++++++++++-----------
 2 files changed, 197 insertions(+), 162 deletions(-)

diff --git a/src/datatrove/pipeline/dedup/minhash.py b/src/datatrove/pipeline/dedup/minhash.py
index 257d0a83..4b641922 100644
--- a/src/datatrove/pipeline/dedup/minhash.py
+++ b/src/datatrove/pipeline/dedup/minhash.py
@@ -18,8 +18,7 @@
 
 # http://en.wikipedia.org/wiki/Mersenne_prime
 _mersenne_prime = np.uint64((1 << 61) - 1)
-_max_hash = np.uint64((1 << 32) - 1)
-_hash_range = 1 << 32
+_max_hash_32b = np.uint64((1 << 32) - 1)
 
 """
 n_grams -> roughly nr of words (this should be small enough to catch fuzzy matches but big enough to not have each shingle be too common)
@@ -28,13 +27,31 @@
 probability of inclusion for s=0.8: 1-(1-0.8^8)^14=0.924
 """
 
-DEFAULT_NR_BUCKETS = 14
-DEFAULT_PER_BUCKET = 8
-DEFAULT_N_GRAMS = 5
-
 SENTINEL = (1 << 32) - 1
 
 
+@dataclass
+class MinhashConfig:
+    n_grams: int = 5
+
+    num_buckets: int = 14
+    hashes_per_bucket: int = 8
+
+    use_64bit_hashes: bool = False
+    seed: int = 1
+
+    @property
+    def hash_dtype(self):
+        return np.uint64 if self.use_64bit_hashes else np.uint32
+
+    @property
+    def hash_format(self):
+        return "Q" if self.use_64bit_hashes else "I"
+
+
+DEFAULT_MINHASH_CONFIG = MinhashConfig()
+
+
 @dataclass
 class HashSig:
     sig: tuple[int]
@@ -52,13 +69,13 @@ def __lt__(self, other):
         return self.to_tuple() < other.to_tuple()
 
 
-def read_sigs(file: InputDataFile, reader_id: int, hashes_per_bucket: int, index_file: bool = False) -> Generator:
-    n = hashes_per_bucket + 1 if not index_file else hashes_per_bucket
-    for data in read_tuples_from_file(file, f"{n}I"):
-        if index_file:
+def read_sigs(file: InputDataFile, reader_id: int, config: MinhashConfig, index_file: bool = False) -> Generator:
+    if index_file:
+        for data in read_tuples_from_file(file, f"{config.hashes_per_bucket}{config.hash_format}"):
             yield HashSig(sig=data, doc_id=-1, file_id=-1, reader_id=reader_id)
-        else:
-            yield HashSig(sig=data[:-1], doc_id=data[-1], file_id=reader_id, reader_id=reader_id)
+    else:
+        for *data, doc_id in read_tuples_from_file(file, f"{config.hashes_per_bucket}{config.hash_format}", "I"):
+            yield HashSig(sig=data, doc_id=doc_id, file_id=reader_id, reader_id=reader_id)
 
 
 class MinhashDedupSignature(PipelineStep):
@@ -68,19 +85,13 @@ class MinhashDedupSignature(PipelineStep):
     def __init__(
         self,
         output_folder: BaseOutputDataFolder,
-        num_buckets: int = DEFAULT_NR_BUCKETS,
-        hashes_per_bucket: int = DEFAULT_PER_BUCKET,
-        n_grams: int = DEFAULT_N_GRAMS,
-        seed: int = 1,
+        config: MinhashConfig = DEFAULT_MINHASH_CONFIG,
         **kwargs,
     ):
         super().__init__(**kwargs)
         self.output_folder = output_folder
-        self.n_grams = n_grams
-        self.num_buckets = num_buckets
-        self.hashes_per_bucket = hashes_per_bucket
-        self.num_hashes = self.num_buckets * self.hashes_per_bucket
-        self.seed = seed
+        self.config = config
+        self.num_hashes = self.config.num_buckets * self.config.hashes_per_bucket
         self._parameters = None
 
     @property
@@ -89,7 +100,7 @@ def parameters(self):
             # Create parameters for a random bijective permutation function
             # that maps a 32-bit hash value to another 32-bit hash value.
             # http://en.wikipedia.org/wiki/Universal_hashing
-            gen = np.random.RandomState(self.seed)
+            gen = np.random.RandomState(self.config.seed)
             self._parameters = gen.randint(
                 1, _mersenne_prime, dtype=np.uint64, size=(1, self.num_hashes)
             ), gen.randint(0, _mersenne_prime, dtype=np.uint64, size=(1, self.num_hashes))
@@ -97,8 +108,12 @@ def parameters(self):
 
     def get_signature(self, shingles):
         a, b = self.parameters
-        phv = np.bitwise_and((shingles * a + b) % _mersenne_prime, _max_hash)
-        return [x.tolist() for x in np.split(np.min(phv, axis=0).astype(np.uint32), self.num_buckets)]
+        phv = (shingles * a + b) % _mersenne_prime
+        if self.config.use_64bit_hashes:
+            phv = np.bitwise_and(phv, _max_hash_32b)
+        return [
+            x.tolist() for x in np.split(np.min(phv, axis=0).astype(self.config.hash_dtype), self.config.num_buckets)
+        ]
 
     def set_up_dl_locks(self, dl_lock, up_lock):
         self.output_folder.set_lock(up_lock)
@@ -107,7 +122,7 @@ def get_shingles(self, text):
         return np.array(
             [
                 [sha1_hash32(" ".join(x).encode("utf-8"))]
-                for x in ngrams(word_tokenize(simplify_content(text)), self.n_grams)
+                for x in ngrams(word_tokenize(simplify_content(text)), self.config.n_grams)
             ],
             dtype=np.uint64,
         )
@@ -115,7 +130,7 @@ def get_shingles(self, text):
     def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1):
         buckets = [
             self.output_folder.open(f"bucket_{bi:03d}/{rank:05d}.minhash.sig", mode="wb")
-            for bi in range(self.num_buckets)
+            for bi in range(self.config.num_buckets)
         ]
         for doc_idx, doc in enumerate(data):
             self.stat_update(StatHints.total)
@@ -124,14 +139,20 @@ def __call__(self, data: DocumentsPipeline, rank: int = 0, world_size: int = 1):
                 sig = self.get_signature(shingles)
                 for bi, (bucket, bucket_sig) in enumerate(zip(buckets, sig)):
                     # print(f"{self.hashes_per_bucket=} {bucket_sig=}")
-                    bucket.write(struct.pack("<%sI" % self.hashes_per_bucket, *bucket_sig))
+                    bucket.write(
+                        struct.pack(f"<{self.config.hashes_per_bucket}{self.config.hash_format}", *bucket_sig)
+                    )
                     bucket.write(struct.pack("<I", doc_idx))
         logger.info("Sorting buckets...")
         for bi, bucket in enumerate(buckets):
             bucket.close()
             fo = self.output_folder.open(bucket.relative_path, mode="r+b", overwrite=True)
-            mmap = np.memmap(fo.file_handler, dtype=[(str(i), np.uint32) for i in range(self.hashes_per_bucket + 1)])
-            mmap.sort(order=[str(i) for i in range(self.hashes_per_bucket + 1)])
+            mmap = np.memmap(
+                fo.file_handler,
+                dtype=[(str(i), self.config.hash_dtype) for i in range(self.config.hashes_per_bucket)]
+                + [(str(self.config.hashes_per_bucket), np.uint32)],
+            )  # doc_id at the end
+            mmap.sort(order=[str(i) for i in range(self.config.hashes_per_bucket + 1)])
             mmap.flush()
             fo.close()
         self.output_folder.close()
@@ -146,17 +167,15 @@ def __init__(
         input_folder: BaseInputDataFolder,
         output_folder: BaseOutputDataFolder,
         index_folder: BaseInputDataFolder = None,
-        hashes_per_bucket: int = DEFAULT_PER_BUCKET,
-        num_buckets: int = DEFAULT_NR_BUCKETS,
+        config: MinhashConfig = DEFAULT_MINHASH_CONFIG,
         only_dedup_in_index: bool = True,
         **kwargs,
     ):
         super().__init__(**kwargs)
         self.input_folder = input_folder
         self.output_folder = output_folder
-        self.num_buckets = num_buckets
-        self.hashes_per_bucket = hashes_per_bucket
         self.index_folder = index_folder
+        self.config = config
         self.only_dedup_in_index = only_dedup_in_index
 
     def set_up_dl_locks(self, dl_lock, up_lock):
@@ -165,15 +184,15 @@ def set_up_dl_locks(self, dl_lock, up_lock):
 
     def __call__(self, data: DocumentsPipeline, bucket: int = 0, world_size: int = 1):
         assert data is None, "You should not use an input block before MinhashDedupBuckets"
-        assert world_size == self.num_buckets, "You must run exactly one task per bucket"
+        assert world_size == self.config.num_buckets, "You must run exactly one task per bucket"
         sig_files = self.input_folder.list_files(suffix=f"bucket_{bucket:03d}")
-        sig_readers = [read_sigs(file, file_i, self.hashes_per_bucket) for file_i, file in enumerate(sig_files)]
+        sig_readers = [read_sigs(file, file_i, self.config) for file_i, file in enumerate(sig_files)]
         index_files = self.index_folder.list_files(suffix=f"bucket_{bucket:03d}") if self.index_folder else None
         if index_files:
             logger.info(f"Found index file(s): {', '.join([file.relative_path for file in index_files])}")
             sig_readers.extend(
                 [
-                    read_sigs(file, len(sig_readers) + file_i, self.hashes_per_bucket, index_file=True)
+                    read_sigs(file, len(sig_readers) + file_i, self.config, index_file=True)
                     for file_i, file in enumerate(index_files)
                 ]
             )
@@ -209,14 +228,14 @@ def __init__(
         self,
         input_folder: BaseInputDataFolder,
         output_folder: BaseOutputDataFolder,
-        num_buckets: int = DEFAULT_NR_BUCKETS,
+        config: MinhashConfig = DEFAULT_MINHASH_CONFIG,
         save_cluster_id: bool = False,
         **kwargs,
     ):
         super().__init__(**kwargs)
         self.input_folder = input_folder
         self.output_folder = output_folder
-        self.num_buckets = num_buckets
+        self.config = config
         self.save_cluster_id = save_cluster_id
 
     def set_up_dl_locks(self, dl_lock, up_lock):
@@ -225,7 +244,7 @@ def set_up_dl_locks(self, dl_lock, up_lock):
 
     def __call__(self, data: DocumentsPipeline, _: int = 0, world_size: int = 1):
         dup_files = self.input_folder.list_files(extension=".dups")
-        assert len(dup_files) == self.num_buckets, "There should be exactly one .dups file per bucket"
+        assert len(dup_files) == self.config.num_buckets, "There should be exactly one .dups file per bucket"
         assert world_size == 1, "World size must be 1 for clustering"
         union_set = {}
 
@@ -334,15 +353,13 @@ def __init__(
         input_folder: BaseInputDataFolder,
         output_folder: BaseOutputDataFolder,
         index_name: str,
-        hashes_per_bucket: int = DEFAULT_PER_BUCKET,
-        num_buckets: int = DEFAULT_NR_BUCKETS,
+        config: MinhashConfig = DEFAULT_MINHASH_CONFIG,
         **kwargs,
     ):
         super().__init__(**kwargs)
         self.input_folder = input_folder
         self.output_folder = output_folder
-        self.num_buckets = num_buckets
-        self.hashes_per_bucket = hashes_per_bucket
+        self.config = config
         self.index_name = index_name
 
     def set_up_dl_locks(self, dl_lock, up_lock):
@@ -351,9 +368,9 @@ def set_up_dl_locks(self, dl_lock, up_lock):
 
     def __call__(self, data: DocumentsPipeline, bucket: int = 0, world_size: int = 1):
         assert data is None, "You should not use an input block before MinhashDedupBuckets"
-        assert world_size == self.num_buckets, "You must run exactly one task per bucket"
+        assert world_size == self.config.num_buckets, "You must run exactly one task per bucket"
         sig_files = self.input_folder.list_files(suffix=f"bucket_{bucket:03d}")
-        sig_readers = [read_sigs(file, file_i, self.hashes_per_bucket) for file_i, file in enumerate(sig_files)]
+        sig_readers = [read_sigs(file, file_i, self.config) for file_i, file in enumerate(sig_files)]
 
         pq = [next(sig_reader) for sig_reader in sig_readers]
         heapq.heapify(pq)
@@ -365,7 +382,7 @@ def __call__(self, data: DocumentsPipeline, bucket: int = 0, world_size: int = 1
         while pq:
             v: HashSig = heapq.heappop(pq)
             if not last or last.sig != v.sig:
-                out_f.write(struct.pack("<%dI" % self.hashes_per_bucket, *v.sig))
+                out_f.write(struct.pack("<%dI" % self.config.hashes_per_bucket, *v.sig))
             last = v
             next_sig = next(sig_readers[v.file_id], None)
             if next_sig:
diff --git a/tests/pipeline/test_minhash.py b/tests/pipeline/test_minhash.py
index 91f2349b..ea8a5bae 100644
--- a/tests/pipeline/test_minhash.py
+++ b/tests/pipeline/test_minhash.py
@@ -6,9 +6,12 @@
 from collections import defaultdict
 from math import floor
 
+import numpy as np
+
 from datatrove.data import Document
 from datatrove.io import LocalInputDataFolder, LocalOutputDataFolder
 from datatrove.pipeline.dedup.minhash import (
+    MinhashConfig,
     MinhashDedupBuckets,
     MinhashDedupCluster,
     MinhashDedupFilter,
@@ -38,121 +41,136 @@ def tearDown(self):
         shutil.rmtree(self.test_dir)
 
     def test_signatures(self):
-        minhash = MinhashDedupSignature(
-            output_folder=LocalOutputDataFolder(os.path.join(self.test_dir, "signatures1")),
-        )
-        shingles = minhash.get_shingles(lorem_ipsum)
-        sig = minhash.get_signature(shingles)
-
-        minhash2 = MinhashDedupSignature(
-            output_folder=LocalOutputDataFolder(os.path.join(self.test_dir, "signatures2"))
-        )
-        # check consistency
-        assert sig == minhash2.get_signature(shingles)
-
-        # check correct number of outputs
-        assert len(sig) == minhash.num_buckets
-        assert all((len(x) == minhash.hashes_per_bucket for x in sig))
-
-        # check similarity approximation
-        for pctd in range(0, 100, 5):
-            dec = pctd / 100
-            endp = floor(len(lorem_ipsum) * dec)
-            textd = lorem_ipsum[:endp] + lorem_ipsum[len(lorem_ipsum) - 1 : endp : -1]
-            sigd = minhash.get_signature(minhash.get_shingles(textd))
-            simil = sum([1 if a == b else 0 for ba, bb in zip(sig, sigd) for a, b in zip(ba, bb)]) / minhash.num_hashes
-            assert dec - 0.21 < simil < dec + 0.21
-
-        # check output file format and order
-        samples = [Document(f"sample {i}, {lorem_ipsum[i:: 10]}", data_id="test") for i in range(100)]
-        minhash(samples)
-        for bi in range(minhash.num_buckets):
-            with open(os.path.join(minhash.output_folder.path, f"bucket_{bi:03d}", "00000.minhash.sig"), "rb") as f:
-                prev = None
-                doc_ids = set()
-                S = struct.calcsize("I")
-                for di in range(100):
-                    data = struct.unpack("<%sI" % minhash.hashes_per_bucket, f.read(minhash.hashes_per_bucket * S))
-                    doc_id = struct.unpack("<I", f.read(S))[0]
-                    # ensure sorted order
-                    assert prev is None or data >= prev
-                    prev = data
-                    assert 0 <= doc_id < 100
-                    doc_ids.add(doc_id)
-                assert len(doc_ids) == 100
+        for use_64bit_hashes in (True, False):
+            config = MinhashConfig(use_64bit_hashes=use_64bit_hashes)
+            minhash = MinhashDedupSignature(
+                output_folder=LocalOutputDataFolder(os.path.join(self.test_dir, "signatures1")), config=config
+            )
+            shingles = minhash.get_shingles(lorem_ipsum)
+            sig = minhash.get_signature(shingles)
+
+            minhash2 = MinhashDedupSignature(
+                output_folder=LocalOutputDataFolder(os.path.join(self.test_dir, "signatures2")), config=config
+            )
+            # check consistency
+            assert sig == minhash2.get_signature(shingles)
+
+            # check correct number of outputs
+            assert len(sig) == minhash.config.num_buckets
+            assert all((len(x) == minhash.config.hashes_per_bucket for x in sig))
+
+            # check similarity approximation
+            for pctd in range(0, 100, 5):
+                dec = pctd / 100
+                endp = floor(len(lorem_ipsum) * dec)
+                textd = lorem_ipsum[:endp] + lorem_ipsum[len(lorem_ipsum) - 1 : endp : -1]
+                sigd = minhash.get_signature(minhash.get_shingles(textd))
+                simil = (
+                    sum([1 if a == b else 0 for ba, bb in zip(sig, sigd) for a, b in zip(ba, bb)]) / minhash.num_hashes
+                )
+                assert dec - 0.21 < simil < dec + 0.21
+
+            # check output file format and order
+            samples = [Document(f"sample {i}, {lorem_ipsum[i:: 10]}", data_id="test") for i in range(100)]
+            minhash(samples)
+            for bi in range(config.num_buckets):
+                with open(
+                    os.path.join(minhash.output_folder.path, f"bucket_{bi:03d}", "00000.minhash.sig"), "rb"
+                ) as f:
+                    prev = None
+                    doc_ids = set()
+                    S = np.dtype(config.hash_dtype).itemsize
+                    for di in range(100):
+                        data = struct.unpack(
+                            f"<%s{config.hash_format}" % config.hashes_per_bucket, f.read(config.hashes_per_bucket * S)
+                        )
+                        doc_id = struct.unpack("<I", f.read(struct.calcsize("I")))[0]
+                        # ensure sorted order
+                        assert prev is None or data >= prev
+                        prev = data
+                        assert 0 <= doc_id < 100
+                        doc_ids.add(doc_id)
+                    assert len(doc_ids) == 100
 
     def test_buckets_and_cluster(self):
-        sigs_folder = os.path.join(self.test_dir, "b_signatures")
-        buckets_folder = os.path.join(self.test_dir, "b_buckets")
-        clusters_folder = os.path.join(self.test_dir, "b_clusters")
-
-        signatures_block = MinhashDedupSignature(output_folder=LocalOutputDataFolder(sigs_folder))
-        buckets_block = MinhashDedupBuckets(
-            input_folder=LocalInputDataFolder(sigs_folder),
-            output_folder=LocalOutputDataFolder(buckets_folder),
-        )
-
-        clusters = [[0, 20, 50], [400, 420], [800, 810, 820, 840, 860], [1200, 1215, 1225, 1245], [1600], [2000]]
-
-        cluster_samples = [
-            Document(content=lorem_ipsum[x : x + 300], data_id=f"{ci}_{xi}", metadata={"ci": ci, "xi": xi})
-            for ci, cluster in enumerate(clusters)
-            for xi, x in enumerate(cluster)
-        ]
-
-        signatures_block(cluster_samples)
-        # test file read
-        for fi, file in enumerate(buckets_block.input_folder.list_files()):
-            last = None
-            for sig in read_sigs(file, fi, buckets_block.hashes_per_bucket):
-                assert 0 <= sig.doc_id < 100
-                assert last is None or sig.sig >= last
-                assert len(sig.sig) == buckets_block.hashes_per_bucket
-                last = sig.sig
-
-        # test duplicate pairs
-        for b in range(buckets_block.num_buckets):
-            buckets_block(None, bucket=b, world_size=buckets_block.num_buckets)
-        bucket_results_folder = LocalInputDataFolder(buckets_folder)
-        dup_files = bucket_results_folder.list_files(extension=".dups")
-        pairs = defaultdict(set)
-        for dup_file in dup_files:
-            with dup_file.open(binary=True) as df:
-                while data := df.read(4 * struct.calcsize("I")):
-                    f1, d1, f2, d2 = struct.unpack("<4I", data)
-                    assert f1 == f2 == 0
-                    assert cluster_samples[d1].metadata["ci"] == cluster_samples[d2].metadata["ci"]
-                    pairs[d1].add(d2)
-                    pairs[d2].add(d1)
-        doc_id = 0
-        for cluster in clusters:
-            for a in range(doc_id, doc_id + len(cluster)):
-                assert len(cluster) < 2 or any(a in pairs[b] for b in range(doc_id, doc_id + len(cluster)) if a != b)
-            doc_id += len(cluster)
-
-        # clustering
-        cluster_block = MinhashDedupCluster(bucket_results_folder, LocalOutputDataFolder(clusters_folder))
-        cluster_block(None)
-
-        cluster_results_folder = LocalInputDataFolder(clusters_folder)
-        remove_ids = set()
-        with cluster_results_folder.list_files()[0].open_binary() as df:
-            while data := df.read(struct.calcsize("I")):
-                remove_ids.add(struct.unpack("<I", data)[0])
-        doc_id = 0
-        kept = set()
-        for ci, cluster in enumerate(clusters):
-            to_remove = 0
-            for xi, a in enumerate(range(doc_id, doc_id + len(cluster))):
-                if a in remove_ids:
-                    to_remove += 1
-                else:
-                    kept.add(f"{ci}_{xi}")
-            doc_id += len(cluster)
-            assert to_remove == len(cluster) - 1
-
-        # filtering
-        filter_block = MinhashDedupFilter(cluster_results_folder)
-        filtered = filter_block(cluster_samples)
-        filtered_ids = {x.data_id for x in filtered}
-        assert filtered_ids == kept
+        for use_64bit_hashes in (True, False):
+            sigs_folder = os.path.join(self.test_dir, "b_signatures")
+            buckets_folder = os.path.join(self.test_dir, "b_buckets")
+            clusters_folder = os.path.join(self.test_dir, "b_clusters")
+            config = MinhashConfig(use_64bit_hashes=use_64bit_hashes)
+
+            signatures_block = MinhashDedupSignature(output_folder=LocalOutputDataFolder(sigs_folder), config=config)
+            buckets_block = MinhashDedupBuckets(
+                input_folder=LocalInputDataFolder(sigs_folder),
+                output_folder=LocalOutputDataFolder(buckets_folder),
+                config=config,
+            )
+
+            clusters = [[0, 20, 50], [400, 420], [800, 810, 820, 840, 860], [1200, 1215, 1225, 1245], [1600], [2000]]
+
+            cluster_samples = [
+                Document(content=lorem_ipsum[x : x + 300], data_id=f"{ci}_{xi}", metadata={"ci": ci, "xi": xi})
+                for ci, cluster in enumerate(clusters)
+                for xi, x in enumerate(cluster)
+            ]
+
+            signatures_block(cluster_samples)
+            # test file read
+            for fi, file in enumerate(buckets_block.input_folder.list_files()):
+                last = None
+                for sig in read_sigs(file, fi, config):
+                    assert 0 <= sig.doc_id < 100
+                    assert last is None or sig.sig >= last
+                    assert len(sig.sig) == config.hashes_per_bucket
+                    last = sig.sig
+
+            # test duplicate pairs
+            for b in range(config.num_buckets):
+                buckets_block(None, bucket=b, world_size=config.num_buckets)
+            bucket_results_folder = LocalInputDataFolder(buckets_folder)
+            dup_files = bucket_results_folder.list_files(extension=".dups")
+            pairs = defaultdict(set)
+            for dup_file in dup_files:
+                with dup_file.open(binary=True) as df:
+                    while data := df.read(4 * struct.calcsize("I")):
+                        f1, d1, f2, d2 = struct.unpack("<4I", data)
+                        assert f1 == f2 == 0
+                        assert cluster_samples[d1].metadata["ci"] == cluster_samples[d2].metadata["ci"]
+                        pairs[d1].add(d2)
+                        pairs[d2].add(d1)
+            doc_id = 0
+            for cluster in clusters:
+                for a in range(doc_id, doc_id + len(cluster)):
+                    assert len(cluster) < 2 or any(
+                        a in pairs[b] for b in range(doc_id, doc_id + len(cluster)) if a != b
+                    )
+                doc_id += len(cluster)
+
+            # clustering
+            cluster_block = MinhashDedupCluster(
+                bucket_results_folder, LocalOutputDataFolder(clusters_folder), config=config
+            )
+            cluster_block(None)
+
+            cluster_results_folder = LocalInputDataFolder(clusters_folder)
+            remove_ids = set()
+            with cluster_results_folder.list_files()[0].open_binary() as df:
+                while data := df.read(struct.calcsize("I")):
+                    remove_ids.add(struct.unpack("<I", data)[0])
+            doc_id = 0
+            kept = set()
+            for ci, cluster in enumerate(clusters):
+                to_remove = 0
+                for xi, a in enumerate(range(doc_id, doc_id + len(cluster))):
+                    if a in remove_ids:
+                        to_remove += 1
+                    else:
+                        kept.add(f"{ci}_{xi}")
+                doc_id += len(cluster)
+                assert to_remove == len(cluster) - 1
+
+            # filtering
+            filter_block = MinhashDedupFilter(cluster_results_folder)
+            filtered = filter_block(cluster_samples)
+            filtered_ids = {x.data_id for x in filtered}
+            assert filtered_ids == kept