<commit_before>#!/usr/bin/env python

from subprocess import Popen, PIPE
import os
import re
import signal
import unittest

class MasterTestCase(unittest.TestCase):

    def setUp(self):
        if os.path.exists("/tmp/mbsocket"):
            os.unlink("/tmp/mbsocket")

        self.uts = Popen(["./unit-test-server", "rtu"], stdout=PIPE, stderr=PIPE)
        self.rec = Popen(["../src/mbrecorder"], stdout=PIPE)
        self.col = Popen(["../src/mbcollect", "-f", "master-test-case.ini"], stdout=PIPE)
        self.is_master = True

    def tearDown(self):
        self.col.send_signal(signal.SIGTERM)
        self.rec.send_signal(signal.SIGTERM)
        self.uts.send_signal(signal.SIGTERM)
        self.col.wait()
        self.rec.wait()
        self.uts.wait()

    def test_output(self):
        # MBCOLLECT
        # First line contains host name
        line = self.col.stdout.readline()
        if self.is_master:
            m = re.search("Slave name (\w+), ID (\d+)", line)
            self.assertEqual(m.group(2), "1")
        else:
            m = re.search("Server name (\w+), IP ([\.\d]+):(\d+)", line)
            self.assertEqual(m.group(2), "127.0.0.1")
            self.assertEqual(m.group(3), "1502")

        server_name = m.group(1)
        self.assertEqual(server_name, "hello")

        # Check parsing of values to read
        c = re.compile("Address (\d+) => (\d+) values \((\w+)\)")

        expected_values = [
            {'a': 0, 'l': 2, 'type': 'floatmsb', 'values': ['12345.599609']},
            {'a': 2, 'l': 2, 'type': 'floatlsb', 'values': ['789.000000']},
            {'a': 0, 'l': 4, 'type': 'floatmsb', 'values': [
                '12345.599609', '0.000000']},
            {'a': 4, 'l': 1, 'type': 'int', 'values': ['1']},
            {'a': 4, 'l': 3, 'type': 'int', 'values': ['1', '2', '3']},
        ]
        for d in expected_values:
            m = c.match(self.col.stdout.readline())
            self.assertEqual(m.group(1), str(d['a']))
            self.assertEqual(m.group(2), str(d['l']))
            self.assertEqual(m.group(3), d['type'])

        line = self.col.stdout.readline()
        if self.is_master:
            self.assertEqual(line, "Opening /dev/ttyUSB0 at 115200 bauds (N, 8, 1)\n")
        else:
            self.assertEqual(line, "Connecting to 127.0.0.1:1502\n")

        # MBRECORDER
        # Ignore remaining output for collector and check recorder
        c = re.compile("mb_%s_(\d+) ([\d\.]+)([\||\n]?)" % server_name)

        for d in expected_values:
            line = self.rec.stdout.readline()

            if d['type'].find('int') == 0:
                l = d['l']
                step = 1
            else:
                l = d['l'] / 2
                step = 2

            start = 0
            for i in range(l):
                m = c.search(line, start)
                self.assertIsNotNone(m, "Item: %d" % i)
                # Address
                self.assertEqual(m.group(1), str(d['a'] + i * step))
                # Value read
                self.assertEqual(m.group(2), str(d['values'][i]))
                # Separator
                if i != l - 1:
                    self.assertEqual(m.group(3), '|')
                else:
                    self.assertEqual(m.group(3), '\n')

                start = m.end()

class ClientTestCase(MasterTestCase):

    def setUp(self):
        if os.path.exists("/tmp/mbsocket"):
            os.unlink("/tmp/mbsocket")

        self.uts = Popen(["./unit-test-server", "tcp"], stdout=PIPE, stderr=PIPE)
        self.rec = Popen(["../src/mbrecorder"], stdout=PIPE)
        self.col = Popen(["../src/mbcollect", "-f", "client-test-case.ini"], stdout=PIPE)
        self.is_master = False

class SlaveTestCase(unittest.TestCase):

    def setUp(self):
        if os.path.exists("/tmp/mbsocket"):
            os.unlink("/tmp/mbsocket")

        self.rec = Popen(["../src/mbrecorder"], stdout=PIPE)
        self.col = Popen(["../src/mbcollect", "-f", "slave-test-case.ini"], stdout=PIPE)
        self.uts = Popen(["./unit-test-client", "rtu"], stdout=PIPE, stderr=PIPE)

    def tearDown(self):
        self.col.send_signal(signal.SIGTERM)
        self.rec.send_signal(signal.SIGTERM)
        self.col.wait()
        self.rec.wait()
        # Client stops himself

    def test_recorder(self):
        # Client write only one register
        self.assertEqual(self.rec.stdout.readline(), "mb_0 1234\n")
        # Another query with two registers
        self.assertEqual(self.rec.stdout.readline(), "mb_1 5678|mb_2 9012\n")

class ServerTestCase(SlaveTestCase):
    """Inhrerits from SlaveTestCase test_recorder and tearDown"""

    def setUp(self):
        if os.path.exists("/tmp/mbsocket"):
            os.unlink("/tmp/mbsocket")

        self.rec = Popen(["../src/mbrecorder"], stdout=PIPE)
        self.col = Popen(["../src/mbcollect", "-f", "server-test-case.ini"], stdout=PIPE)
        self.uts = Popen(["./unit-test-client", "tcp"], stdout=PIPE, stderr=PIPE)

if __name__ == '__main__':
    unittest.main()
<commit_msg>Improve unit tests<commit_after>#!/usr/bin/env python

from subprocess import Popen, PIPE
import os
import re
import signal
import unittest

class CommonMSTestCase(unittest.TestCase):

    def setUp(self):
        self.expected_values = [
            {'a': 0, 'l': 2, 'type': 'floatmsb', 'values': ['12345.599609']},
            {'a': 2, 'l': 2, 'type': 'floatlsb', 'values': ['789.000000']},
            {'a': 0, 'l': 4, 'type': 'floatmsb', 'values': [
                '12345.599609', '0.000000']},
            {'a': 4, 'l': 1, 'type': 'int', 'values': ['1']},
            {'a': 4, 'l': 3, 'type': 'int', 'values': ['1', '2', '3']},
        ]

    def check_parsing(self, process):
        # Check parsing of values to read
        c = re.compile("Address (\d+) => (\d+) values \((\w+)\)")

        for d in self.expected_values:
            m = c.match(process.stdout.readline())
            self.assertEqual(m.group(1), str(d['a']))
            self.assertEqual(m.group(2), str(d['l']))
            self.assertEqual(m.group(3), d['type'])

    def check_recorder(self, recorder, server_name):
        # MBRECORDER
        # Ignore remaining output for collector and check recorder
        c = re.compile("mb_%s_(\d+) ([\d\.]+)([\||\n]?)" % server_name)

        for d in self.expected_values:
            line = recorder.stdout.readline()

            if d['type'].find('int') == 0:
                l = d['l']
                step = 1
            else:
                l = d['l'] / 2
                step = 2

            start = 0
            for i in range(l):
                m = c.search(line, start)
                self.assertIsNotNone(m, "Item: %d" % i)
                # Address
                self.assertEqual(m.group(1), str(d['a'] + i * step))
                # Value read
                self.assertEqual(m.group(2), str(d['values'][i]))
                # Separator
                if i != l - 1:
                    self.assertEqual(m.group(3), '|')
                else:
                    self.assertEqual(m.group(3), '\n')

                start = m.end()

class MasterTestCase(CommonMSTestCase):

    def setUp(self):
        super(MasterTestCase, self).setUp()
        if os.path.exists("/tmp/mbsocket"):
            os.unlink("/tmp/mbsocket")

        self.uts = Popen(["./unit-test-server", "rtu"], stdout=PIPE, stderr=PIPE)
        self.rec = Popen(["../src/mbrecorder"], stdout=PIPE)
        self.col = Popen(["../src/mbcollect", "-f", "master-test-case.ini"], stdout=PIPE)

    def tearDown(self):
        self.col.send_signal(signal.SIGTERM)
        self.rec.send_signal(signal.SIGTERM)
        self.uts.send_signal(signal.SIGTERM)
        self.col.wait()
        self.rec.wait()
        self.uts.wait()

    def test_output(self):
        line = self.col.stdout.readline()
        m = re.search("Slave name (\w+), ID (\d+)", line)
        self.assertEqual(m.group(2), "1")

        server_name = m.group(1)
        self.assertEqual(server_name, "hello")

        self.check_parsing(self.col)

        line = self.col.stdout.readline()
        self.assertEqual(line, "Opening /dev/ttyUSB0 at 115200 bauds (N, 8, 1)\n")

        self.check_recorder(self.rec, server_name)

class ClientTestCase(CommonMSTestCase):

    def setUp(self):
        super(ClientTestCase, self).setUp()

        if os.path.exists("/tmp/mbsocket"):
            os.unlink("/tmp/mbsocket")

        self.uts = Popen(["./unit-test-server", "tcp"], stdout=PIPE, stderr=PIPE)
        self.rec = Popen(["../src/mbrecorder"], stdout=PIPE)
        self.col = Popen(["../src/mbcollect", "-f", "client-test-case.ini"], stdout=PIPE)

    def test_output(self):
        line = self.col.stdout.readline()
        m = re.search("Server name (\w+), IP ([\.\d]+):(\d+)", line)
        self.assertEqual(m.group(2), "127.0.0.1")
        self.assertEqual(m.group(3), "1502")

        server_name = m.group(1)
        self.assertEqual(server_name, "hello")

        self.check_parsing(self.col)

        line = self.col.stdout.readline()
        self.assertEqual(line, "Connecting to 127.0.0.1:1502\n")
        self.check_recorder(self.rec, server_name)

