<file_sep>import torch
from torch.autograd import Variable

from parameterized import parameterized

from embzip.data import load_embeddings_txt, make_vocab
from embzip.model import gumbel_softmax, EmbeddingCompressor
from embzip.data import load_embeddings_txt, make_vocab, load_hdf5, save_hdf5
from embzip.model import (gumbel_softmax, EmbeddingCompressor,
                          FactorizedEmbeddingsOutput,
                          FactorizedEmbeddingsInput)


@parameterized([
    ('data/glove.6B.300d.txt', 100),
])
def test_glove_txt(path, maxsize):
    gloves = load_embeddings_txt(path, maxsize)
    assert len(gloves) <= maxsize
    for k, v in gloves.items():
        assert isinstance(k, str)
        assert isinstance(v, torch.autograd.Variable)
        assert v.data.shape[0] == 1


@parameterized([
    ('data/glove.6B.300d.txt', 100),
])
def test_vocab(path, maxsize):
    gloves = load_embeddings_txt(path, maxsize)
    vocab = make_vocab(gloves)

    assert len(vocab['idx2word']) == \
           len(vocab['word2idx']) == \
           vocab['emb_tables'].data.shape[0]


@parameterized([
    (None, 3, 4),
    (2, 3, 4),
])
def test_one_hot(batch, n_dic, n_codes):
    size = torch.Size([n_dic])
    if batch is not None:
        size = torch.Size([batch]) + size

    t = torch.LongTensor(size).random_() % n_codes

    o = one_hot(t, n_codes)
    print(o)

    exp_size = size + torch.Size([n_codes])
    assert exp_size == o.size()


@parameterized([
    (2, 3, 4)
])
def test_gumbel_softmax(b, m, k):

    logits = Variable(1 * torch.Tensor(b * m, k).uniform_())

    gsm = gumbel_softmax(logits, tau=1, hard=True)
    print(gsm)
    print(gsm.view(b, m * k))

    assert torch.equal(torch.ones(b*m), torch.sum(gsm, -1).data)


@parameterized([
    (2, 5, 2, 4)
])
def test_compressor(batch, emb_size, n_dic, n_codes):
    ec = EmbeddingCompressor(emb_size, n_dic, n_codes)
    inp = Variable(torch.Tensor(batch, emb_size).uniform_())
    out = ec(inp)
    print(out)


@parameterized([
    ('data/glove.6B.300d.txt', 100),
])
def test_save_hdf5(path, num_embs):
    gloves, _ = load_embeddings_txt(path, num_embs)
    ec = EmbeddingCompressor(300, 8, 16)

    orig_embs, orig_vocab_map = save_hdf5('foo.h5', ec, gloves)
    embs, vocab_map = load_hdf5('foo.h5')

    assert(np.array_equal(orig_embs, embs))
    assert orig_vocab_map == vocab_map


@parameterized([
    ([[0]], [[0, 1], [1, 2]]),
    ([[0, 1], [1, 2]], [[1, 2], [3, 2], [0, 3]]),
    ([[0, 1], [2, 1]], [[0, 1], [1, 2], [2, 3]])
])
def test_fe_input(indices, index_map):
    vocab_map = {i: index_map[i] for i in range(len(index_map))}
    emb_table = np.arange(10).reshape(5, 2)
    indices = Variable(torch.LongTensor(indices))
    print(indices)

    fei = FactorizedEmbeddingsInput(emb_table, vocab_map)
    out = fei(indices)
    print(out)

    # figure out expanded embeddings
    index_map_t = torch.LongTensor(index_map)
    index_map_t = torch.autograd.Variable(index_map_t)

    exp_emb_table = torch.nn.Embedding(*emb_table.shape)
    exp_emb_table.weight.data.copy_(torch.from_numpy(emb_table))

    final_emb_table = exp_emb_table(index_map_t).sum(-2).data
    et = torch.nn.Embedding(*final_emb_table.shape)
    et.weight.data.copy_(final_emb_table)

    exp_out = et(indices)
    print(exp_out)
    assert torch.equal(exp_out.data, out.data)

