<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-

from collections import defaultdict
import collections
import itertools
import logging
import backport
from zope.index.text.textindex import TextIndex
from zope.index.text.cosineindex import CosineIndex
from zope.index.text.lexicon import Lexicon
from zope.index.text.lexicon import Splitter
from zope.index.text.stopdict import get_stopdict
import time
import dedupe.tfidf as tfidf

logger = logging.getLogger(__name__)

class Blocker:
    '''Takes in a record and returns all blocks that record belongs to'''
    def __init__(self, 
                 predicates, 
                 stop_words = None) :

        if stop_words is None :
            stop_words = defaultdict(lambda : set(get_stopdict()))

        self.predicates = predicates

        self.stop_words = stop_words

        self.tfidf_fields = defaultdict(set)

        for full_predicate in predicates :
            for predicate in full_predicate :
                if predicate.type == "TfidfPredicate" :
                    self.tfidf_fields[predicate.field].add(predicate)

    #@profile
    def __call__(self, records):

        start_time = time.time()
        predicates = [(':' + str(i), predicate)
                      for i, predicate
                      in enumerate(self.predicates)]

        for i, record in enumerate(records) :
            record_id, instance = record
    
            for pred_id, predicate in predicates :
                block_keys = predicate(record_id, instance)
                for block_key in block_keys :
                    yield block_key + pred_id, record_id
            
            if i % 10000 == 0 :
                logger.info('%(iteration)d, %(elapsed)f2 seconds',
                             {'iteration' :i,
                              'elapsed' :time.time() - start_time})



class DedupeBlocker(Blocker) :
    

    def tfIdfBlock(self, data, field): 
        '''Creates TF/IDF canopy of a given set of data'''

        splitter = Splitter()

        stop_word_remover = CustomStopWordRemover(self.stop_words[field])



        indices = {}
        for predicate in self.tfidf_fields[field] :
            indices[predicate] = TextIndex(Lexicon(splitter, stop_word_remover))
            indices[predicate].index = CosineIndex(indices[predicate].lexicon)
            pipeline = indices[predicate].lexicon._pipeline

        index_to_id = {}
        base_tokens = {}


        for i, (record_id, doc) in enumerate(data, 1) :
            index_to_id[i] = record_id
            last = [doc]
            for each in pipeline :
                last = each.process(last)
            base_tokens[i] = ' OR '.join(last)
            for index in indices.values() :
                index.index_doc(i, doc)

        logger.info(time.asctime())                

        for predicate in self.tfidf_fields[field] :
            logger.info("Canopy: %s", str(predicate))
            canopy = tfidf.makeCanopy(indices[predicate],
                                      base_tokens, 
                                      predicate.threshold)
            predicate.canopy = dict((index_to_id[k], index_to_id[v])
                                    for k, v
                                    in canopy.iteritems())
        
        logger.info(time.asctime())                
               
class RecordLinkBlocker(Blocker) :
    def tfIdfBlock(self, data_1, data_2, field): 
        '''Creates TF/IDF canopy of a given set of data'''

        splitter = Splitter()

        stop_word_remover = CustomStopWordRemover(self.stop_words[field])

        indices = {}
        for predicate in self.tfidf_fields[field] :
            indices[predicate] = TextIndex(Lexicon(splitter, stop_word_remover))
            indices[predicate].index = CosineIndex(indices[predicate].lexicon)
            pipeline = indices[predicate].lexicon._pipeline

        index_to_id = {}
        base_tokens = {}

        i = 1

        for record_id, doc in data_1 :
            index_to_id[i] = record_id
            last = [doc]
            for each in pipeline :
                last = each.process(last)
            base_tokens[i] = ' OR '.join(last)
            i += 1

        for record_id, doc in data_2  :
            index_to_id[i] = record_id
            for index in indices.values() :
                index.index_doc(i, doc)
            i += 1

        for predicate in self.tfidf_fields[field] :
            logger.info("Canopy: %s", str(predicate))
            canopy = tfidf.makeCanopy(indices[predicate],
                                      base_tokens, 
                                      predicate.threshold)
            predicate.canopy = dict((index_to_id[k], index_to_id[v])
                                    for k, v
                                    in canopy.iteritems())


def blockTraining(training_pairs,
                  predicate_set,
                  eta=.1,
                  epsilon=.1,
                  matching = "Dedupe"):
    '''
    Takes in a set of training pairs and predicates and tries to find
    a good set of blocking rules.
    '''

    # Setup

    training_dupes = set(training_pairs['match'])
    training_distinct = set(training_pairs['distinct'])

    if matching == "RecordLink" :
        coverage = RecordLinkCoverage(predicate_set,
                                      training_dupes | training_distinct)

    else :
        coverage = DedupeCoverage(predicate_set,
                                  training_dupes | training_distinct)


    print len(coverage.overlap.keys())

    # Compound Predicates
    compound_predicates = itertools.combinations(coverage.overlap, 2)
    intersection = set.intersection

    for compound_predicate in compound_predicates :
        compound_predicate = CompoundPredicate(compound_predicate)
        predicate_1, predicate_2 = compound_predicate
        
        coverage.overlap[compound_predicate] =\
            intersection(coverage.overlap[predicate_1],
                         coverage.overlap[predicate_2])

        i = 0
        for blocks in itertools.product(coverage.blocks[predicate_1].values(),
                                        coverage.blocks[predicate_2].values()) :
            coverage.blocks[compound_predicate][i] =\<commit_msg>remove errant print statement<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-