class SlaveTestCase(unittest.TestCase):

    def setUp(self):
        if os.path.exists("/tmp/mbsocket"):
            os.unlink("/tmp/mbsocket")

        self.rec = Popen(["../src/mbrecorder"], stdout=PIPE)
        self.col = Popen(["../src/mbcollect", "-f", "slave-test-case.ini"], stdout=PIPE)
        self.uts = Popen(["./unit-test-client", "rtu"], stdout=PIPE, stderr=PIPE)

    def tearDown(self):
        self.col.send_signal(signal.SIGTERM)
        self.rec.send_signal(signal.SIGTERM)
        # Client stops himself
        self.col.wait()
        self.rec.wait()
        self.uts.wait()

    def test_recorder(self):
        # Client write only one register
        self.assertEqual(self.rec.stdout.readline(), "mb_0 1234\n")
        # Another query with two registers
        self.assertEqual(self.rec.stdout.readline(), "mb_1 5678|mb_2 9012\n")

class ServerTestCase(SlaveTestCase):
    """Inhrerits from SlaveTestCase test_recorder and tearDown"""

    def setUp(self):
        if os.path.exists("/tmp/mbsocket"):
            os.unlink("/tmp/mbsocket")

        self.rec = Popen(["../src/mbrecorder"], stdout=PIPE)
        self.col = Popen(["../src/mbcollect", "-f", "server-test-case.ini"], stdout=PIPE)
        self.uts = Popen(["./unit-test-client", "tcp"], stdout=PIPE, stderr=PIPE)

if __name__ == '__main__':
    unittest.main()
<|endoftext|>


<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})

        # clear canopies to reduce memory usage
        for predicate_set in self.tfidf_fields.values() :
            for predicate in predicate_set :
                predicate.canopy = {}

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
            stringify = predicate.stringify

        index_to_id = {}
        base_tokens = {}

        for i, (record_id, doc) in enumerate(data, 1) :
            doc = stringify(doc)
            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
            stringify = predicate.stringify

        index_to_id = {}
        base_tokens = {}

        i = 1

        for record_id, doc in data_1 :
            doc = stringify(doc)
            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  :
            doc = stringify(doc)
            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] =\
                intersection(*blocks)
            i += 1

    predicate_set = coverage.overlap.keys()

    # Only consider predicates that cover at least one duplicate pair
    dupe_coverage = coverage.predicateCoverage(predicate_set,
                                               training_dupes)
    predicate_set = dupe_coverage.keys()

    # Within blocks, we will compare every combination of
    # records. Therefore, we want to avoid predicates that make large
    # blocks.
    for pred in predicate_set[:] :
        blocks = coverage.blocks[pred].itervalues()
        if any(len(block) >= 100 for block in blocks if block) :
            predicate_set.remove(pred)

    # As an efficency, we can throw away the predicates that cover too
    # many distinct pairs
    coverage_threshold = eta * len(training_distinct)
    logger.info("coverage threshold: %s", coverage_threshold)

    distinct_coverage = coverage.predicateCoverage(predicate_set,
                                                   training_distinct)

    for pred, pairs in distinct_coverage.items() :
        if len(pairs) > coverage_threshold :
            predicate_set.remove(pred)

    distinct_coverage = coverage.predicateCoverage(predicate_set,
                                                   training_distinct)

    final_predicate_set = findOptimumBlocking(training_dupes,
                                              predicate_set,
                                              distinct_coverage,
                                              epsilon,
                                              coverage)

    logger.info('Final predicate set:')
    for predicate in final_predicate_set :
        logger.info(predicate)

    if final_predicate_set:
        return final_predicate_set, coverage.stop_words
    else:
        raise ValueError('No predicate found! We could not learn a single good predicate. Maybe give Dedupe more training data')

def findOptimumBlocking(uncovered_dupes,
                        predicate_set,
                        distinct_coverage,
                        epsilon,
                        coverage):

    # Greedily find the predicates that, at each step, covers the
    # most duplicates and covers the least distinct pairs, due to
    # Chvatal, 1979
    #
    # We would like to optimize the ratio of the probability of of a
    # predicate covering a duplicate pair versus the probability of
    # covering a distinct pair. If we have a uniform prior belief
    # about those probabilities, we can estimate these probabilities as
    #
    # (predicate_covered_dupe_pairs + 1) / (all_dupe_pairs + 2)
    #
    # (predicate_covered_distinct_pairs + 1) / (all_distinct_pairs + 2)
    #
    # When we are trying to find the best predicate among a set of
    # predicates, the denominators factor out and our coverage
    # estimator becomes
    #
    # (predicate_covered_dupe_pairs + 1)/ (predicate_covered_distinct_pairs + 1)

    dupe_coverage = coverage.predicateCoverage(predicate_set,
                                               uncovered_dupes)

    uncovered_dupes = set(uncovered_dupes)

    final_predicate_set = []
    while len(uncovered_dupes) > epsilon:

        best_cover = 0
        best_predicate = None
        for predicate in dupe_coverage :
            dupes = len(dupe_coverage[predicate])
            distinct = len(distinct_coverage[predicate])
            cover = (dupes + 1.0)/(distinct + 1.0)
            if cover > best_cover:
                best_cover = cover
                best_predicate = predicate
                best_distinct = distinct
                best_dupes = dupes

        if not best_predicate:
            logger.warning('Ran out of predicates')
            break

        final_predicate_set.append(best_predicate)
        predicate_set.remove(best_predicate)

        uncovered_dupes = uncovered_dupes - dupe_coverage[best_predicate]
        dupe_coverage = coverage.predicateCoverage(predicate_set,
                                                   uncovered_dupes)

        logger.debug(best_predicate)
        logger.debug('cover: %(cover)f, found_dupes: %(found_dupes)d, '
                      'found_distinct: %(found_distinct)d, '
                      'uncovered dupes: %(uncovered)d',
                      {'cover' : best_cover,
                       'found_dupes' : best_dupes,
                       'found_distinct' : best_distinct,
                       'uncovered' : len(uncovered_dupes)
                       })

    return final_predicate_set

class Coverage(object) :

    def coveredBy(self, predicates, record_ids, pairs) :
        self.overlap = defaultdict(set)
        self.blocks = defaultdict(lambda : defaultdict(set))

        for pair in pairs :
            record_1, record_2 = pair
            record_1_id = record_ids[record_1]
            record_2_id = record_ids[record_2]
            for predicate in predicates :
                field_predicate_1 = predicate(record_1_id, record_1)
                if field_predicate_1:
                    field_predicate_2 = predicate(record_2_id, record_2)
                    if field_predicate_2 :
                        field_preds = (set(field_predicate_2)
                                       & set(field_predicate_1))
                        if field_preds :
                            self.overlap[predicate].add(pair)
                            for field_pred in field_preds :
                                self.blocks[predicate][field_pred].add(pair)

        for predicate, coverage in self.blocks.items() :
            for field_pred, block in coverage.items() :
                if len(block) < 100 :
                    del self.blocks[predicate][field_pred]

    def predicateCoverage(self,
                          predicate_set,
                          pairs) :

        coverage = defaultdict(set)
        pairs = set(pairs)

        for predicate in predicate_set :
            covered_pairs = pairs.intersection(self.overlap[predicate])
            if covered_pairs :
                coverage[predicate] = covered_pairs

        return coverage

class DedupeCoverage(Coverage) :
    def __init__(self, predicate_set, pairs) :

        records = set(itertools.chain(*pairs))

        id_records = dict(itertools.izip(itertools.count(), records))
        record_ids = dict(itertools.izip(records, itertools.count()))

        blocker = DedupeBlocker(predicate_set)

        for field in blocker.tfidf_fields :
            data = [(record_id, record[field])
                    for record_id, record
                    in id_records.items()]
            stop_words = stopWords(data)
            blocker.stop_words[field].update(stop_words)
            blocker.tfIdfBlock(data, field)

        self.stop_words = blocker.stop_words
        self.coveredBy(blocker.predicates, record_ids, pairs)

class RecordLinkCoverage(Coverage) :
    def __init__(self, predicate_set, pairs) :

        data_1 = set([])
        data_2 = set([])

        for record_1, record_2 in pairs :
            data_1.add(record_1)
            data_2.add(record_2)

        i = 0
        id_records_1 = {}
        id_records_2 = {}

        for record in data_1 :
            id_records_1[i] = record
            i += 1

        for record in data_2 :
            id_records_2[i] = record
            i += 1

        id_records = id_records_1.copy()
        id_records.update(id_records_2)

        record_ids = dict((record, record_id)
                          for record_id, record
                          in id_records.items())

        blocker = RecordLinkBlocker(predicate_set)

        for field in blocker.tfidf_fields :
            fields_1 = ((record_id, record[field])
                        for record_id, record
                        in id_records_1.items())
            fields_2 = [(record_id, record[field])
                        for record_id, record
                        in id_records_2.items()]

            stop_words = stopWords(fields_2)
            blocker.stop_words[field].update(stop_words)

            blocker.tfIdfBlock(fields_1, fields_2, field)

        self.stop_words = blocker.stop_words
        self.coveredBy(blocker.predicates, record_ids, pairs)

def stopWords(data) :
    index = TextIndex(Lexicon(Splitter()))

    for i, (_, doc) in enumerate(data, 1) :
        index.index_doc(i, doc)

    doc_freq = [(len(index.index._wordinfo[wid]), word)
                for word, wid in index.lexicon.items()]

    doc_freq.sort(reverse=True)

    N = float(index.index.documentCount())
    threshold = int(max(1000, N * 0.05))

    stop_words = set([])

    for frequency, word in doc_freq :
        if frequency > threshold :
            stop_words.add(word)
        else :
            break

    return stop_words

class Predicate(object) :
    def __iter__(self) :
        yield self

    def __repr__(self) :
        return "%s: %s" % (self.type, self.__name__)

class SimplePredicate(Predicate) :
    type = "SimplePredicate"

    def __init__(self, func, field) :
        self.func = func
        self.__name__ = "(%s, %s)" % (func.__name__, field)
        self.field = field

    def __call__(self, record_id, record) :
        column = record[self.field]
        return self.func(column)

