AlgebraNet - Learning Algebra from Scratch

%matplotlib inline
%config InlineBackend.figure_format='retina'

import operator

import numpy as np
import matplotlib.pyplot as plt
plt.rcParams["figure.figsize"] = (8, 8)
plt.rcParams["font.size"] = 14

from sklearn.cross_validation import train_test_split
from sklearn.utils import check_random_state

from keras.models import Sequential
from keras.layers import Dense, RepeatVector, TimeDistributed, Activation
from keras.layers.recurrent import LSTM
from keras.optimizers import SGD

OPS = ['+', '*', '-', '/']
ABC = '1234567890' + ''.join(OPS) + '. '

def encode(characters):
    """Encode a string as an array of integers."""
    char_idx = dict((c, i) for i, c in enumerate(ABC))
    X_ = np.zeros((len(characters), len(ABC)))
    for i, c in enumerate(characters):
        try:
            X_[i, char_idx[c]] = 1
        except:
            print('chars:', characters)
            print(i, c, char_idx[c], X_.shape)
    return X_


def decode(X, calc_argmax=True):
    """Decode an array of integers to a string."""
    idx_char = dict((i, c) for i, c in enumerate(ABC))
    X = X.argmax(axis=-1)    
    return ''.join(idx_char[i] for i in X)


# Generate data
def make_maths_problem(n_samples=1000, n_digits=3, invert=True, random_state=None):
    rng = check_random_state(random_state)
    
    X_ = []
    y_ = []
    
    assert all(op in ABC for op in OPS), "Not all operations are in the alphabet"
    
    math_op = {'+': operator.add, 
               '-': operator.sub,
               '*': operator.mul,
               '/': operator.truediv}
    n_ops = len(OPS)
    # the one represents the character for the operator
    max_len = 2 * n_digits + 1

    while len(X_) < n_samples:
        a, b = rng.randint(10**n_digits, size=2)
        op = rng.choice(OPS)
        a_op_b = "%i%s%i" % (a, op, b)
        if a_op_b not in X_:
            answer = math_op[op](a, b)

            if answer == np.inf or np.isnan(answer):
                continue
            
            if invert:
                X_.append(a_op_b.ljust(max_len)[::-1])
            else:
                X_.append(a_op_b.ljust(max_len))
            # make sure the string is white space padded but no longer
            # than the max length (could happen when representing 1/3)
            y_.append(str(answer).ljust(2 * n_digits)[:2 * n_digits])
    
    X = np.zeros((n_samples, max_len, len(ABC)),
                 dtype=np.bool)
    y = np.zeros((n_samples, 2 * n_digits, len(ABC)),
                 dtype=np.bool)
    for i, x in enumerate(X_):
        X[i] = encode(x)
        
    for i, x in enumerate(y_):
        y[i] = encode(x)

    return (X, y)

encode("23+42")

X, y = make_maths_problem(10, n_digits=3)

for i in range(10):
    # undo question reversal
    q = decode(X[i][::-1])
    correct = decode(y[i])
    print(q, '=', correct)

X, y = make_maths_problem(2*50000, n_digits=3, random_state=234)
X_train,X_test, y_train,y_test = train_test_split(X, y, train_size=0.8)

def rnn_model(hidden_size=128, n_digits=3, abc_size=len(ABC)):
    model = Sequential()
    # encoder
    model.add(LSTM(hidden_size, input_shape=(None, abc_size)))
    # input for the decoder
    # this sets the length of the output sequence and has to
    # match the length we used when constructing the examples
    model.add(RepeatVector(2 * n_digits))

    # decoder model
    model.add((LSTM(hidden_size, return_sequences=True)))

    # For each step of the output sequence, decide which character should be chosen
    model.add(TimeDistributed(Dense(abc_size)))
    model.add(Activation('softmax'))

    model.compile(loss='categorical_crossentropy',
                  optimizer='adam',
                  metrics=['accuracy'])

    return model

BATCH_SIZE = 128

model = rnn_model()

# use same 10 examples for all iterations
n_test = 10
idx = np.random.randint(len(X_test), size=n_test)
Xs, ys = X_test[idx], y_test[idx]

for iteration in range(3):
    print()
    print('-' * 50)
    print('After', 10 * (iteration + 1), 'epochs:')
    model.fit(X_train, y_train, batch_size=BATCH_SIZE, nb_epoch=10, verbose=False,
              validation_data=(X_test, y_test))

    probas = model.predict_proba(Xs, verbose=0)

    for i in range(n_test):
        # undo question reversal
        q = decode(Xs[i][::-1])
        correct = decode(ys[i])
        guess = decode(probas[i])
        print(q, '=', correct, '=?', guess)

def fractional_difference(n_test=200):
    idx = np.random.randint(len(X_test), size=n_test)
    Xs, ys = X_test[idx], y_test[idx]
    probas = model.predict_proba(Xs, verbose=0)
    true_ = np.array([float(decode(ys[i])) for i in range(n_test)])
    predicted_ = np.array([float(decode(probas[i])) for i in range(n_test)])

    _ = plt.hist((true_ - predicted_)/true_, range=(-0.1, 0.1), bins=20)
    plt.xlabel('(true-guess)/true')
    plt.xlim(-0.11, 0.11)
    
fractional_difference()