@parameterized([
    ([[1, 1]], [[0, 1], [1, 2]]),
    ([[0, 1], [1, 2]], [[1, 2], [3, 2], [0, 3]]),
    ([[0, 1]], [[0, 1], [1, 2], [2, 3]]),
    ([[0, 1], [1, 1]], [[0, 1], [1, 2], [2, 3]]),
    ([[0, 1], [1, 1], [1, 0]], [[0, 1], [1, 2], [2, 3]])
])
def test_fe_output(pe, index_map):
    e = 2
    mk = 5
    v = len(index_map)
    mk_table = np.arange(10).reshape(e, mk)
    mk_table_T = np.arange(10).reshape(mk, e)
    vocab_map = {i: index_map[i] for i in range(len(index_map))}
    feo = FactorizedEmbeddingsOutput(np.array(mk_table), vocab_map)

    print(mk_table_T)
    emb_table = np.vstack([mk_table_T[index_map[i]].sum(0) for i in range(v)])
    print('e', emb_table)

    pe = Variable(torch.Tensor(pe))

    out = feo(pe)

    out_exp = pe.data @ torch.from_numpy(emb_table).float().t()
    print('o', out_exp)

    assert torch.equal(out.data, out_exp)
<file_sep>import torch
import torch.nn.functional as F
from torch.autograd import Variable
from torch import nn


def sample_gumbel(shape, eps=1e-10, out=None):
    """
    Sample from Gumbel(0, 1)
    based on
    https://github.com/ericjang/gumbel-softmax/blob/3c8584924603869e90ca74ac20a6a03d99a91ef9/Categorical%20VAE.ipynb ,
    (MIT license)
    """
    U = out.resize_(shape).uniform_() if out is not None else torch.rand(shape)
    return - torch.log(eps - torch.log(U + eps))


def gumbel_softmax_sample(logits, tau=1, eps=1e-10):
    """
    Draw a sample from the Gumbel-Softmax distribution
    based on
    https://github.com/ericjang/gumbel-softmax/blob/3c8584924603869e90ca74ac20a6a03d99a91ef9/Categorical%20VAE.ipynb
    (MIT license)
    """
    gumbel_noise = sample_gumbel(logits.size(), eps=eps, out=logits.data.new())
    y = logits + Variable(gumbel_noise)
    #y = F.log_softmax(y.view(-1, logits.size(-1)) / tau)
    y = F.softmax(y.view(-1, logits.size(-1)) / tau)
    return y.view_as(logits)


def gumbel_softmax(logits, tau=1, hard=False, eps=1e-10):
    """
    Sample from the Gumbel-Softmax distribution and optionally discretize.
    Args:
      logits: [batch_size, n_class] unnormalized log-probs
      tau: non-negative scalar temperature
      hard: if True, take argmax, but differentiate w.r.t. soft sample y
    Returns:
      [batch_size, n_class] sample from the Gumbel-Softmax distribution.
      If hard=True, then the returned sample will be one-hot, otherwise it will
      be a probability distribution that sums to 1 across classes
    Constraints:
    - this implementation only works on batch_size x num_features tensor for now
    based on
    https://github.com/ericjang/gumbel-softmax/blob/3c8584924603869e90ca74ac20a6a03d99a91ef9/Categorical%20VAE.ipynb ,
    (MIT license)
    """
    shape = logits.size()
    assert len(shape) == 2
    y_soft = gumbel_softmax_sample(logits, tau=tau, eps=eps)
    if hard:
        _, k = y_soft.data.max(-1)
        # this bit is based on
        # https://discuss.pytorch.org/t/stop-gradients-for-st-gumbel-softmax/530/5
        y_hard = logits.data.new(*shape).zero_().scatter_(-1, k.view(-1, 1), 1.0)
        # this cool bit of code achieves two things:
        # - makes the output value exactly one-hot (since we add then
        #   subtract y_soft value)
        # - makes the gradient equal to y_soft gradient (since we strip
        #   all other gradients)
        y = Variable(y_hard - y_soft.data) + y_soft
    else:
        y = y_soft
    return y