from collections import defaultdict
import collections
import itertools
import logging
import backport
from zope.index.text.textindex import TextIndex
from zope.index.text.cosineindex import CosineIndex
from zope.index.text.lexicon import Lexicon
from zope.index.text.lexicon import Splitter
from zope.index.text.stopdict import get_stopdict
import time
import dedupe.tfidf as tfidf

logger = logging.getLogger(__name__)

class Blocker:
    '''Takes in a record and returns all blocks that record belongs to'''
    def __init__(self, 
                 predicates, 
                 stop_words = None) :

        if stop_words is None :
            stop_words = defaultdict(lambda : set(get_stopdict()))

        self.predicates = predicates

        self.stop_words = stop_words

        self.tfidf_fields = defaultdict(set)

        for full_predicate in predicates :
            for predicate in full_predicate :
                if predicate.type == "TfidfPredicate" :
                    self.tfidf_fields[predicate.field].add(predicate)

    #@profile
    def __call__(self, records):

        start_time = time.time()
        predicates = [(':' + str(i), predicate)
                      for i, predicate
                      in enumerate(self.predicates)]

        for i, record in enumerate(records) :
            record_id, instance = record
    
            for pred_id, predicate in predicates :
                block_keys = predicate(record_id, instance)
                for block_key in block_keys :
                    yield block_key + pred_id, record_id
            
            if i % 10000 == 0 :
                logger.info('%(iteration)d, %(elapsed)f2 seconds',
                             {'iteration' :i,
                              'elapsed' :time.time() - start_time})



class DedupeBlocker(Blocker) :
    

    def tfIdfBlock(self, data, field): 
        '''Creates TF/IDF canopy of a given set of data'''

        splitter = Splitter()

        stop_word_remover = CustomStopWordRemover(self.stop_words[field])



        indices = {}
        for predicate in self.tfidf_fields[field] :
            indices[predicate] = TextIndex(Lexicon(splitter, stop_word_remover))
            indices[predicate].index = CosineIndex(indices[predicate].lexicon)
            pipeline = indices[predicate].lexicon._pipeline

        index_to_id = {}
        base_tokens = {}


        for i, (record_id, doc) in enumerate(data, 1) :
            index_to_id[i] = record_id
            last = [doc]
            for each in pipeline :
                last = each.process(last)
            base_tokens[i] = ' OR '.join(last)
            for index in indices.values() :
                index.index_doc(i, doc)

        logger.info(time.asctime())                

        for predicate in self.tfidf_fields[field] :
            logger.info("Canopy: %s", str(predicate))
            canopy = tfidf.makeCanopy(indices[predicate],
                                      base_tokens, 
                                      predicate.threshold)
            predicate.canopy = dict((index_to_id[k], index_to_id[v])
                                    for k, v
                                    in canopy.iteritems())
        
        logger.info(time.asctime())                
               
class RecordLinkBlocker(Blocker) :
    def tfIdfBlock(self, data_1, data_2, field): 
        '''Creates TF/IDF canopy of a given set of data'''

        splitter = Splitter()

        stop_word_remover = CustomStopWordRemover(self.stop_words[field])

        indices = {}
        for predicate in self.tfidf_fields[field] :
            indices[predicate] = TextIndex(Lexicon(splitter, stop_word_remover))
            indices[predicate].index = CosineIndex(indices[predicate].lexicon)
            pipeline = indices[predicate].lexicon._pipeline

        index_to_id = {}
        base_tokens = {}

        i = 1

        for record_id, doc in data_1 :
            index_to_id[i] = record_id
            last = [doc]
            for each in pipeline :
                last = each.process(last)
            base_tokens[i] = ' OR '.join(last)
            i += 1

        for record_id, doc in data_2  :
            index_to_id[i] = record_id
            for index in indices.values() :
                index.index_doc(i, doc)
            i += 1

        for predicate in self.tfidf_fields[field] :
            logger.info("Canopy: %s", str(predicate))
            canopy = tfidf.makeCanopy(indices[predicate],
                                      base_tokens, 
                                      predicate.threshold)
            predicate.canopy = dict((index_to_id[k], index_to_id[v])
                                    for k, v
                                    in canopy.iteritems())


def blockTraining(training_pairs,
                  predicate_set,
                  eta=.1,
                  epsilon=.1,
                  matching = "Dedupe"):
    '''
    Takes in a set of training pairs and predicates and tries to find
    a good set of blocking rules.
    '''

    # Setup

    training_dupes = set(training_pairs['match'])
    training_distinct = set(training_pairs['distinct'])

    if matching == "RecordLink" :
        coverage = RecordLinkCoverage(predicate_set,
                                      training_dupes | training_distinct)

    else :
        coverage = DedupeCoverage(predicate_set,
                                  training_dupes | training_distinct)

    # Compound Predicates
    compound_predicates = itertools.combinations(coverage.overlap, 2)
    intersection = set.intersection

    for compound_predicate in compound_predicates :
        compound_predicate = CompoundPredicate(compound_predicate)
        predicate_1, predicate_2 = compound_predicate
        
        coverage.overlap[compound_predicate] =\
            intersection(coverage.overlap[predicate_1],
                         coverage.overlap[predicate_2])

        i = 0
        for blocks in itertools.product(coverage.blocks[predicate_1].values(),
                                        coverage.blocks[predicate_2].values()) :
            coverage.blocks[compound_predicate][i] =\
{'fraction_non_alphanumeric': 0.08637556669052732, 'fraction_numerical': 0.004135846655531695, 'mean_word_length': 1.751422319474836, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 22, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4581082', 'n_tokens_mistral': 3604, 'n_tokens_neox': 3395, 'n_words': 759}