Source code for data_juicer.ops.deduplicator.document_minhash_deduplicator

# 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


[docs]@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. """
[docs] 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
[docs] 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
[docs] 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