class TfidfPredicate(Predicate):
    type = "TfidfPredicate"

    def __init__(self, threshold, field):
        self.__name__ = '(%s, %s)' % (threshold, field)
        self.field = field
        self.canopy = {}
        self.threshold = threshold

    def __call__(self, record_id, record) :
        center = self.canopy.get(record_id)
        if center is not None :
            return (unicode(center),)
        else :
            return ()

        return call

    def stringify(self, doc) :
        return doc

    def __getstate__(self):
        result = self.__dict__.copy()
        result['canopy'] = {}
        return result

class TfidfSetPredicate(TfidfPredicate) :
    type = "TfidfPredicate"

    def stringify(self, doc) :
        return ' '.join('_'.join(str(each).split()) for each in doc)

class CompoundPredicate(Predicate) :
    type = "CompoundPredicate"

    def __init__(self, predicates) :
        self.predicates = predicates
        self.__name__ = '(%s)' % ', '.join([str(pred)
                                            for pred in
                                            predicates])

    def __iter__(self) :
        for pred in self.predicates :
            yield pred

    def __call__(self, record_id, record) :
        predicate_keys = (predicate(record_id, record)
                          for predicate in self.predicates)
        return (':'.join(block_key)
                for block_key
                in itertools.product(*predicate_keys))

class CustomStopWordRemover(object):
    def __init__(self, stop_words) :
        self.stop_words = stop_words

    def process(self, lst):
        return [w for w in lst if not w in self.stop_words]

<commit_msg>remove unreachable code<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})

        # clear canopies to reduce memory usage
        for predicate_set in self.tfidf_fields.values() :
            for predicate in predicate_set :
                predicate.canopy = {}

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
            stringify = predicate.stringify

        index_to_id = {}
        base_tokens = {}

        for i, (record_id, doc) in enumerate(data, 1) :
            doc = stringify(doc)
            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
            stringify = predicate.stringify

        index_to_id = {}
        base_tokens = {}

        i = 1

        for record_id, doc in data_1 :
            doc = stringify(doc)
            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  :
            doc = stringify(doc)
            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] =\
                intersection(*blocks)
            i += 1

    predicate_set = coverage.overlap.keys()

    # Only consider predicates that cover at least one duplicate pair
    dupe_coverage = coverage.predicateCoverage(predicate_set,
                                               training_dupes)
    predicate_set = dupe_coverage.keys()

    # Within blocks, we will compare every combination of
    # records. Therefore, we want to avoid predicates that make large
    # blocks.
    for pred in predicate_set[:] :
        blocks = coverage.blocks[pred].itervalues()
        if any(len(block) >= 100 for block in blocks if block) :
            predicate_set.remove(pred)

    # As an efficency, we can throw away the predicates that cover too
    # many distinct pairs
    coverage_threshold = eta * len(training_distinct)
    logger.info("coverage threshold: %s", coverage_threshold)

    distinct_coverage = coverage.predicateCoverage(predicate_set,
                                                   training_distinct)

    for pred, pairs in distinct_coverage.items() :
        if len(pairs) > coverage_threshold :
            predicate_set.remove(pred)

    distinct_coverage = coverage.predicateCoverage(predicate_set,
                                                   training_distinct)

    final_predicate_set = findOptimumBlocking(training_dupes,
                                              predicate_set,
                                              distinct_coverage,
                                              epsilon,
                                              coverage)

    logger.info('Final predicate set:')
    for predicate in final_predicate_set :
        logger.info(predicate)

    if final_predicate_set:
        return final_predicate_set, coverage.stop_words
    else:
        raise ValueError('No predicate found! We could not learn a single good predicate. Maybe give Dedupe more training data')

def findOptimumBlocking(uncovered_dupes,
                        predicate_set,
                        distinct_coverage,
                        epsilon,
                        coverage):

    # Greedily find the predicates that, at each step, covers the
    # most duplicates and covers the least distinct pairs, due to
    # Chvatal, 1979
    #
    # We would like to optimize the ratio of the probability of of a
    # predicate covering a duplicate pair versus the probability of
    # covering a distinct pair. If we have a uniform prior belief
    # about those probabilities, we can estimate these probabilities as
    #
    # (predicate_covered_dupe_pairs + 1) / (all_dupe_pairs + 2)
    #
    # (predicate_covered_distinct_pairs + 1) / (all_distinct_pairs + 2)
    #
    # When we are trying to find the best predicate among a set of
    # predicates, the denominators factor out and our coverage
    # estimator becomes
    #
    # (predicate_covered_dupe_pairs + 1)/ (predicate_covered_distinct_pairs + 1)

    dupe_coverage = coverage.predicateCoverage(predicate_set,
                                               uncovered_dupes)

    uncovered_dupes = set(uncovered_dupes)

    final_predicate_set = []
    while len(uncovered_dupes) > epsilon:

        best_cover = 0
        best_predicate = None
        for predicate in dupe_coverage :
            dupes = len(dupe_coverage[predicate])
            distinct = len(distinct_coverage[predicate])
            cover = (dupes + 1.0)/(distinct + 1.0)
            if cover > best_cover:
                best_cover = cover
                best_predicate = predicate
                best_distinct = distinct
                best_dupes = dupes

        if not best_predicate:
            logger.warning('Ran out of predicates')
            break

        final_predicate_set.append(best_predicate)
        predicate_set.remove(best_predicate)

        uncovered_dupes = uncovered_dupes - dupe_coverage[best_predicate]
        dupe_coverage = coverage.predicateCoverage(predicate_set,
                                                   uncovered_dupes)

        logger.debug(best_predicate)
        logger.debug('cover: %(cover)f, found_dupes: %(found_dupes)d, '
                      'found_distinct: %(found_distinct)d, '
                      'uncovered dupes: %(uncovered)d',
                      {'cover' : best_cover,
                       'found_dupes' : best_dupes,
                       'found_distinct' : best_distinct,
                       'uncovered' : len(uncovered_dupes)
                       })

    return final_predicate_set

class Coverage(object) :

    def coveredBy(self, predicates, record_ids, pairs) :
        self.overlap = defaultdict(set)
        self.blocks = defaultdict(lambda : defaultdict(set))

        for pair in pairs :
            record_1, record_2 = pair
            record_1_id = record_ids[record_1]
            record_2_id = record_ids[record_2]
            for predicate in predicates :
                field_predicate_1 = predicate(record_1_id, record_1)
                if field_predicate_1:
                    field_predicate_2 = predicate(record_2_id, record_2)
                    if field_predicate_2 :
                        field_preds = (set(field_predicate_2)
                                       & set(field_predicate_1))
                        if field_preds :
                            self.overlap[predicate].add(pair)
                            for field_pred in field_preds :
                                self.blocks[predicate][field_pred].add(pair)

        for predicate, coverage in self.blocks.items() :
            for field_pred, block in coverage.items() :
                if len(block) < 100 :
                    del self.blocks[predicate][field_pred]

    def predicateCoverage(self,
                          predicate_set,
                          pairs) :

        coverage = defaultdict(set)
        pairs = set(pairs)

        for predicate in predicate_set :
            covered_pairs = pairs.intersection(self.overlap[predicate])
            if covered_pairs :
                coverage[predicate] = covered_pairs

        return coverage

class DedupeCoverage(Coverage) :
    def __init__(self, predicate_set, pairs) :

        records = set(itertools.chain(*pairs))

        id_records = dict(itertools.izip(itertools.count(), records))
        record_ids = dict(itertools.izip(records, itertools.count()))

        blocker = DedupeBlocker(predicate_set)

        for field in blocker.tfidf_fields :
            data = [(record_id, record[field])
                    for record_id, record
                    in id_records.items()]
            stop_words = stopWords(data)
            blocker.stop_words[field].update(stop_words)
            blocker.tfIdfBlock(data, field)

        self.stop_words = blocker.stop_words
        self.coveredBy(blocker.predicates, record_ids, pairs)

class RecordLinkCoverage(Coverage) :
    def __init__(self, predicate_set, pairs) :

        data_1 = set([])
        data_2 = set([])

        for record_1, record_2 in pairs :
            data_1.add(record_1)
            data_2.add(record_2)

        i = 0
        id_records_1 = {}
        id_records_2 = {}

        for record in data_1 :
            id_records_1[i] = record
            i += 1

        for record in data_2 :
            id_records_2[i] = record
            i += 1

        id_records = id_records_1.copy()
        id_records.update(id_records_2)

        record_ids = dict((record, record_id)
                          for record_id, record
                          in id_records.items())

        blocker = RecordLinkBlocker(predicate_set)

        for field in blocker.tfidf_fields :
            fields_1 = ((record_id, record[field])
                        for record_id, record
                        in id_records_1.items())
            fields_2 = [(record_id, record[field])
                        for record_id, record
                        in id_records_2.items()]

            stop_words = stopWords(fields_2)
            blocker.stop_words[field].update(stop_words)

            blocker.tfIdfBlock(fields_1, fields_2, field)

        self.stop_words = blocker.stop_words
        self.coveredBy(blocker.predicates, record_ids, pairs)

def stopWords(data) :
    index = TextIndex(Lexicon(Splitter()))

    for i, (_, doc) in enumerate(data, 1) :
        index.index_doc(i, doc)

    doc_freq = [(len(index.index._wordinfo[wid]), word)
                for word, wid in index.lexicon.items()]

    doc_freq.sort(reverse=True)

    N = float(index.index.documentCount())
    threshold = int(max(1000, N * 0.05))

    stop_words = set([])

    for frequency, word in doc_freq :
        if frequency > threshold :
            stop_words.add(word)
        else :
            break

    return stop_words

class Predicate(object) :
    def __iter__(self) :
        yield self

    def __repr__(self) :
        return "%s: %s" % (self.type, self.__name__)

