document_minhash_deduplicator.py

# Some code here has been modified from:
# https://github.com/bigcode-project/bigcode-dataset/blob/main/near_deduplication/minhash_deduplication.py
# --------------------------------------------------------

import hashlib
import struct
from collections import defaultdict
from typing import Optional

import numpy as np
import regex
from loguru import logger
from pydantic import Field, PositiveInt
from tqdm import tqdm
from typing_extensions import Annotated

from data_juicer.utils.constant import HashKeys
from data_juicer.utils.lazy_loader import LazyLoader
from data_juicer.utils.model_utils import prepare_sentencepiece_model

from ..base_op import OPERATORS, Deduplicator
from ..common.helper_func import UnionFind, split_on_whitespace

integrate = LazyLoader('integrate', 'scipy.integrate')

OP_NAME = 'document_minhash_deduplicator'

MERSENNE_PRIME = np.uint64((1 << 61) - 1)
MAX_HASH = np.uint64((1 << 32) - 1)


def sha1_hash32(data):
    """
    Directly taken from datasketch package to avoid dependency.

    Parameters
    ----------
    data : bytes

    Returns
    -------
    int
    """
    return struct.unpack('<I', hashlib.sha1(data).digest()[:4])[0]


def optimal_param(
    threshold: float,
    num_perm: int,
    false_positive_weight: float = 0.5,
    false_negative_weight: float = 0.5,
):
    """
    Compute the optimal `MinHashLSH` parameter that minimizes the weighted sum
    of probabilities of false positive and false negative, taken from
    datasketch.

    :param threshold: float. The threshold for similarity
    :param num_perm: int. The number of permutations
    :param false_positive_weight: float. The weight of false positive
    :param false_negative_weight: float. The weight of false negative
    :return: Tuple[int, int]. The optimal `b` and `r` parameters. The number of
        bands, and the number of rows per band respectively
    """

    def false_positive_probability(th: float, band: int, rows: int):
        """Source: `datasketch.lsh`"""

        def proba(s):
            return 1 - (1 - s**float(rows))**float(band)

        a, _ = integrate.quad(proba, 0.0, th)
        return a

    def false_negative_probability(th: float, band: int, rows: int):
        """Source: `datasketch.lsh`"""

        def proba(s):
            return 1 - (1 - (1 - s**float(rows))**float(band))

        a, _ = integrate.quad(proba, th, 1.0)
        return a

    # object: minimize the weighted FP and FN ratio
    min_error = float('inf')
    opt = (0, 0)
    for b in range(1, num_perm + 1):
        max_r = int(num_perm / b)
        for r in range(1, max_r + 1):
            fp = false_positive_probability(threshold, b, r)
            fn = false_negative_probability(threshold, b, r)
            error = fp * false_positive_weight + fn * false_negative_weight
            if error < min_error:
                min_error = error
                opt = (b, r)
    return opt


