learn-corpus

#!/usr/bin/python
# -*- coding: utf-8 -*-
import time
import os
import sys
HERE = os.path.dirname(__file__)
sys.path.append(HERE)
import argparse
import charmodel
import random
import json
import numpy as np
import errno
import arguments
import language

from meta import save_opinions

VALIDATION_ABORT_THRESHOLD = 6.0


def full_entropy_matrix(net, texts, ignore_start):
    entropies_cache = {}
    for tid, text in texts.items():
        entropies = net.test(text, ignore_start, 0)
        entropies_cache[tid] = entropies
    return entropies_cache


def get_problem_and_control_sets(entropy_cache, problems):
    classified_tids = set()
    for records in problems.values():
        classified_tids.update(x[2] for x in records)

    control_tids = set(x for x in entropy_cache
                       if x not in classified_tids)

    print ("found %d classified texts and %d control (%d total)" %
           (len(classified_tids), len(control_tids), len(entropy_cache)))

    return classified_tids, control_tids


def entropies_to_opinions(entropy_cache, problems, text_len_lut):
    # Get the cross-entropy between all subnets against the texts of
    # each problem. Because texts are often in more than one problem,
    # we keep a cache (this is sort-of expensive).
    affinities = {}
    names = {}
    all_control_texts = {}
    all_control_models = {}
    all_text_lengths = {}

    classified_tids, control_tids = \
        get_problem_and_control_sets(entropy_cache, problems)

    control_text_lut = {}
    control_model_lut = {k: 0 for k in classified_tids}
    control_n = max(1.0, len(control_tids))

    for tid in classified_tids:
        entropies = entropy_cache[tid]
        control_text_lut[tid] = (sum(entropies[k] for k in control_tids) /
                                 control_n)

    for tid in control_tids:
        entropies = entropy_cache[tid]
        for k in classified_tids:
            control_model_lut[k] += entropies[k]

    for k in classified_tids:
        control_model_lut[k] /= control_n

    for pid, records in problems.items():
        names[pid], _, text_ids = zip(*records)

        entropy_matrix = []
        control_texts = []
        control_models = []
        lengths = []
        for fn, ffn, tid in records:
            entropies = entropy_cache[tid]
            local_entropies = [entropies[x] for x in text_ids]

            entropy_matrix.append(local_entropies)
            # control_texts is list of lists, the right shape for
            # column-oriented subtraction.
            control_texts.append([control_text_lut[tid]])
            control_models.append(control_model_lut[tid])
            lengths.append(text_len_lut[tid])

        affinities[pid] = np.array(entropy_matrix)
        all_control_texts[pid] = np.array(control_texts)
        all_control_models[pid] = np.array(control_models)
        all_text_lengths[pid] = lengths
    return {
        'affinities': affinities,
        'names': names,
        'control_texts': all_control_texts,
        'control_models': all_control_models,
        'text_lengths': all_text_lengths
    }


def opine(net, texts, problems, ignore_start):
    print "beginning opine"
    start = time.time()
    entropy_cache = full_entropy_matrix(net, texts, ignore_start)
    print "got raw entropies in %d:%02d" % divmod(time.time() - start, 60)
    text_len_lut = {k: len(v) for k, v in texts.iteritems()}
    return entropies_to_opinions(entropy_cache, problems, text_len_lut)


def calc_ventropy_change(ventropies, prev_ventropies):
    ve_sum = 0.0
    ve_sum2 = 0.0
    ve_diff_sum = 0.0
    ve_diff_sum2 = 0.0
    for prev, e in zip(prev_ventropies, ventropies):
        diff = prev - e
        ve_diff_sum += diff
        ve_diff_sum2 += diff * diff
        ve_sum += e
        ve_sum2 += e * e
    ve_scale = 1.0 / len(ventropies)
    ve_mean = ve_sum * ve_scale
    ve_stddev = (ve_sum2 * ve_scale - ve_mean * ve_mean) ** 0.5
    ve_diff_mean = ve_diff_sum * ve_scale
    ve_diff_stddev = (ve_diff_sum2 * ve_scale -
                      ve_diff_mean * ve_diff_mean) ** 0.5

    return (ve_mean, ve_stddev, ve_diff_mean, ve_diff_stddev)


def train(net, texts, leakage, epochs, leakage_decay,
          learn_rate_decay, ignore_start, validation_text,
          validation_reverse_length, opinion_args,
          net_filename, confab_n=3, confab_len=79,
          confab_interval=0, confab_caps_marker=None):
    prev_ventropies = net.test(validation_text, ignore_start, 1)
    v_history = []

    for name in net.class_names[:confab_n]:
        print name,

    net.start_confab(confab_interval, confab_n, confab_len,
                     confab_caps_marker)

    for i in range(epochs):
        if not confab_interval:
            print ("starting epoch %d with learn-rate %s, "
                   "leakage %s" % (i + 1, net.learn_rate, leakage))
        start = time.time()

        for name, text in texts.items():
            net.train(text, name, leakage=leakage,
                      ignore_start=ignore_start)
        if net_filename:
            net.save(net_filename)

        middle = time.time()
        if not confab_interval:
            print "training took %d:%02d" % divmod(middle - start, 60)

        ventropies = net.test(validation_text, ignore_start, 1)
        vstats = calc_ventropy_change(ventropies, prev_ventropies)
        prev_ventropies = ventropies
        print ("validation entropy %.3f±%.3f  diff % .3f±%.3f" % vstats)
        end = time.time()
        if not confab_interval:
            print "validation took %d:%02d" % divmod(end - middle, 60)
        if vstats[0] > VALIDATION_ABORT_THRESHOLD:
            raise RuntimeError('Validation error %.3f > %.3f; aborting' %
                               (vstats[0], VALIDATION_ABORT_THRESHOLD))

        if validation_reverse_length:
            v_history.append(vstats[2] < 0.0)
            v_history = v_history[-validation_reverse_length:]
            if all(v_history):
                print ("stopping because validation has been getting worse "
                       "for %d epochs" % validation_reverse_length)
                return

        if opinion_args is not None:
            interval, opinion_dest, problems = opinion_args
            if (i + 1) % interval == 0:
                print "starting periodic opinion"
                start = time.time()
                fn = opinion_dest % i
                opinion = opine(net, texts, problems, ignore_start)
                save_opinions(fn, problems=problems, **opinion)
                print "opinion took %d:%02d" % divmod(time.time() - start, 60)

        leakage *= leakage_decay
        net.learn_rate *= learn_rate_decay

    net.stop_confab()