class SimplePredicate(Predicate) :
    type = "SimplePredicate"

    def __init__(self, func, field) :
        self.func = func
        self.__name__ = "(%s, %s)" % (func.__name__, field)
        self.field = field

    def __call__(self, record_id, record) :
        column = record[self.field]
        return self.func(column)

class TfidfPredicate(Predicate):
    type = "TfidfPredicate"

    def __init__(self, threshold, field):
        self.__name__ = '(%s, %s)' % (threshold, field)
        self.field = field
        self.canopy = {}
        self.threshold = threshold

    def __call__(self, record_id, record) :
        center = self.canopy.get(record_id)
        if center is not None :
            return (unicode(center),)
        else :
            return ()

    def stringify(self, doc) :
        return doc

    def __getstate__(self):
        result = self.__dict__.copy()
        result['canopy'] = {}
        return result

class TfidfSetPredicate(TfidfPredicate) :
    type = "TfidfPredicate"

    def stringify(self, doc) :
        return ' '.join('_'.join(str(each).split()) for each in doc)

class CompoundPredicate(Predicate) :
    type = "CompoundPredicate"

    def __init__(self, predicates) :
        self.predicates = predicates
        self.__name__ = '(%s)' % ', '.join([str(pred)
                                            for pred in
                                            predicates])

    def __iter__(self) :
        for pred in self.predicates :
            yield pred

    def __call__(self, record_id, record) :
        predicate_keys = (predicate(record_id, record)
                          for predicate in self.predicates)
        return (':'.join(block_key)
                for block_key
                in itertools.product(*predicate_keys))

class CustomStopWordRemover(object):
    def __init__(self, stop_words) :
        self.stop_words = stop_words

    def process(self, lst):
        return [w for w in lst if not w in self.stop_words]

<|endoftext|>


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

# provides functions for selecting a sample of training data

import itertools
import logging
import collections
import functools
import random
from abc import ABC, abstractmethod
import math

from typing import (Dict, Sequence, Iterable, Tuple, List,
                    FrozenSet)

from . import blocking, core
from .predicates import CompoundPredicate, Predicate

logger = logging.getLogger(__name__)

Cover = Dict[Predicate, FrozenSet[int]]

class BlockLearner(ABC):
    def learn(self, matches, recall):
        '''
        Takes in a set of training pairs and predicates and tries to find
        a good set of blocking rules.
        '''
        comparison_cover = self.comparison_cover  # type: ignore
        match_cover = self.cover(matches)  # type: ignore

        for key in list(match_cover.keys() - comparison_cover.keys()):
            del match_cover[key]

        coverable_dupes = frozenset.union(*match_cover.values())
        uncoverable_dupes = [pair for i, pair in enumerate(matches)
                             if i not in coverable_dupes]

        epsilon = int((1.0 - recall) * len(matches))

        if len(uncoverable_dupes) > epsilon:
            logger.warning(OUT_OF_PREDICATES_WARNING)
            logger.debug(uncoverable_dupes)
            epsilon = 0
        else:
            epsilon -= len(uncoverable_dupes)

        candidate_cover = self.generate_candidates_rf(match_cover,
                                                      comparison_cover)

        searcher = BranchBound(len(coverable_dupes) - epsilon, 2500)
        final_predicates = searcher.search(candidate_cover)

        logger.info('Final predicate set:')
        for predicate in final_predicates:
            logger.info(predicate)

        return final_predicates

    def generate_candidates_bilenko(self,
                                    match_cover: dict,
                                    comparison_cover: dict) -> dict:
        predicates = list(match_cover)
        candidates = {}
        K = 3

        for i, predicate in enumerate(predicates):
            current_match_cover = match_cover[predicate]
            current_comparison_cover = comparison_cover[predicate]
            predicate.count = self.estimate(current_comparison_cover)
            candidates[predicate] = current_match_cover
            remaining = predicates
            predicate = CompoundPredicate(predicate,)
            for _ in range(K):
                if not remaining:
                    break
                best_p = max(remaining,
                             key=lambda x: (len(current_match_cover & match_cover[x]) /
                                            (self.estimate(current_comparison_cover & comparison_cover[x]) or float('inf'))))
                predicate += best_p
                current_match_cover &= match_cover[best_p]
                current_comparison_cover &= comparison_cover[best_p]
                predicate.count = self.estimate(current_comparison_cover)  # type: ignore
                candidates[predicate] = current_match_cover
                remaining.remove(best_p)

        return candidates

    def generate_candidates_rf(self,
                               match_cover: Cover,
                               comparison_cover: Cover) -> Cover:
        predicates = list(match_cover)
        matches = list(frozenset.union(*match_cover.values()))
        pred_sample_size = max(int(math.sqrt(len(predicates))), 5)
        candidates = {}
        K = 3

        n_samples = 5000
        for _ in range(n_samples):
            sample_predicates = random.sample(predicates,
                                              pred_sample_size)
            resampler = Resampler(matches)
            sample_match_cover = {pred: resampler(pairs)
                                  for pred, pairs
                                  in match_cover.items()}

            # initialize variables that will be
            # the base for the constructing k-conjunctions
            candidate = None
            covered_comparisons = InfiniteSet()
            covered_matches = InfiniteSet()
            covered_sample_matches = InfiniteSet()

            def score(predicate: Predicate) -> float:
                try:
                    return (len(covered_sample_matches &
                                sample_match_cover[predicate]) /
                            self.estimate(covered_comparisons &
                                          comparison_cover[predicate]))
                except ZeroDivisionError:
                    return 0.

            for _ in range(K):
                next_predicate = max(sample_predicates, key=score)
                if candidate:
                    candidate += next_predicate
                else:
                    candidate = next_predicate

                covered_comparisons &= comparison_cover[next_predicate]
                candidate.count = self.estimate(covered_comparisons)  # type: ignore

                covered_matches &= match_cover[next_predicate]
                candidates[candidate] = covered_matches

                covered_sample_matches &= sample_match_cover[next_predicate]

                sample_predicates.remove(next_predicate)

        return candidates  # type: ignore

    def cover(self, pairs):
        predicate_cover = {}
        for predicate in self.blocker.predicates:  # type: ignore
            coverage = frozenset(
                i for i, (record_1, record_2)
                in enumerate(pairs)
                if (set(predicate(record_1)) &
                    set(predicate(record_2, target=True))))
            if coverage:
                predicate_cover[predicate] = coverage

        return predicate_cover

    @abstractmethod
    def estimate(self, comparisons) -> float:
        ...

class DedupeBlockLearner(BlockLearner):

    def __init__(self, predicates, sampled_records, data):

        N = sampled_records.original_length
        N_s = len(sampled_records)

        self.r = (N * (N - 1)) / (N_s * (N_s - 1))

        self.blocker = blocking.Fingerprinter(predicates)
        self.blocker.index_all(data)

        self.comparison_cover = self.coveredPairs(self.blocker, sampled_records)

    @staticmethod
    def coveredPairs(blocker, records):
        cover = {}

        pair_enumerator = core.Enumerator()
        n_records = len(records)

        for predicate in blocker.predicates:
            pred_cover = collections.defaultdict(set)

            for id, record in records.items():
                blocks = predicate(record)
                for block in blocks:
                    pred_cover[block].add(id)

            if not pred_cover:
                continue

            max_cover = max(len(v) for v in pred_cover.values())
            if max_cover == n_records:
                continue

            pairs = frozenset(
                pair_enumerator[pair]
                for block in pred_cover.values()
                for pair in itertools.combinations(sorted(block), 2))
            cover[predicate] = pairs

        return cover

    def estimate(self, comparisons):
        # Result due to Stefano Allesina and Jacopo Grilli,
        # details forthcoming
        #
        # This estimates the total number of comparisons a blocking
        # rule will produce.

        return self.r * len(comparisons)

class RecordLinkBlockLearner(BlockLearner):

    def __init__(self, predicates, sampled_records_1, sampled_records_2, data_2):

        r_a = ((sampled_records_1.original_length) /
               len(sampled_records_1))
        r_b = ((sampled_records_2.original_length) /
               len(sampled_records_2))

        self.r = r_a * r_b

        self.blocker = blocking.Fingerprinter(predicates)
        self.blocker.index_all(data_2)

        self.comparison_cover = self.coveredPairs(self.blocker,
                                                  sampled_records_1,
                                                  sampled_records_2)

    def coveredPairs(self, blocker, records_1, records_2):
        cover = {}

        pair_enumerator = core.Enumerator()

        for predicate in blocker.predicates:
            cover[predicate] = collections.defaultdict(lambda: (set(), set()))
            for id, record in records_2.items():
                blocks = predicate(record, target=True)
                for block in blocks:
                    cover[predicate][block][1].add(id)

            current_blocks = set(cover[predicate])
            for id, record in records_1.items():
                blocks = set(predicate(record))
                for block in blocks & current_blocks:
                    cover[predicate][block][0].add(id)

        for predicate, blocks in cover.items():
            pairs = frozenset(
                pair_enumerator[pair]
                for A, B in blocks.values()
                for pair in itertools.product(A, B))
            cover[predicate] = pairs

        return cover

    def estimate(self, comparisons):
        # https://stats.stackexchange.com/a/465060/82

        return self.r * len(comparisons)