class FactorizedEmbeddingsInput(nn.Module):
    def __init__(self, emb_table, vocab_map):
        '''
            emb_table: numpy array (V x E)
            vocab_map: {word -> [indices]}
        '''
        super().__init__()
        self.emb_tables = nn.Embedding(*emb_table.shape)
        self.emb_tables.weight.data.copy_(torch.from_numpy(emb_table))
        self.vocab_map = vocab_map
        self.index_map = torch.LongTensor(list(vocab_map.values()))

    def forward(self, indices):
        # B x S -> BS x M
        exp_indices = self.index_map[indices.data.view(-1)]
        # exp_indices = exp_indices.view(indices.size(0), indices.size(1), -1)

        # BS x M -> BS x M x E
        exp_embs = self.emb_tables(Variable(exp_indices))

        # BS x M x E -> BS x E
        exp_embs = exp_embs.sum(-2)

        # BS x E -> B x S x E
        exp_embs = exp_embs.view(indices.size(0), indices.size(1), -1)

        return exp_embs

    @classmethod
    def load_hdf5(cls, h5_file):
        print('Loading compressed embeddings in %s' % h5_file)
        f = h5py.File(h5_file, 'r')
        # embeddings
        embs = np.array(f['embeddings'])
        words = list(f['vocab'])
        indices = list(f['indices'])
        # print(embs.dtype, words, indices)
        vocab_map = {k: v.tolist() for k, v in zip(words, indices)}
        return cls(embs, vocab_map)



class FactorizedEmbeddingsOutput(nn.Module):
    def __init__(self, emb_table, vocab_map):
        '''
            emb_table: numpy array (V x E)
            vocab_map: {word -> [indices]}
        '''
        super().__init__()
        self.linear = nn.Linear(*emb_table.shape, bias=False)
        self.linear.weight.data.copy_(torch.from_numpy(emb_table))
        self.vocab_map = vocab_map
        self.index_map = torch.LongTensor(list(vocab_map.values()))

    def forward(self, x):
        # B x E -> B x MK
        mk_scores = self.linear(x)

        # B x MK -> B x V
        v_size = self.index_map.size(0)
        v_scores = [mk_scores[:, self.index_map[i]].sum(1)
                    for i in range(v_size)]
        return torch.stack(v_scores, 1)

    @classmethod
    def load_hdf5(cls, h5_file):
        print('Loading compressed embeddings in %s' % h5_file)
        f = h5py.File(h5_file, 'r')
        # embeddings
        embs = np.array(f['embeddings'])
        words = list(f['vocab'])
        indices = list(f['indices'])
        # print(embs.dtype, words, indices)
        vocab_map = {k: v.tolist() for k, v in zip(words, indices)}
        return cls(embs, vocab_map)


class EmbeddingCompressor(nn.Module):
    def __init__(self, emb_size, n_tables, n_codes, temp=1, hard=True):
        """
            emb_size: (E) embedding size (both in input and output)
            n_tables: (M) number of embedding tables
            n_codes:  (K) number of codes in each table
            temp:     temperature of gumbel-softmax
            hard:     switch between hard and soft gumbel-softmax
        """
        super().__init__()
        self.emb_size = emb_size
        self.n_codes = n_codes
        self.n_tables = n_tables
        self.temp = temp
        self.hard = hard
        mk_half = n_tables * n_codes // 2
        self.linear1 = nn.Linear(emb_size, mk_half)
        self.linear2 = nn.Linear(mk_half, n_tables * n_codes)
        self.emb_tables = nn.Linear(n_tables * n_codes, emb_size, bias=False)

    def get_indices(self, x):
        """Given a batch of embeddings of shape B x E
           returns a list-of-lists of non-zero indices, one per row
        """
        # B x E -> B x MK
        one_hot = self.encode(x)
        return [torch.nonzero(one_hot[i]).squeeze().cpu().tolist()
                for i in range(one_hot.size(0))]

    def encode(self, x):
        """Given a batch of embeddings of shape B x E
           returns a batch of one-hot embeddings of shape B x M x K
        """
        # B x E -> B x MK/2
        h = nn.Tanh()(self.linear1(x))

        # B x MK/2 -> B x MK
        a = nn.functional.softplus(self.linear2(h))

        # B x MK -> BM x K
        a = a.view(-1, self.n_codes)

        # BM x K -> BM x K
        d = gumbel_softmax(a, tau=self.temp, hard=self.hard)

        # BM x K -> B x MK
        return d.view(-1, self.n_tables * self.n_codes)

    def forward(self, x):
        # B x E -> B x MK
        one_hot = self.encode(x)

        # B x MK -> B x E
        return self.emb_tables(one_hot)
<file_sep>import math
import torch
import numpy as np
import h5py

from collections import OrderedDict


def load_embeddings_txt(path, max=None, vsize=0):
    train = OrderedDict()
    valid = OrderedDict()
    with open(path, 'rt', encoding='utf-8') as ef:
        for i, line in enumerate(ef):
            if i >= max+vsize:
                break
            tokens = line.split()
            word = tokens[0]
            # 1 x EmbSize
            vector = np.array(tokens[1:], dtype=np.float32)[None, :]
            if i < max:
                train[word] = vector
            if i > max:
                valid[word] = vector
    return train, valid