def get_net_and_corpus(srcdir, controldir, control_n, lang,
                       reverse, validation_dir, word_df_threshold,
                       net_kwargs):
    print srcdir
    texts, problems = language.load_corpus(srcdir, lang)

    if control_n and controldir:
        control_texts, _ = language.load_control_texts(controldir, lang)

        if len(control_texts) > control_n:
            # sort, to avoid hashing indeterminacy
            c_items = sorted(control_texts.items())
            random.shuffle(c_items)
            control_texts = dict(c_items[:control_n])

        texts.update(control_texts)

        print "using %d control texts; wanted %d" % (len(control_texts),
                                                     control_n)

    if word_df_threshold:
        language.word_df_filter(texts, word_df_threshold)

    # remap into net enumeration encoding
    alphabet = charmodel.Alphabet(''.join(texts.values()), ignore_case=False,
                                  threshold=1e-10)
    print >>sys.stderr, "alphabet is %s" % (alphabet.alphabet,)

    remapped = {}
    for k, v in texts.iteritems():
        t = alphabet.encode_text(v)
        if reverse:
            t = ''.join(reversed(t))
        remapped[k] = t

    if validation_dir:
        # validation text doesn't go in texts.
        v_texts, _ = language.load_control_texts(validation_dir, lang)
        v_text = alphabet.encode_text('\n'.join(v_texts.values()))
        if reverse:
            v_text = ''.join(reversed(v_text))

    textnames = sorted(texts.keys())

    metadata = json.dumps({
        'alphabet': alphabet.alphabet,
        'collapse_chars': alphabet.collapsed_chars,
        'version': 2,
        'class_names': textnames,
        'case_insensitive': False,
        'utf8': True,
        'collapse_space': False,
        'reverse': reverse,
    }, sort_keys=True)

    net = charmodel.Net(alphabet, textnames, metadata=metadata, **net_kwargs)
    return net, remapped, problems, v_text


def main():
    parser = argparse.ArgumentParser()
    rnn_args = parser.add_argument_group('RNN arguments')

    arguments.add_common_args(parser.add_argument)
    arguments.add_rnn_args(rnn_args.add_argument)

    args = parser.parse_args()

    if args.enable_fp_exceptions:
        charmodel.enable_fp_exceptions()

    if args.rng_seed != -1:
        random.seed(args.rng_seed)

    if args.save is None:
        if args.output_dir is None:
            print "I don't know where to write anything!"
            sys.exit(1)

        args.save = os.path.join(args.output_dir,
                                 '%s-affinities.pickle' % args.lang)

        try:
            os.makedirs(args.output_dir)
            print "made output directory %r" % args.output_dir
        except OSError, e:
            if e.errno != errno.EEXIST:
                raise

    if args.epochs is None:
        if args.stop_on_validation_reverse:
            args.epochs = 100
        else:
            args.epochs = 1

    net_kwargs = {}
    train_kwargs = {}
    for k, v in vars(args).items():
        if k in ("bptt_depth",
                 "hidden_size",
                 "rng_seed",
                 "log_file",
                 "verbose",
                 "learn_rate",
                 "temporal_pgm_dump",
                 "periodic_pgm_dump",
                 "periodic_pgm_period",
                 "basename",
                 "activation",
                 "learning_method",
                 "batch_size",
                 "filename",
                 "presynaptic_noise",
                 "init_method",
                 ):
            net_kwargs[k] = v

        if (v is not None and
            k in ("confab_interval",
                  "confab_n",
                  "confab_len",
                  "confab_caps_marker",
                 )):
            train_kwargs[k] = v

    net, texts, problems, v_text = get_net_and_corpus(args.input_dir,
                                                      args.control_corpus,
                                                      args.control_n,
                                                      args.lang,
                                                      args.reverse,
                                                      args.validation_corpus,
                                                      args.word_df_threshold,
                                                      net_kwargs)

    if args.opinion_every:
        opinion_args = (args.opinion_every,
                        args.save + '-%d',
                        problems)
    else:
        opinion_args = None

    train(net, texts, args.leakage, args.epochs, args.leakage_decay,
          args.learn_rate_decay, args.ignore_start, v_text,
          args.stop_on_validation_reverse, opinion_args,
          args.filename, **train_kwargs)

    opinion = opine(net, texts, problems, args.ignore_start)
    save_opinions(args.save, problems=problems, **opinion)

main()