<commit_before>from __future__ import annotations

import logging
import random
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING

import numpy
import numpy.typing
import sklearn.linear_model

import dedupe.core as core
import dedupe.training as training

if TYPE_CHECKING:
    from typing import Dict, Iterable, Literal, Mapping

    from dedupe._typing import Data, Labels, LabelsLike
    from dedupe._typing import RecordDictPair as TrainingExample
    from dedupe._typing import RecordDictPairs as TrainingExamples
    from dedupe._typing import RecordIDPair
    from dedupe.datamodel import DataModel
    from dedupe.predicates import Predicate


logger = logging.getLogger(__name__)


class HasCandidates:
    """Has a list of pairs that we could ask the user to label."""

    _candidates: TrainingExamples

    @property
    def candidates(self) -> TrainingExamples:
        return self._candidates

    def __len__(self) -> int:
        return len(self.candidates)


class Learner(ABC, HasCandidates):
    """A single learner that is used by DisagreementLearner."""

    _fitted: bool = False

    @abstractmethod
    def fit(self, pairs: TrainingExamples, y: LabelsLike) -> None:
        """Train on the given data."""

    @abstractmethod
    def candidate_scores(self) -> numpy.typing.NDArray[numpy.float_]:
        """For each of self.candidates, return our current guess [0,1] of if a match."""

    @abstractmethod
    def remove(self, index: int) -> None:
        """Remove a pair from self.candidates."""

    @staticmethod
    def _verify_fit_args(pairs: TrainingExamples, y: LabelsLike) -> list[Literal[0, 1]]:
        """Helper method to verify the arguments given to fit()"""
        if len(pairs) == 0:
            raise ValueError("pairs must have length of at least 1")
        y = list(y)
        if len(pairs) != len(y):
            raise ValueError(
                f"pairs and y must be same length. Got {len(pairs)} and {len(y)}"
            )
        return y


class MatchLearner(Learner):
    def __init__(self, data_model: DataModel, candidates: TrainingExamples):
        self.data_model = data_model
        self._candidates = candidates.copy()
        self._classifier = sklearn.linear_model.LogisticRegression()
        self._distances = self._calc_distances(self.candidates)

    def fit(self, pairs: TrainingExamples, y: LabelsLike) -> None:
        y = self._verify_fit_args(pairs, y)
        self._classifier.fit(self._calc_distances(pairs), numpy.array(y))
        self._fitted = True

    def remove(self, index: int) -> None:
        self._candidates.pop(index)
        self._distances = numpy.delete(self._distances, index, axis=0)

    def candidate_scores(self) -> numpy.typing.NDArray[numpy.float_]:
        if not self._fitted:
            raise ValueError("Must call fit() before candidate_scores()")
        scores: numpy.typing.NDArray[numpy.float_] = self._classifier.predict_proba(
            self._distances
        )[:, 1].reshape(-1, 1)
        return scores

    def _calc_distances(
        self, pairs: TrainingExamples
    ) -> numpy.typing.NDArray[numpy.float_]:
        return self.data_model.distances(pairs)