class BranchBound(object):
    def __init__(self, target, max_calls):
        self.calls: int = max_calls
        self.target: int = target
        self.cheapest_score = float('inf')
        self.original_cover: Cover
        self.cheapest: Tuple[Predicate, ...]

    def search(self,
               candidates: Cover,
               partial: Tuple[Predicate, ...] = ()) -> Tuple[Predicate, ...]:
        if self.calls <= 0:
            return self.cheapest

        if self.cheapest_score == float('inf'):
            self.original_cover = candidates.copy()
            self.cheapest = tuple(candidates)
            self.cheapest_score = self.score(candidates)

        self.calls -= 1

        covered = self.covered(partial)
        score = self.score(partial)

        if covered >= self.target:
            if score < self.cheapest_score:
                self.cheapest = partial
                self.cheapest_score = score

        else:
            window = self.cheapest_score - score

            candidates = {p: cover
                          for p, cover in candidates.items()
                          if p.count < window}  # type: ignore

            reachable = self.reachable(candidates) + covered

            if candidates and reachable >= self.target:

                order_by = functools.partial(self.order_by, candidates)

                best = max(candidates, key=order_by)

                remaining = self.uncovered_by(candidates,
                                              candidates[best])
                self.search(remaining, partial + (best,))
                del remaining

                reduced = self.remove_dominated(candidates, best)
                self.search(reduced, partial)
                del reduced

        return self.cheapest

    @staticmethod
    def order_by(candidates: Cover, p: Predicate) -> Tuple[int, float]:
        return (len(candidates[p]), -p.count)  # type: ignore

    @staticmethod
    def score(partial: Iterable[Predicate]) -> float:
        return sum(p.count for p in partial)  # type: ignore

    def covered(self, partial: Tuple[Predicate, ...]) -> int:
        if partial:
            return len(frozenset.union(*(self.original_cover[p]
                                         for p in partial)))
        else:
            return 0

    @staticmethod
    def reachable(dupe_cover: Cover) -> int:
        if dupe_cover:
            return len(frozenset.union(*dupe_cover.values()))
        else:
            return 0

    @staticmethod
    def remove_dominated(coverage: Cover, dominator: Predicate) -> Cover:
        dominant_cover = coverage[dominator]

        for pred, cover in coverage.copy().items():
            if (dominator.count <= pred.count and  # type: ignore
                    dominant_cover >= cover):
                del coverage[pred]

        return coverage

    @staticmethod
    def uncovered_by(coverage: Cover, covered: frozenset) -> Cover:
        remaining = {}
        for predicate, uncovered in coverage.items():
            still_uncovered = uncovered - covered
            if still_uncovered:
                remaining[predicate] = still_uncovered

        return remaining

class InfiniteSet(object):

    def __and__(self, item):
        return item

    def __rand__(self, item):
        return item

class Resampler(object):

    def __init__(self, sequence: Sequence):

        sampled = random.choices(sequence, k=len(sequence))

        c = collections.Counter(sampled)
        max_value = len(sequence) + 1

        self.replacements: Dict[int, List[int]] = {}
        for k, v in c.items():
            self.replacements[k] = [v]
            if v > 1:
                for _ in range(v - 1):
                    self.replacements[k].append(max_value)
                    max_value += 1

    def __call__(self, iterable: Iterable) -> frozenset:

        result = itertools.chain.from_iterable(self.replacements[k]
                                               for k in iterable
                                               if k in self.replacements)
        return frozenset(result)

OUT_OF_PREDICATES_WARNING = "Ran out of predicates: Dedupe tries to find blocking rules that will work well with your data. Sometimes it can't find great ones, and you'll get this warning. It means that there are some pairs of true records that dedupe may never compare. If you are getting bad results, try increasing the `max_comparison` argument to the train method"  # noqa: E501
<commit_msg>remove bilenko mode candidate_gen<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-

# provides functions for selecting a sample of training data

import itertools
import logging
import collections
import functools
import random
from abc import ABC, abstractmethod
import math

from typing import (Dict, Sequence, Iterable, Tuple, List,
                    FrozenSet)

from . import blocking, core
from .predicates import CompoundPredicate, Predicate

logger = logging.getLogger(__name__)

Cover = Dict[Predicate, FrozenSet[int]]

class BlockLearner(ABC):
    def learn(self, matches, recall):
        '''
        Takes in a set of training pairs and predicates and tries to find
        a good set of blocking rules.
        '''
        comparison_cover = self.comparison_cover  # type: ignore
        match_cover = self.cover(matches)  # type: ignore

        for key in list(match_cover.keys() - comparison_cover.keys()):
            del match_cover[key]

        coverable_dupes = frozenset.union(*match_cover.values())
        uncoverable_dupes = [pair for i, pair in enumerate(matches)
                             if i not in coverable_dupes]

        epsilon = int((1.0 - recall) * len(matches))

        if len(uncoverable_dupes) > epsilon:
            logger.warning(OUT_OF_PREDICATES_WARNING)
            logger.debug(uncoverable_dupes)
            epsilon = 0
        else:
            epsilon -= len(uncoverable_dupes)

        candidate_cover = self.generate_candidates(match_cover,
                                                   comparison_cover)

        searcher = BranchBound(len(coverable_dupes) - epsilon, 2500)
        final_predicates = searcher.search(candidate_cover)

        logger.info('Final predicate set:')
        for predicate in final_predicates:
            logger.info(predicate)

        return final_predicates

    def generate_candidates(self,
                            match_cover: Cover,
                            comparison_cover: Cover) -> Cover:
        predicates = list(match_cover)
        matches = list(frozenset.union(*match_cover.values()))
        pred_sample_size = max(int(math.sqrt(len(predicates))), 5)
        candidates = {}
        K = 3

        n_samples = 5000
        for _ in range(n_samples):
            sample_predicates = random.sample(predicates,
                                              pred_sample_size)
            resampler = Resampler(matches)
            sample_match_cover = {pred: resampler(pairs)
                                  for pred, pairs
                                  in match_cover.items()}

            # initialize variables that will be
            # the base for the constructing k-conjunctions
            candidate = None
            covered_comparisons = InfiniteSet()
            covered_matches = InfiniteSet()
            covered_sample_matches = InfiniteSet()

            def score(predicate: Predicate) -> float:
                try:
                    return (len(covered_sample_matches &
                                sample_match_cover[predicate]) /
                            self.estimate(covered_comparisons &
                                          comparison_cover[predicate]))
                except ZeroDivisionError:
                    return 0.

            for _ in range(K):
                next_predicate = max(sample_predicates, key=score)
                if candidate:
                    candidate += next_predicate
                else:
                    candidate = next_predicate

                covered_comparisons &= comparison_cover[next_predicate]
                candidate.count = self.estimate(covered_comparisons)  # type: ignore

                covered_matches &= match_cover[next_predicate]
                candidates[candidate] = covered_matches

                covered_sample_matches &= sample_match_cover[next_predicate]

                sample_predicates.remove(next_predicate)

        return candidates  # type: ignore

    def cover(self, pairs):
        predicate_cover = {}
        for predicate in self.blocker.predicates:  # type: ignore
            coverage = frozenset(
                i for i, (record_1, record_2)
                in enumerate(pairs)
                if (set(predicate(record_1)) &
                    set(predicate(record_2, target=True))))
            if coverage:
                predicate_cover[predicate] = coverage

        return predicate_cover

    @abstractmethod
    def estimate(self, comparisons) -> float:
        ...

class DedupeBlockLearner(BlockLearner):

    def __init__(self, predicates, sampled_records, data):

        N = sampled_records.original_length
        N_s = len(sampled_records)

        self.r = (N * (N - 1)) / (N_s * (N_s - 1))

        self.blocker = blocking.Fingerprinter(predicates)
        self.blocker.index_all(data)

        self.comparison_cover = self.coveredPairs(self.blocker, sampled_records)

    @staticmethod
    def coveredPairs(blocker, records):
        cover = {}

        pair_enumerator = core.Enumerator()
        n_records = len(records)

        for predicate in blocker.predicates:
            pred_cover = collections.defaultdict(set)

            for id, record in records.items():
                blocks = predicate(record)
                for block in blocks:
                    pred_cover[block].add(id)

            if not pred_cover:
                continue

            max_cover = max(len(v) for v in pred_cover.values())
            if max_cover == n_records:
                continue

            pairs = frozenset(
                pair_enumerator[pair]
                for block in pred_cover.values()
                for pair in itertools.combinations(sorted(block), 2))
            cover[predicate] = pairs

        return cover

    def estimate(self, comparisons):
        # Result due to Stefano Allesina and Jacopo Grilli,
        # details forthcoming
        #
        # This estimates the total number of comparisons a blocking
        # rule will produce.

        return self.r * len(comparisons)

class RecordLinkBlockLearner(BlockLearner):

    def __init__(self, predicates, sampled_records_1, sampled_records_2, data_2):

        r_a = ((sampled_records_1.original_length) /
               len(sampled_records_1))
        r_b = ((sampled_records_2.original_length) /
               len(sampled_records_2))

        self.r = r_a * r_b

        self.blocker = blocking.Fingerprinter(predicates)
        self.blocker.index_all(data_2)

        self.comparison_cover = self.coveredPairs(self.blocker,
                                                  sampled_records_1,
                                                  sampled_records_2)

    def coveredPairs(self, blocker, records_1, records_2):
        cover = {}

        pair_enumerator = core.Enumerator()

        for predicate in blocker.predicates:
            cover[predicate] = collections.defaultdict(lambda: (set(), set()))
            for id, record in records_2.items():
                blocks = predicate(record, target=True)
                for block in blocks:
                    cover[predicate][block][1].add(id)

            current_blocks = set(cover[predicate])
            for id, record in records_1.items():
                blocks = set(predicate(record))
                for block in blocks & current_blocks:
                    cover[predicate][block][0].add(id)

        for predicate, blocks in cover.items():
            pairs = frozenset(
                pair_enumerator[pair]
                for A, B in blocks.values()
                for pair in itertools.product(A, B))
            cover[predicate] = pairs

        return cover

    def estimate(self, comparisons):
        # https://stats.stackexchange.com/a/465060/82

        return self.r * len(comparisons)