def make_vocab(word2vector):
    if not word2vector:
        return None
    idx2word = []
    word2idx = {}
    all_vars = []
    for word, vector in word2vector.items():
        word2idx[word] = len(idx2word)
        idx2word.append(word)
        all_vars.append(torch.from_numpy(vector))
    embeddings = torch.cat(all_vars, 0)
    embeddings = torch.autograd.Variable(embeddings, requires_grad=False)
    return dict(idx2word=idx2word,
                word2idx=word2idx,
                embeddings=embeddings,
                n_embs=embeddings.size(0),
                emb_size=embeddings.size(1))


def print_compression_stats(vocab, n_tables, n_codes):
    n_words = vocab['n_embs']
    emb_size = vocab['emb_size']
    orig_mb = n_words * emb_size * 4 / 2 ** 20
    comp_mb = n_tables * n_codes * emb_size * 4 / 2 ** 20
    ratio = -   100 * (orig_mb - comp_mb) / orig_mb

    print('[ORIG] {:4} words with embedding size {:3}: {:3.3} MB'.format(n_words, emb_size, orig_mb))
    print('[COMP] {:4} words ({:2} tables, {:3} codes):  {:3.3} MB ({:.2f}%)'.format(n_words, n_tables, n_codes, comp_mb, ratio))


def check_training(old_params, new_params):
    if old_params is None:
        return
    old_sums = [p.sum().data[0] for p in old_params]
    new_sums = [p.sum().data[0] for p in new_params]
    for old, new in zip(old_sums, new_sums):
        if old == new:
            print(old_sums)
            print(new_sums)
            print('WARNING: some parameter did not seem to change!')


def euclidian_dist(x, y):
    return math.sqrt((x - y).pow(2).sum())


def dump_reconstructed_embeddings(out_file, model, emb_dict):
    print('Saving reconstructed emb_tables for train set in %s' % out_file)
    with open(out_file, 'wt') as f:
        for word, emb in emb_dict.items():
            #print(word)
            emb = torch.autograd.Variable(torch.from_numpy(emb))
            emb_comp = model(emb).squeeze().cpu().data.numpy().tolist()
            for x in emb_comp:
                word += ' ' + str(x)
            f.write(word + '\n')


def save_hdf5(out_file, model, emb_dict):
    print('Saving compressed embedding in %s' % out_file)
    # embeddings
    f = h5py.File(h5_file, 'w')
    embs = model.emb_tables.weight.cpu().data.numpy()
    f.create_dataset('embeddings', data=embs)
    # words
    words = list(emb_dict.keys)

    # indices

def load_hdf5(h5_file):
    print('Loading compressed embeddings in %s' % h5_file)
    f = h5py.File(h5_file, 'r')
    # embeddings
    embs = np.array(f['embeddings'])
    words = list(f['vocab'])
    indices = list(f['indices'])
    #print(embs.dtype, words, indices)
    vocab_map = {k: v.tolist() for k, v in zip(words, indices)}
    return embs, vocab_map


def load_artichoke(h5_file):
    print('Loading Artichoke file: %s' % h5_file)
    f = h5py.File(h5_file, 'r')
    words = list(f['/data/vocab/7fd01463c400/values'])
    embs = f['/data/parameters/embedding/W']
    embs = np.split(embs, embs.shape[0])

    train_d = {k: v for k, v in zip(words, embs)}
    #print(words)
    f.close()
    return train_d, None


def save_artichoke(h5_file, model, emb_dict):
    print('Saving Artichoke file: %s' % h5_file)
    f = h5py.File(h5_file, 'r+')
    embs = []
    for word, emb in emb_dict.items():
        # print(word)
        emb = torch.autograd.Variable(torch.from_numpy(emb))
        emb_comp = model(emb).squeeze().cpu().data.numpy().tolist()
        embs.append(emb_comp)

    f_embs = f['/data/parameters/embedding/W']
    f_embs[...] = embs
    f.close()
<file_sep>import argparse
import math
import os
import torch

from torch import nn
from time import time