class BlockLearner(Learner):
    block_learner: training.BlockLearner

    def __init__(self):
        self.current_predicates: tuple[Predicate, ...] = ()
        self._cached_scores: numpy.typing.NDArray[numpy.float_] | None = None
        self._old_dupes: TrainingExamples = []

    def fit(self, pairs: TrainingExamples, y: LabelsLike) -> None:
        y = self._verify_fit_args(pairs, y)
        dupes = [pair for label, pair in zip(y, pairs) if label]

        new_dupes = [pair for pair in dupes if pair not in self._old_dupes]
        new_uncovered = not all(self._predict(new_dupes))

        if new_uncovered:
            self.current_predicates = self.block_learner.learn(
                dupes, recall=1.0, index_predicates=True, candidate_types="simple"
            )
            self._cached_scores = None
            self._old_dupes = dupes
        self._fitted = True

    def candidate_scores(self) -> numpy.typing.NDArray[numpy.float_]:
        if not self._fitted:
            raise ValueError("Must call fit() before candidate_scores()")
        if self._cached_scores is None:
            labels = self._predict(self.candidates)
            self._cached_scores = numpy.array(labels).reshape(-1, 1)

        return self._cached_scores

    def learn_predicates(
        self, dupes: TrainingExamples, recall: float, index_predicates: bool
    ) -> tuple[Predicate, ...]:
        return self.block_learner.learn(
            dupes,
            recall=recall,
            index_predicates=index_predicates,
            candidate_types="random forest",
        )

    def _predict(self, pairs: TrainingExamples) -> Labels:
        labels: Labels = []
        for record_1, record_2 in pairs:

            for predicate in self.current_predicates:
                keys = predicate(record_2, target=True)
                if keys:
                    if set(predicate(record_1)) & set(keys):
                        labels.append(1)
                        break
            else:
                labels.append(0)

        return labels

    def remove(self, index: int) -> None:
        self._candidates.pop(index)
        if self._cached_scores is not None:
            self._cached_scores = numpy.delete(self._cached_scores, index, axis=0)

    def _sample_indices(
        self, sample_size: int, max_cover: int
    ) -> Iterable[RecordIDPair]:

        weights: Dict[RecordIDPair, float] = {}
        for predicate, covered in self.block_learner.comparison_cover.items():
            # each predicate gets to vote for every record pair it covers. the
            # strength of that vote is in inverse proportion to the number of
            # records the predicate covers.
            #
            # if a predicate only covers a few record pairs, the value of
            # the vote it puts on those few pairs will be worth more than
            # a predicate that covers almost all the record pairs
            proportion = len(covered) / max_cover
            weight: float = numpy.exp(-1000 * proportion)
            if weight and proportion != 1:
                for pair in covered:
                    weights[pair] = weights.get(pair, 0.0) + weight

        sample_ids: Iterable[RecordIDPair]
        if sample_size < len(weights):
            # consider using a reservoir sampling strategy, which would
            # be more memory efficient and probably about as fast
            normalized_weights = numpy.fromiter(weights.values(), dtype=float) / sum(
                weights.values()
            )<commit_msg>reduce decay factor for sampling weights<commit_after>from __future__ import annotations

import logging
import random
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING

import numpy
import numpy.typing
import sklearn.linear_model

import dedupe.core as core
import dedupe.training as training

if TYPE_CHECKING:
    from typing import Dict, Iterable, Literal, Mapping

    from dedupe._typing import Data, Labels, LabelsLike
    from dedupe._typing import RecordDictPair as TrainingExample
    from dedupe._typing import RecordDictPairs as TrainingExamples
    from dedupe._typing import RecordIDPair
    from dedupe.datamodel import DataModel
    from dedupe.predicates import Predicate


logger = logging.getLogger(__name__)


class HasCandidates:
    """Has a list of pairs that we could ask the user to label."""

    _candidates: TrainingExamples

    @property
    def candidates(self) -> TrainingExamples:
        return self._candidates

    def __len__(self) -> int:
        return len(self.candidates)


class Learner(ABC, HasCandidates):
    """A single learner that is used by DisagreementLearner."""

    _fitted: bool = False

    @abstractmethod
    def fit(self, pairs: TrainingExamples, y: LabelsLike) -> None:
        """Train on the given data."""

    @abstractmethod
    def candidate_scores(self) -> numpy.typing.NDArray[numpy.float_]:
        """For each of self.candidates, return our current guess [0,1] of if a match."""

    @abstractmethod
    def remove(self, index: int) -> None:
        """Remove a pair from self.candidates."""

    @staticmethod
    def _verify_fit_args(pairs: TrainingExamples, y: LabelsLike) -> list[Literal[0, 1]]:
        """Helper method to verify the arguments given to fit()"""
        if len(pairs) == 0:
            raise ValueError("pairs must have length of at least 1")
        y = list(y)
        if len(pairs) != len(y):
            raise ValueError(
                f"pairs and y must be same length. Got {len(pairs)} and {len(y)}"
            )
        return y