@OPERATORS.register_module(OP_NAME)
class DocumentMinhashDeduplicator(Deduplicator):
    """
    Deduplicator to deduplicate samples at document-level using MinHashLSH.

    Different from simhash, minhash is stored as bytes, so they won't be
    kept in the final dataset.
    """

    def __init__(
        self,
        tokenization: str = 'space',
        window_size: PositiveInt = 5,
        lowercase: bool = True,
        ignore_pattern: Optional[str] = None,
        num_permutations: PositiveInt = 256,
        jaccard_threshold: Annotated[float, Field(ge=0, le=1)] = 0.7,
        num_bands: Optional[PositiveInt] = None,
        num_rows_per_band: Optional[PositiveInt] = None,
        tokenizer_model: Optional[str] = None,
        *args,
        **kwargs,
    ):
        """
        Initialization method.

        :param tokenization: tokenization method for sample texts. It
            should be one of [space, punctuation, character,
            sentencepiece]. For English-like languages, we recommend
            to use 'space', for Chinese-like languages, we recommend
            to use 'character', and for multiple languages, we recommend
            to use 'sentencepiece'. If using 'sentencepiece', please
            provided the model path in the 'tokenizer_model' field.
        :param window_size: window size of shingling
        :param lowercase: whether to convert text to lower case first
        :param ignore_pattern: whether to ignore sub-strings with
            specific pattern when computing minhash
        :param num_permutations: number of permutations in minhash
            computing
        :param jaccard_threshold: the min jaccard similarity threshold
            in near-duplicate detection. When the jaccard similarity of
            two sample texts is >= this threshold, they are regarded as
            similar samples and this op will only keep one of them after
            deduplication
        :param num_bands: number of bands in LSH. Default it's None, and
            it will be determined by an optimal params computation
            algorithm by minimize the weighted sum of probs of False
            Positives and False Negatives
        :param num_rows_per_band: number of rows in each band in LSH.
            Default it's None, and it will be determined by an optimal
            params computation algorithm
        :param tokenizer_model: path for the sentencepiece model, used for
            sentencepiece tokenization.
        """
        super().__init__(*args, **kwargs)
        # about minhash computation
        self.tokenization = tokenization
        self.window_size = window_size
        self.lowercase = lowercase
        self.ignore_pattern = ignore_pattern
        if self.ignore_pattern:
            self.ignore_pattern = regex.compile(self.ignore_pattern)

        # check parameters
        if self.ignore_pattern and self.tokenization == 'punctuation':
            logger.warning('Be careful that tokenization with punctuations '
                           'won\'t work if the ignore pattern includes '
                           'punctuations.')
        self.punctuation_pattern = regex.compile(r'\p{P}')

        if self.tokenization == 'sentencepiece':
            if tokenizer_model is None:
                raise ValueError("To use 'sentencepiece' tokenization, "
                                 "'tokenizer_model' is required.")
            self.tokenizer = prepare_sentencepiece_model(tokenizer_model)
        else:
            self.tokenizer = None

        # about deduplication
        self.num_permutation = num_permutations
        self.jaccard_threshold = jaccard_threshold
        self.num_bands = num_bands
        self.num_rows_per_band = num_rows_per_band

        # initialize deduplication parameters
        # check number of bands and rows
        if self.num_bands is None or self.num_rows_per_band is None:
            self.num_bands, self.num_rows_per_band = optimal_param(
                self.jaccard_threshold,
                self.num_permutation,
            )

        # compute hash ranges and create hash tables
        self.hash_ranges = [(i * self.num_rows_per_band,
                             (i + 1) * self.num_rows_per_band)
                            for i in range(self.num_bands)]
        self.hash_tables = [defaultdict(set) for _ in range(self.num_bands)]

        # generate permutations
        gen = np.random.RandomState(seed=42)
        self.perm_a, self.perm_b = np.array(
            [(
                gen.randint(1, MERSENNE_PRIME, dtype=np.uint64),
                gen.randint(0, MERSENNE_PRIME, dtype=np.uint64),
            ) for _ in range(self.num_permutation)],
            dtype=np.uint64,
        ).T

    def compute_hash(self, sample):
        """
        Compute minhash values for the sample.

        :param sample: input sample
        :return: sample with minhash value.
        """
        # check if it's computed already
        if HashKeys.minhash in sample:
            return sample

        text = sample[self.text_key]

        if self.lowercase:
            text = text.lower()
        if self.ignore_pattern:
            text = self.ignore_pattern.sub('', text)

        # get tokens for different tokenization method
        tokens = set()
        if self.tokenization == 'character':
            tokens = {
                str.encode(text[i:i + self.window_size])
                for i in range(len(text) - self.window_size)
            }
        elif self.tokenization == 'punctuation':
            tokens = self.punctuation_pattern.split(text)
            tokens = {
                str.encode(' '.join(tokens[i:i + self.window_size]))
                for i in range(len(tokens) - self.window_size)
            }
        elif self.tokenization == 'space':
            tokens = split_on_whitespace(text)
            tokens = {
                str.encode(' '.join(tokens[i:i + self.window_size]))
                for i in range(len(tokens) - self.window_size)
            }
        elif self.tokenization == 'sentencepiece':
            tokens = self.tokenizer.encode(text, out_type=str)
            tokens = {
                str.encode(''.join(tokens[i:i + self.window_size]))
                for i in range(len(tokens) - self.window_size)
            }
        else:
            raise NotImplementedError(
                f'Unimplemented tokenization method [{self.tokenization}]')

        # compute minhash value
        hv = np.array([sha1_hash32(token) for token in tokens],
                      dtype=np.uint64)
        phv = np.bitwise_and(
            ((hv * np.tile(self.perm_a,
                           (len(hv), 1)).T).T + self.perm_b) % MERSENNE_PRIME,
            MAX_HASH)
        hash_values = np.vstack([
            phv,
            np.ones(self.num_permutation, dtype=np.uint64) * MAX_HASH
        ]).min(axis=0)
        sample[HashKeys.minhash] = [
            bytes(hash_values[start:end].byteswap().data)
            for start, end in self.hash_ranges
        ]
        return sample

    def process(self, dataset, show_num=0):
        """
        For doc-level, dataset --> dataset.

        :param dataset: input dataset
        :param show_num: number of traced samples used when tracer is
            open.
        :return: deduplicated dataset and the sampled duplicate pairs.
        """
        # no need to deduplicate because too few samples
        if len(dataset) <= 1:
            return dataset, {}

        minhashes = dataset[HashKeys.minhash]
        # remove bytes minhash column otherwise unexpected error would occur
        # when exporting the processed dataset
        dataset = dataset.remove_columns([HashKeys.minhash])

        # make clusters -- construct the minhash lookup tables of seg to ids
        logger.info(f'Start clustering for {len(dataset)} samples...')
        batch_size = 10000
        for i in tqdm(range(0, len(minhashes), batch_size),
                      dynamic_ncols=True,
                      desc='Iterating MinHashes of samples...'):
            batch = minhashes[i:i + batch_size]
            for idx, hs in enumerate(batch):
                for h, hashtable in zip(hs, self.hash_tables):
                    hashtable[h].add(idx + i)

        # using UnionFind set to union samples within the same clusters
        union_find = UnionFind()
        for table in tqdm(self.hash_tables,
                          dynamic_ncols=True,
                          desc='Clustering'):
            for cluster in table.values():
                if len(cluster) <= 1:
                    continue
                idx = min(cluster)
                for x in cluster:
                    union_find.union(x, idx)
        logger.info(f'There are {len(set(union_find.parent.values()))} '
                    f'clusters that includes multiple near-duplicate samples.')

        # record the duplicate sample pairs
        dup_pairs = {}
        if show_num > 0:
            for i in range(len(dataset)):
                cluster_idx = union_find.find(i)
                if cluster_idx not in dup_pairs and cluster_idx != i:
                    dup_pairs[cluster_idx] = [
                        dataset[cluster_idx],
                        dataset[i],
                    ]
                if len(dup_pairs) >= show_num:
                    break

        # filtering -- only keep those samples whose parent index is itself,
        # including:
        # 1. samples that form a cluster by themselves
        # 2. the first sample in a cluster that includes multiple samples
        def _filter_minhash_dup_helper(sample, index):
            return union_find.find(index) == index

        dataset = dataset.filter(
            _filter_minhash_dup_helper,
            with_indices=True,
        )
        logger.info(f'Keep {len(dataset)} samples after MinHash dedup.')

        return dataset, dup_pairs