class BranchBound(object):
    def __init__(self, target, max_calls):
        self.calls: int = max_calls
        self.target: int = target
        self.cheapest_score = float('inf')
        self.original_cover: Cover
        self.cheapest: Tuple[Predicate, ...]

    def search(self,
               candidates: Cover,
               partial: Tuple[Predicate, ...] = ()) -> Tuple[Predicate, ...]:
        if self.calls <= 0:
            return self.cheapest

        if self.cheapest_score == float('inf'):
            self.original_cover = candidates.copy()
            self.cheapest = tuple(candidates)
            self.cheapest_score = self.score(candidates)

        self.calls -= 1

        covered = self.covered(partial)
        score = self.score(partial)

        if covered >= self.target:
            if score < self.cheapest_score:
                self.cheapest = partial
                self.cheapest_score = score

        else:
            window = self.cheapest_score - score

            candidates = {p: cover
                          for p, cover in candidates.items()
                          if p.count < window}  # type: ignore

            reachable = self.reachable(candidates) + covered

            if candidates and reachable >= self.target:

                order_by = functools.partial(self.order_by, candidates)

                best = max(candidates, key=order_by)

                remaining = self.uncovered_by(candidates,
                                              candidates[best])
                self.search(remaining, partial + (best,))
                del remaining

                reduced = self.remove_dominated(candidates, best)
                self.search(reduced, partial)
                del reduced

        return self.cheapest

    @staticmethod
    def order_by(candidates: Cover, p: Predicate) -> Tuple[int, float]:
        return (len(candidates[p]), -p.count)  # type: ignore

    @staticmethod
    def score(partial: Iterable[Predicate]) -> float:
        return sum(p.count for p in partial)  # type: ignore

    def covered(self, partial: Tuple[Predicate, ...]) -> int:
        if partial:
            return len(frozenset.union(*(self.original_cover[p]
                                         for p in partial)))
        else:
            return 0

    @staticmethod
    def reachable(dupe_cover: Cover) -> int:
        if dupe_cover:
            return len(frozenset.union(*dupe_cover.values()))
        else:
            return 0

    @staticmethod
    def remove_dominated(coverage: Cover, dominator: Predicate) -> Cover:
        dominant_cover = coverage[dominator]

        for pred, cover in coverage.copy().items():
            if (dominator.count <= pred.count and  # type: ignore
                    dominant_cover >= cover):
                del coverage[pred]

        return coverage

    @staticmethod
    def uncovered_by(coverage: Cover, covered: frozenset) -> Cover:
        remaining = {}
        for predicate, uncovered in coverage.items():
            still_uncovered = uncovered - covered
            if still_uncovered:
                remaining[predicate] = still_uncovered

        return remaining

class InfiniteSet(object):

    def __and__(self, item):
        return item

    def __rand__(self, item):
        return item

class Resampler(object):

    def __init__(self, sequence: Sequence):

        sampled = random.choices(sequence, k=len(sequence))

        c = collections.Counter(sampled)
        max_value = len(sequence) + 1

        self.replacements: Dict[int, List[int]] = {}
        for k, v in c.items():
            self.replacements[k] = [v]
            if v > 1:
                for _ in range(v - 1):
                    self.replacements[k].append(max_value)
                    max_value += 1

    def __call__(self, iterable: Iterable) -> frozenset:

        result = itertools.chain.from_iterable(self.replacements[k]
                                               for k in iterable
                                               if k in self.replacements)
        return frozenset(result)

OUT_OF_PREDICATES_WARNING = "Ran out of predicates: Dedupe tries to find blocking rules that will work well with your data. Sometimes it can't find great ones, and you'll get this warning. It means that there are some pairs of true records that dedupe may never compare. If you are getting bad results, try increasing the `max_comparison` argument to the train method"  # noqa: E501
<|endoftext|>


<commit_before>"""Asynchronous Advantage Actor-Critic (A3C) algorithm for reinforcement learning."""

from deepchem.models import TensorGraph
from deepchem.models.tensorgraph import TFWrapper
from deepchem.models.tensorgraph.layers import Feature, Weights, Label, Layer
import numpy as np
import tensorflow as tf
import copy
import multiprocessing
import os
import re
import threading

class A3CLoss(Layer):
  """This layer computes the loss function for A3C."""

  def __init__(self, value_weight, entropy_weight, **kwargs):
    super(A3CLoss, self).__init__(**kwargs)
    self.value_weight = value_weight
    self.entropy_weight = entropy_weight

  def create_tensor(self, **kwargs):
    reward, action, prob, value = [layer.out_tensor for layer in self.in_layers]
    prob = prob + np.finfo(np.float32).eps
    log_prob = tf.log(prob)
    policy_loss = -tf.reduce_sum(
        (reward - value) * tf.reduce_sum(action * log_prob))
    value_loss = tf.reduce_sum(tf.square(reward - value))
    entropy = -tf.reduce_sum(prob * log_prob)
    self.out_tensor = policy_loss + self.value_weight * value_loss - self.entropy_weight * entropy
    return self.out_tensor

def _create_feed_dict(features, state):
  return dict((f.out_tensor, np.expand_dims(s, axis=0))
              for f, s in zip(features, state))

class A3C(object):
  """
  Implements the Asynchronous Advantage Actor-Critic (A3C) algorithm for reinforcement learning.

  This algorithm requires the policy to output two quantities: a vector giving the probability of
  taking each action, and an estimate of the value function for the current state.  It optimizes
  both outputs at once using a loss that is the sum of three terms:

  1. The policy loss, which seeks to maximize the discounted reward for each action.
  2. The value loss, which tries to make the value estimate match the actual discounted reward
     that was attained at each step.
  3. An entropy term to encourage exploration.

  This class only supports environments with discrete action spaces, not continuous ones.  The
  "action" argument passed to the environment is an integer, giving the index of the action to perform.
  """

  def __init__(self,
               env,
               policy,
               max_rollout_length=20,
               discount_factor=0.99,
               value_weight=1.0,
               entropy_weight=0.01,
               model_dir=None):
    """Create an object for optimizing a policy.

    Parameters
    ----------
    env: Environment
      the Environment to interact with
    policy: Policy
      the Policy to optimize.  Its create_layers() method must return a map containing the
      keys 'action_prob' and 'value', corresponding to the action probabilities and value estimate
    max_rollout_length: int
      the maximum length of rollouts to generate
    discount_factor: float
      the discount factor to use when computing rewards
    value_weight: float
      a scale factor for the value loss term in the loss function
    entropy_weight: float
      a scale factor for the entropy term in the loss function
    model_dir: str
      the directory in which the model will be saved.  If None, a temporary directory will be created.
    """
    self._env = env
    self._policy = policy
    self.max_rollout_length = max_rollout_length
    self.discount_factor = discount_factor
    self.value_weight = value_weight
    self.entropy_weight = entropy_weight
    self.optimizer = TFWrapper(
        tf.train.AdamOptimizer, learning_rate=0.001, beta1=0.9, beta2=0.999)
    (self._graph, self._features, rewards, actions, self._action_prob,
     self._value) = self._build_graph(None, 'global', model_dir)
    with self._graph._get_tf("Graph").as_default():
      self._session = tf.Session()

  def _build_graph(self, tf_graph, scope, model_dir):
    """Construct a TensorGraph containing the policy and loss calculations."""
    features = [Feature(shape=[None] + list(s)) for s in self._env.state_shape]
    policy_layers = self._policy.create_layers(features)
    action_prob = policy_layers['action_prob']
    value = policy_layers['value']
    rewards = Weights(shape=(None, 1))
    actions = Label(shape=(None, self._env.n_actions))
    loss = A3CLoss(
        self.value_weight,
        self.entropy_weight,
        in_layers=[rewards, actions, action_prob, value])
    graph = TensorGraph(
        batch_size=self.max_rollout_length,
        use_queue=False,
        graph=tf_graph,
        model_dir=model_dir)
    for f in features:
      graph._add_layer(f)
    graph.add_output(action_prob)
    graph.add_output(value)
    graph.set_loss(loss)
    graph.set_optimizer(self.optimizer)
    with graph._get_tf("Graph").as_default():
      with tf.variable_scope(scope):
        graph.build()
    return graph, features, rewards, actions, action_prob, value

  def fit(self, total_steps, max_checkpoints_to_keep=5,
          checkpoint_interval=600):
    """Train the policy.

    Parameters
    ----------
    total_steps: int
      the total number of time steps to perform on the environment, across all rollouts
      on all threads
    max_checkpoints_to_keep: int
      the maximum number of checkpoint files to keep.  When this number is reached, older
      files are deleted.
    checkpoint_interval: float
      the time interval at which to save checkpoints, measured in seconds
    """
    with self._graph._get_tf("Graph").as_default():
      step_count = [0]
      workers = []
      threads = []
      for i in range(multiprocessing.cpu_count()):
        workers.append(_Worker(self, i))
      self._session.run(tf.global_variables_initializer())
      for worker in workers:
        thread = threading.Thread(
            name=worker.scope,
            target=lambda: worker.run(step_count, total_steps))
        threads.append(thread)
        thread.start()
      saver = tf.train.Saver(max_to_keep=max_checkpoints_to_keep)
      checkpoint_index = 0
      while True:
        threads = [t for t in threads if t.isAlive()]
        if len(threads) > 0:
          threads[0].join(checkpoint_interval)
        checkpoint_index += 1
        saver.save(
            self._session, self._graph.save_file, global_step=checkpoint_index)
        if len(threads) == 0:
          break

  def predict(self, state):
    """Compute the policy's output predictions for a state.

    Parameters
    ----------
    state: array
      the state of the environment for which to generate predictions

    Returns
    -------
    the array of action probabilities, and the estimated value function
    """
    with self._graph._get_tf("Graph").as_default():
      feed_dict = _create_feed_dict(self._features, state)
      return self._session.run(
          [self._action_prob.out_tensor, self._value.out_tensor],
          feed_dict=feed_dict)

  def select_action(self, state, deterministic=False):
    """Select an action to perform based on the environment's state.

    Parameters
    ----------
    state: array
      the state of the environment for which to select an action
    deterministic: bool
      if True, always return the best action (that is, the one with highest probability).
      If False, randomly select an action based on the computed probabilities.

    Returns
    -------
    the index of the selected action
    """
    with self._graph._get_tf("Graph").as_default():
      feed_dict = _create_feed_dict(self._features, state)
      probabilities = self._session.run(
          self._action_prob.out_tensor, feed_dict=feed_dict)
      if deterministic:
        return probabilities.argmax()
      else:
        return np.random.choice(
            np.arange(self._env.n_actions), p=probabilities[0])

  def restore(self):
    """Reload the model parameters from the most recent checkpoint file."""
    last_checkpoint = tf.train.latest_checkpoint(self._graph.model_dir)
    if last_checkpoint is None:
      raise ValueError('No checkpoint found')
    with self._graph._get_tf("Graph").as_default():
      saver = tf.train.Saver()
      saver.restore(self._session, last_checkpoint)