from embzip.data import load_embeddings_txt, make_vocab, print_compression_stats, euclidian_dist, check_training, dump_reconstructed_embeddings
from embzip.model import EmbeddingCompressor

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--path',
                        required=True,
                        help='Path of embedding file')
    parser.add_argument('--n_embs',
                        type=int,
                        default=75000,
                        help='Number of emb_tables to consider')
    parser.add_argument('--n_tables',
                        type=int,
                        required=True,
                        help='Number of embedding tables')
    parser.add_argument('--n_codes',
                        type=int,
                        required=True,
                        help='Number of emb_tables in each table')
    parser.add_argument('--samples',
                        type=int,
                        default=int(2e7),
                        help='Number of embeddings to train')
    parser.add_argument('--report_every',
                        type=int,
                        default=1000,
                        help='Reporting interval in batches')
    parser.add_argument('--batch_size',
                        type=int,
                        default=128,
                        help='Number of emb_tables in each batch')
    parser.add_argument('--lr',
                        type=float,
                        default=0.0001,
                        help='Learning rate for optimizer')
    parser.add_argument('--temp',
                        type=float,
                        default=1,
                        help='Gumbel-Softmax temperature')
    parser.add_argument('--valid',
                        type=int,
                        default=1000,
                        help='Number of emb_tables for validation')
    parser.add_argument('--cuda',
                        action='store_true',
                        help='Run on GPU')
    parser.add_argument('--hard',
                        action='store_true',
                        help='Use hard gumbel softmax')
    args = parser.parse_args()

    # load emb_tables
    train_g, valid_g = load_embeddings_txt(args.path, args.n_embs, args.valid)
    train_vocab = make_vocab(train_g)
    valid_vocab = make_vocab(valid_g) if valid_g else make_vocab(train_g)

    # stats
    print_compression_stats(train_vocab, args.n_tables, args.n_codes)

    # model
    ec = EmbeddingCompressor(train_vocab['emb_size'], args.n_tables, args.n_codes, args.temp, args.hard)

    # optimizer
    optim = torch.optim.Adam(ec.parameters(), lr=args.lr)

    # criterion
    criterion = torch.nn.MSELoss(size_average=True)

    losses = []
    old_params = None

    print('[CUDA]', args.cuda)
    if args.cuda:
        ec.cuda()
        ec = nn.DataParallel(ec)
        criterion.cuda()

    samples = 0

    try:
        while samples < args.samples:

            # shuffle emb_tables
            embeddings = train_vocab['embeddings'][torch.randperm(train_vocab['n_embs'])]

            # train
            for k in range(0, train_vocab['n_embs'], args.batch_size):
                t0 = time()
                samples += args.batch_size
                batch = embeddings[k:k+args.batch_size]
                if args.cuda:
                    batch = batch.cuda()

                y_pred = ec(batch)

                loss = criterion(y_pred, batch)
                losses.append(loss.data[0])

                optim.zero_grad()
                loss.backward()
                optim.step()

                t1 = args.batch_size * 1000 * (time() - t0)

                if samples % args.report_every == 0:
                    # check validation set and report
                    avg_train_loss = sum(losses) / len(losses)
                    losses = []

                    # validate
                    v_batch = valid_vocab['embeddings']
                    if args.cuda:
                        v_batch = v_batch.cuda()

                    v_y = ec(v_batch)

                    v_loss = criterion(v_y, v_batch)
                    v_loss = v_loss.data[0]
                    avg_euc = sum([euclidian_dist(v_y.data[i], v_batch.data[i]) for i in range(v_y.size(0))]) / v_y.size(0)
                    print('[%d] train: %.5f | valid: %.5f | euc_dist: %.2f | %.2f emb/sec' % (samples, avg_train_loss, v_loss, avg_euc, t1))
                    #print([p.sum().data[0] for p in ec.parameters()])

    except KeyboardInterrupt:
        print('Training stopped!')

    # dump full size reconstructed embeddings to file
    dump_reconstructed_embeddings(args.path + '.comp', ec, train_g)

    # save compressed embeddings in hdf5 format
    #:save_hdf5(args.path + '.h5;', ec, train_g)
<file_sep># MAC
http://download.pytorch.org/whl/torch-0.2.0.post3-cp36-cp36m-macosx_10_7_x86_64.whl

# LINUX
# http://download.pytorch.org/whl/cu80/torch-0.2.0.post3-cp36-cp36m-manylinux1_x86_64.whl

torchvision

jupyter
pandas
sklearn
matplotlib
psutil
scipy
scikit-learn
deap
update_checker
tqdm
stopit
# xgboost
scikit-mdr
skrebate
tpot
h5py

nose
parametrized