class MatchLearner(Learner):
    def __init__(self, data_model: DataModel, candidates: TrainingExamples):
        self.data_model = data_model
        self._candidates = candidates.copy()
        self._classifier = sklearn.linear_model.LogisticRegression()
        self._distances = self._calc_distances(self.candidates)

    def fit(self, pairs: TrainingExamples, y: LabelsLike) -> None:
        y = self._verify_fit_args(pairs, y)
        self._classifier.fit(self._calc_distances(pairs), numpy.array(y))
        self._fitted = True

    def remove(self, index: int) -> None:
        self._candidates.pop(index)
        self._distances = numpy.delete(self._distances, index, axis=0)

    def candidate_scores(self) -> numpy.typing.NDArray[numpy.float_]:
        if not self._fitted:
            raise ValueError("Must call fit() before candidate_scores()")
        scores: numpy.typing.NDArray[numpy.float_] = self._classifier.predict_proba(
            self._distances
        )[:, 1].reshape(-1, 1)
        return scores

    def _calc_distances(
        self, pairs: TrainingExamples
    ) -> numpy.typing.NDArray[numpy.float_]:
        return self.data_model.distances(pairs)


class BlockLearner(Learner):
    block_learner: training.BlockLearner

    def __init__(self):
        self.current_predicates: tuple[Predicate, ...] = ()
        self._cached_scores: numpy.typing.NDArray[numpy.float_] | None = None
        self._old_dupes: TrainingExamples = []

    def fit(self, pairs: TrainingExamples, y: LabelsLike) -> None:
        y = self._verify_fit_args(pairs, y)
        dupes = [pair for label, pair in zip(y, pairs) if label]

        new_dupes = [pair for pair in dupes if pair not in self._old_dupes]
        new_uncovered = not all(self._predict(new_dupes))

        if new_uncovered:
            self.current_predicates = self.block_learner.learn(
                dupes, recall=1.0, index_predicates=True, candidate_types="simple"
            )
            self._cached_scores = None
            self._old_dupes = dupes
        self._fitted = True

    def candidate_scores(self) -> numpy.typing.NDArray[numpy.float_]:
        if not self._fitted:
            raise ValueError("Must call fit() before candidate_scores()")
        if self._cached_scores is None:
            labels = self._predict(self.candidates)
            self._cached_scores = numpy.array(labels).reshape(-1, 1)

        return self._cached_scores

    def learn_predicates(
        self, dupes: TrainingExamples, recall: float, index_predicates: bool
    ) -> tuple[Predicate, ...]:
        return self.block_learner.learn(
            dupes,
            recall=recall,
            index_predicates=index_predicates,
            candidate_types="random forest",
        )

    def _predict(self, pairs: TrainingExamples) -> Labels:
        labels: Labels = []
        for record_1, record_2 in pairs:

            for predicate in self.current_predicates:
                keys = predicate(record_2, target=True)
                if keys:
                    if set(predicate(record_1)) & set(keys):
                        labels.append(1)
                        break
            else:
                labels.append(0)

        return labels

    def remove(self, index: int) -> None:
        self._candidates.pop(index)
        if self._cached_scores is not None:
            self._cached_scores = numpy.delete(self._cached_scores, index, axis=0)

    def _sample_indices(
        self, sample_size: int, max_cover: int
    ) -> Iterable[RecordIDPair]:

        weights: Dict[RecordIDPair, float] = {}
        for predicate, covered in self.block_learner.comparison_cover.items():
            # each predicate gets to vote for every record pair it covers. the
            # strength of that vote is in inverse proportion to the number of
            # records the predicate covers.
            #
            # if a predicate only covers a few record pairs, the value of
            # the vote it puts on those few pairs will be worth more than
            # a predicate that covers almost all the record pairs
            proportion = len(covered) / max_cover
            weight: float = numpy.exp(-10 * proportion)
            if weight and proportion != 1:
                for pair in covered:
                    weights[pair] = weights.get(pair, 0.0) + weight

        sample_ids: Iterable[RecordIDPair]
        if sample_size < len(weights):
            # consider using a reservoir sampling strategy, which would
            # be more memory efficient and probably about as fast
            normalized_weights = numpy.fromiter(weights.values(), dtype=float) / sum(
                weights.values()
            )
{'fraction_non_alphanumeric': 0.0889218228973694, 'fraction_numerical': 0.004001482030381623, 'mean_word_length': 2.4551971326164876, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 35, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '184258', 'n_tokens_mistral': 3932, 'n_tokens_neox': 3712, 'n_words': 1093}