class _Worker(object):
  """A Worker object is created for each training thread."""

  def __init__(self, a3c, index):
    self.a3c = a3c
    self.index = index
    self.scope = 'worker%d' % index
    self.env = copy.deepcopy(a3c._env)
    self.env.reset()
    self.graph, self.features, self.rewards, self.actions, self.action_prob, self.value = a3c._build_graph(
        a3c._graph._get_tf('Graph'), self.scope, None)
    with a3c._graph._get_tf("Graph").as_default():
      local_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
                                     self.scope)
      global_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
                                      'global')
      gradients = tf.gradients(self.graph.loss.out_tensor, local_vars)
      grads_and_vars = list(zip(gradients, global_vars))
      self.train_op = a3c._graph._get_tf('Optimizer').apply_gradients(
          grads_and_vars)
      self.update_local_variables = tf.group(
          * [tf.assign(v1, v2) for v1, v2 in zip(local_vars, global_vars)])

  def run(self, step_count, total_steps):
    with self.graph._get_tf("Graph").as_default():
      session = self.a3c._session
      while step_count[0] < total_steps:
        session.run(self.update_local_variables)
        episode_states, episode_actions, episode_rewards = self.create_rollout()
        feed_dict = {}
        for f, s in zip(self.features, episode_states):
          feed_dict[f.out_tensor] = s
        feed_dict[self.rewards.out_tensor] = episode_rewards
        feed_dict[self.actions.out_tensor] = episode_actions
        session.run(self.train_op, feed_dict=feed_dict)
        step_count[0] += len(episode_actions)

  def create_rollout(self):
    """Generate a rollout."""
    n_actions = self.env.n_actions
    session = self.a3c._session
    states = [[] for i in range(len(self.features))]
    actions = []
    rewards = []
    for i in range(self.a3c.max_rollout_length):
      if self.env.terminated:
        break
      state = self.env.state
      for j in range(len(state)):
        states[j].append(state[j])
      feed_dict = _create_feed_dict(self.features, state)
      probabilities = session.run(
          self.action_prob.out_tensor, feed_dict=feed_dict)
      action = np.random.choice(np.arange(n_actions), p=probabilities[0])
      actions.append(np.zeros(n_actions))
      actions[i][action] = 1.0
      rewards.append(self.env.step(action))
    if not self.env.terminated:
      # Add an estimate of the reward for the rest of the episode.
      feed_dict = _create_feed_dict(self.features, self.env.state)
      rewards[-1] += self.a3c.discount_factor * session.run(
          self.value.out_tensor, feed_dict)
    for j in range(len(rewards) - 1, 0, -1):
      rewards[j - 1] += self.a3c.discount_factor * rewards[j]
    if self.env.terminated:
      self.env.reset()
    return np.array(states), np.array(actions), np.array(rewards).reshape(
        (len(rewards), 1))
<commit_msg>Improved A3C checkpointing<commit_after>"""Asynchronous Advantage Actor-Critic (A3C) algorithm for reinforcement learning."""

from deepchem.models import TensorGraph
from deepchem.models.tensorgraph import TFWrapper
from deepchem.models.tensorgraph.layers import Feature, Weights, Label, Layer
import numpy as np
import tensorflow as tf
import copy
import multiprocessing
import os
import re
import threading

class A3CLoss(Layer):
  """This layer computes the loss function for A3C."""

  def __init__(self, value_weight, entropy_weight, **kwargs):
    super(A3CLoss, self).__init__(**kwargs)
    self.value_weight = value_weight
    self.entropy_weight = entropy_weight

  def create_tensor(self, **kwargs):
    reward, action, prob, value = [layer.out_tensor for layer in self.in_layers]
    prob = prob + np.finfo(np.float32).eps
    log_prob = tf.log(prob)
    policy_loss = -tf.reduce_sum(
        (reward - value) * tf.reduce_sum(action * log_prob))
    value_loss = tf.reduce_sum(tf.square(reward - value))
    entropy = -tf.reduce_sum(prob * log_prob)
    self.out_tensor = policy_loss + self.value_weight * value_loss - self.entropy_weight * entropy
    return self.out_tensor

def _create_feed_dict(features, state):
  return dict((f.out_tensor, np.expand_dims(s, axis=0))
              for f, s in zip(features, state))

class A3C(object):
  """
  Implements the Asynchronous Advantage Actor-Critic (A3C) algorithm for reinforcement learning.

  This algorithm requires the policy to output two quantities: a vector giving the probability of
  taking each action, and an estimate of the value function for the current state.  It optimizes
  both outputs at once using a loss that is the sum of three terms:

  1. The policy loss, which seeks to maximize the discounted reward for each action.
  2. The value loss, which tries to make the value estimate match the actual discounted reward
     that was attained at each step.
  3. An entropy term to encourage exploration.

  This class only supports environments with discrete action spaces, not continuous ones.  The
  "action" argument passed to the environment is an integer, giving the index of the action to perform.
  """

  def __init__(self,
               env,
               policy,
               max_rollout_length=20,
               discount_factor=0.99,
               value_weight=1.0,
               entropy_weight=0.01,
               model_dir=None):
    """Create an object for optimizing a policy.

    Parameters
    ----------
    env: Environment
      the Environment to interact with
    policy: Policy
      the Policy to optimize.  Its create_layers() method must return a map containing the
      keys 'action_prob' and 'value', corresponding to the action probabilities and value estimate
    max_rollout_length: int
      the maximum length of rollouts to generate
    discount_factor: float
      the discount factor to use when computing rewards
    value_weight: float
      a scale factor for the value loss term in the loss function
    entropy_weight: float
      a scale factor for the entropy term in the loss function
    model_dir: str
      the directory in which the model will be saved.  If None, a temporary directory will be created.
    """
    self._env = env
    self._policy = policy
    self.max_rollout_length = max_rollout_length
    self.discount_factor = discount_factor
    self.value_weight = value_weight
    self.entropy_weight = entropy_weight
    self.optimizer = TFWrapper(
        tf.train.AdamOptimizer, learning_rate=0.001, beta1=0.9, beta2=0.999)
    (self._graph, self._features, rewards, actions, self._action_prob,
     self._value) = self._build_graph(None, 'global', model_dir)
    with self._graph._get_tf("Graph").as_default():
      self._session = tf.Session()

  def _build_graph(self, tf_graph, scope, model_dir):
    """Construct a TensorGraph containing the policy and loss calculations."""
    features = [Feature(shape=[None] + list(s)) for s in self._env.state_shape]
    policy_layers = self._policy.create_layers(features)
    action_prob = policy_layers['action_prob']
    value = policy_layers['value']
    rewards = Weights(shape=(None, 1))
    actions = Label(shape=(None, self._env.n_actions))
    loss = A3CLoss(
        self.value_weight,
        self.entropy_weight,
        in_layers=[rewards, actions, action_prob, value])
    graph = TensorGraph(
        batch_size=self.max_rollout_length,
        use_queue=False,
        graph=tf_graph,
        model_dir=model_dir)
    for f in features:
      graph._add_layer(f)
    graph.add_output(action_prob)
    graph.add_output(value)
    graph.set_loss(loss)
    graph.set_optimizer(self.optimizer)
    with graph._get_tf("Graph").as_default():
      with tf.variable_scope(scope):
        graph.build()
    return graph, features, rewards, actions, action_prob, value

  def fit(self, total_steps, max_checkpoints_to_keep=5,
          checkpoint_interval=600, restore=False):
    """Train the policy.

    Parameters
    ----------
    total_steps: int
      the total number of time steps to perform on the environment, across all rollouts
      on all threads
    max_checkpoints_to_keep: int
      the maximum number of checkpoint files to keep.  When this number is reached, older
      files are deleted.
    checkpoint_interval: float
      the time interval at which to save checkpoints, measured in seconds
    restore: bool
      if True, restore the model from the most recent checkpoint and continue training
      from there.  If False, retrain the model from scratch.
    """
    with self._graph._get_tf("Graph").as_default():
      step_count = [0]
      workers = []
      threads = []
      for i in range(multiprocessing.cpu_count()):
        workers.append(_Worker(self, i))
      self._session.run(tf.global_variables_initializer())
      if restore:
        self.restore()
      for worker in workers:
        thread = threading.Thread(
            name=worker.scope,
            target=lambda: worker.run(step_count, total_steps))
        threads.append(thread)
        thread.start()
      variables = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='global')
      saver = tf.train.Saver(variables, max_to_keep=max_checkpoints_to_keep)
      checkpoint_index = 0
      while True:
        threads = [t for t in threads if t.isAlive()]
        if len(threads) > 0:
          threads[0].join(checkpoint_interval)
        checkpoint_index += 1
        saver.save(
            self._session, self._graph.save_file, global_step=checkpoint_index)
        if len(threads) == 0:
          break

  def predict(self, state):
    """Compute the policy's output predictions for a state.

    Parameters
    ----------
    state: array
      the state of the environment for which to generate predictions

    Returns
    -------
    the array of action probabilities, and the estimated value function
    """
    with self._graph._get_tf("Graph").as_default():
      feed_dict = _create_feed_dict(self._features, state)
      return self._session.run(
          [self._action_prob.out_tensor, self._value.out_tensor],
          feed_dict=feed_dict)

  def select_action(self, state, deterministic=False):
    """Select an action to perform based on the environment's state.

    Parameters
    ----------
    state: array
      the state of the environment for which to select an action
    deterministic: bool
      if True, always return the best action (that is, the one with highest probability).
      If False, randomly select an action based on the computed probabilities.

    Returns
    -------
    the index of the selected action
    """
    with self._graph._get_tf("Graph").as_default():
      feed_dict = _create_feed_dict(self._features, state)
      probabilities = self._session.run(
          self._action_prob.out_tensor, feed_dict=feed_dict)
      if deterministic:
        return probabilities.argmax()
      else:
        return np.random.choice(
            np.arange(self._env.n_actions), p=probabilities[0])

  def restore(self):
    """Reload the model parameters from the most recent checkpoint file."""
    last_checkpoint = tf.train.latest_checkpoint(self._graph.model_dir)
    if last_checkpoint is None:
      raise ValueError('No checkpoint found')
    with self._graph._get_tf("Graph").as_default():
      variables = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='global')
      saver = tf.train.Saver(variables)
      saver.restore(self._session, last_checkpoint)

class _Worker(object):
  """A Worker object is created for each training thread."""

  def __init__(self, a3c, index):
    self.a3c = a3c
    self.index = index
    self.scope = 'worker%d' % index
    self.env = copy.deepcopy(a3c._env)
    self.env.reset()
    self.graph, self.features, self.rewards, self.actions, self.action_prob, self.value = a3c._build_graph(
        a3c._graph._get_tf('Graph'), self.scope, None)
    with a3c._graph._get_tf("Graph").as_default():
      local_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
                                     self.scope)
      global_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
                                      'global')
      gradients = tf.gradients(self.graph.loss.out_tensor, local_vars)
      grads_and_vars = list(zip(gradients, global_vars))
      self.train_op = a3c._graph._get_tf('Optimizer').apply_gradients(
          grads_and_vars)
      self.update_local_variables = tf.group(
          * [tf.assign(v1, v2) for v1, v2 in zip(local_vars, global_vars)])

  def run(self, step_count, total_steps):
    with self.graph._get_tf("Graph").as_default():
      session = self.a3c._session
      while step_count[0] < total_steps:
        session.run(self.update_local_variables)
        episode_states, episode_actions, episode_rewards = self.create_rollout()
        feed_dict = {}
        for f, s in zip(self.features, episode_states):
          feed_dict[f.out_tensor] = s
        feed_dict[self.rewards.out_tensor] = episode_rewards
        feed_dict[self.actions.out_tensor] = episode_actions
        session.run(self.train_op, feed_dict=feed_dict)
        step_count[0] += len(episode_actions)

  def create_rollout(self):
    """Generate a rollout."""
    n_actions = self.env.n_actions
    session = self.a3c._session
    states = [[] for i in range(len(self.features))]
    actions = []
    rewards = []
    for i in range(self.a3c.max_rollout_length):
      if self.env.terminated:
        break
      state = self.env.state
      for j in range(len(state)):
        states[j].append(state[j])
      feed_dict = _create_feed_dict(self.features, state)
      probabilities = session.run(
          self.action_prob.out_tensor, feed_dict=feed_dict)
      action = np.random.choice(np.arange(n_actions), p=probabilities[0])
      actions.append(np.zeros(n_actions))
      actions[i][action] = 1.0
      rewards.append(self.env.step(action))
    if not self.env.terminated:
      # Add an estimate of the reward for the rest of the episode.
      feed_dict = _create_feed_dict(self.features, self.env.state)
      rewards[-1] += self.a3c.discount_factor * session.run(
          self.value.out_tensor, feed_dict)
    for j in range(len(rewards) - 1, 0, -1):
      rewards[j - 1] += self.a3c.discount_factor * rewards[j]
    if self.env.terminated:
      self.env.reset()
    return np.array(states), np.array(actions), np.array(rewards).reshape(
        (len(rewards), 1))
<|endoftext|>


<commit_before># !/usr/local/bin/python3.4.2
# ----Copyright (c) 2017 Carnegie Hall | The MIT License (MIT)----
# ----For the full license terms, please visit https://github.com/CarnegieHall/linked-data/blob/master/LICENSE----

## Argument[0] is script to run
## Argument[1] is path to graph

import os
import rdflib
from rdflib import Graph
from rdflib import Graph, Literal, Namespace, OWL, RDF, URIRef, XSD
from rdflib.namespace import DCTERMS, FOAF, OWL, RDF, RDFS, SKOS
from rdflib.plugins.serializers.nt import NTSerializer
from rdflib.plugins.serializers.turtle import TurtleSerializer
import sys

def fetch_name(filepath):
    return os.path.splitext(os.path.basename(filepath))[0]

chnames = Namespace('http://data.carnegiehall.org/names/')
chvenues = Namespace('http://data.carnegiehall.org/venues/')
chinstruments = Namespace('http://data.carnegiehall.org/instruments/')
chworks = Namespace('http://data.carnegiehall.org/works/')
event = Namespace('http://purl.org/NET/c4dm/event.owl#')
gndo = Namespace('http://d-nb.info/standards/elementset/gnd#')
mo = Namespace ('http://purl.org/ontology/mo/')
schema = Namespace('http://schema.org/')
dbp = Namespace('http://dbpedia.org/ontology/')
dbpedia = Namespace('http://dbpedia.org/resource/')
geonames = Namespace('http://sws.geonames.org/ontology#')
mbz = Namespace('http://musicbrainz.org/artist/')
lcnaf = Namespace('http://id.loc.gov/authorities/names/')
lcMarcRel = Namespace('http://id.loc.gov/vocabulary/relators/')

filePath_1 = sys.argv[1]
fileName = fetch_name(filePath_1)

turtle = 'turtle'
ttl = 'ttl'
newSerialization = input('Enter new serialization: ')
newExtension = input('Enter new file extension: ')
print('The new serialization will be', newSerialization, 'with extension', newExtension)

g = Graph()

g.parse(filePath_1, format=str(newSerialization))

graph_path = os.path.join(
    os.path.dirname(__file__), os.pardir, 'Graphs', (fileName + '.' + newExtension))

g.bind("chinstruments", chinstruments)
g.bind("chnames", chnames)
g.bind("chworks", chworks)
g.bind("dbpedia-owl", dbp)
g.bind("dbpedia", dbpedia)
g.bind("dcterms", DCTERMS)
g.bind("event", event)
g.bind("foaf", FOAF)
g.bind("geonames", geonames)
g.bind("gndo", gndo)
g.bind("lcnaf", lcnaf)
g.bind("lcMarRel", lcMarcRel)
g.bind("mbz", mbz)
g.bind("mo", mo)
g.bind("rdf", RDF)
g.bind("rdfs", RDFS)
g.bind("skos", SKOS)
g.bind("schema", schema)

g.serialize(destination=graph_path,format=str(newSerialization))

print('Finished converting', fileName, 'to', newSerialization)
<commit_msg>Fixed incorrect URL in geonames ns<commit_after># !/usr/local/bin/python3.4.2
# ----Copyright (c) 2017 Carnegie Hall | The MIT License (MIT)----
# ----For the full license terms, please visit https://github.com/CarnegieHall/linked-data/blob/master/LICENSE----

## Argument[0] is script to run
## Argument[1] is path to graph

import os
import rdflib
from rdflib import Graph
from rdflib import Graph, Literal, Namespace, OWL, RDF, URIRef, XSD
from rdflib.namespace import DCTERMS, FOAF, OWL, RDF, RDFS, SKOS
from rdflib.plugins.serializers.nt import NTSerializer
from rdflib.plugins.serializers.turtle import TurtleSerializer
import sys

def fetch_name(filepath):
    return os.path.splitext(os.path.basename(filepath))[0]

chnames = Namespace('http://data.carnegiehall.org/names/')
chvenues = Namespace('http://data.carnegiehall.org/venues/')
chinstruments = Namespace('http://data.carnegiehall.org/instruments/')
chworks = Namespace('http://data.carnegiehall.org/works/')
event = Namespace('http://purl.org/NET/c4dm/event.owl#')
gndo = Namespace('http://d-nb.info/standards/elementset/gnd#')
mo = Namespace ('http://purl.org/ontology/mo/')
schema = Namespace('http://schema.org/')
dbp = Namespace('http://dbpedia.org/ontology/')
dbpedia = Namespace('http://dbpedia.org/resource/')
geonames = Namespace('http://www.geonames.org/ontology#')
mbz = Namespace('http://musicbrainz.org/artist/')
lcnaf = Namespace('http://id.loc.gov/authorities/names/')
lcMarcRel = Namespace('http://id.loc.gov/vocabulary/relators/')

filePath_1 = sys.argv[1]
fileName = fetch_name(filePath_1)

turtle = 'turtle'
ttl = 'ttl'
newSerialization = input('Enter new serialization: ')
newExtension = input('Enter new file extension: ')
print('The new serialization will be', newSerialization, 'with extension', newExtension)

g = Graph()

g.parse(filePath_1, format=str(newSerialization))

graph_path = os.path.join(
    os.path.dirname(__file__), os.pardir, 'Graphs', (fileName + '.' + newExtension))

g.bind("chinstruments", chinstruments)
g.bind("chnames", chnames)
g.bind("chworks", chworks)
g.bind("dbpedia-owl", dbp)
g.bind("dbpedia", dbpedia)
g.bind("dcterms", DCTERMS)
g.bind("event", event)
g.bind("foaf", FOAF)
g.bind("geonames", geonames)
g.bind("gndo", gndo)
g.bind("lcnaf", lcnaf)
g.bind("lcMarRel", lcMarcRel)
g.bind("mbz", mbz)
g.bind("mo", mo)
g.bind("rdf", RDF)
g.bind("rdfs", RDFS)
g.bind("skos", SKOS)
g.bind("schema", schema)

g.serialize(destination=graph_path,format=str(newSerialization))

print('Finished converting', fileName, 'to', newSerialization)
<|endoftext|>