Exploring Global Continuous Bag of Words (CBOW) Optimization¶

%%bash
#pip3 install 'git+ssh://git@github.com/IBM/regenerator@0.2.0' --user

import copy
import collections, itertools
import math
import random, re

import regenerator
import munch
import nltk

import scipy as sp
import scipy.spatial.distance as spdist

import torch as th
import torch.nn.functional as thfunc

import matplotlib.pyplot as plt
import matplotlib
%matplotlib inline
matplotlib.rcParams['figure.figsize'] = (9, 8)

class Tokenizer:                                                                                    
    def __init__(self):
        #self.nltk_word_tokenizer = nltk.tokenize.WordPunctTokenizer()                              
        self.nltk_word_tokenizer = nltk.tokenize.TreebankWordTokenizer()                            
        self.nltk_sent_tokenizer = nltk.data.load('tokenizers/punkt/english.pickle')
                                                                                                    
    def __call__(self, text):                                                                       
        return [self.nltk_word_tokenizer.tokenize(sent)                                             
                for sent in self.nltk_sent_tokenizer.tokenize(text)]

tokenizer = Tokenizer()

raw_stream = regenerator.Stream.from_txt('wikipedia.txt')

token_stream = (
    raw_stream
    .map(str.lower)
    .map(str.strip)
    .map(tokenizer)
    .unchunk()
)

# nsent
#len(trimmed_token_stream)

# nword
#len(trimmed_token_stream.unchunk())

class Vocab:
    pad, unk, num = '<pad>', '<unk>', '<num>'
    special = pad, unk, num

    special_idx = range(3)
    pad_idx, unk_idx, num_idx = special_idx

    def __init__(self, words, max_size=None, num_regex=r'^[\−\-\+]?[\$¢€£¥₹₤]?\.?\d+([\.,:/]\d+)*$'):
        self.num_regex = re.compile(num_regex)

        words = filter(lambda word: not self.num_regex.match(word), words)
        
        if max_size:
            counter = collections.Counter(words)
            words = (word for word, count in counter.most_common(max_size - len(self.special)))

        words = sorted(words)

        def idx_factory():
            idx, idx_factory.idx = idx_factory.idx, idx_factory.idx + 1
            return idx
        idx_factory.idx = 0

        vocab = collections.defaultdict(idx_factory)
        for word in itertools.chain(self.special, words):
            vocab[word]

        self.vocab = dict(vocab)

        self.reverse = {idx: word for word, idx in self.vocab.items()}
        
    @classmethod
    def from_corpus(cls, corpus, *args, **kwargs):
        return cls((word for sent in corpus for word in sent), *args, **kwargs)
    
    @classmethod
    def from_file(cls, filename, *args, **kwargs):
        with open(filename, mode='r', encoding='utf8') as fh:
            words = [line.rstrip('\n') for line in fh]
            
        return cls(words, *args, **kwargs)

    def save(self, filename):
        with open(filename, mode='w', encoding='utf8') as fh:
            fh.writelines(word + '\n' for word in self)

    def __getitem__(self, word):
        #if word.isdigit():
        if self.num_regex.match(word):
            return self.num_idx

        return self.vocab.get(word, self.unk_idx)

    def __iter__(self):
        '''Return iterator over words in order of indices.
        '''
        return (word for idx, word in sorted((idx, word) for word, idx in self.vocab.items()))

    def __len__(self):
        return len(self.vocab)

#max_vocab_size = 65_536
max_vocab_size = 32_768

vocab = Vocab.from_corpus(trimmed_token_stream, max_size=max_vocab_size)

vocab.save('global_cbow_vocab.txt')

token_index_stream = trimmed_token_stream.map(
    lambda sent: [vocab[word] for word in sent])

token_rev_index_stream = token_index_stream.map(
    lambda sent: [vocab.reverse[word] for word in sent])

train_index_stream, valid_index_stream = token_index_stream.random_split(frac=0.95)

# len(train_index_stream), len(valid_index_stream)

with open('global_cbow_train_idxs.txt', mode='w') as fh:
    for sent in train_index_stream:
        fh.write(' '.join(map(str, sent)) + '\n')

with open('global_cbow_valid_idxs.txt', mode='w') as fh:
    for sent in valid_index_stream:
        fh.write(' '.join(map(str, sent)) + '\n')

#import time

#start = time.time()
#with open('global_cbow_train_idxs.txt', mode='r') as fh:
#    sents = [[int(word) for word in line.strip().split()] for line in fh]
#time.time() - start

#sents

#sentinel = (-1).to_bytes(4, byteorder='big', signed=True)
#int.from_bytes(sentinel, byteorder='big', signed=False)

#nbytes = 4 # 32 bit max
#sentinel = (-1).to_bytes(nbytes, byteorder='big', signed=True)

#with open('global_cbow_train_idxs.bin', mode='wb') as fh:
#    for sent in train_index_stream:
#        for word in sent:
#            fh.write(word.to_bytes(nbytes, byteorder='big', signed=False))
#    fh.write(sentinel)

#start = time.time()

#with open('global_cbow_train_idxs.bin', mode='rb') as fh:
#    sents = []
#    cur_sent = []

#    while data := fh.read(nbytes):
#        if data == sentinel:
#            sents.append(cur_sent)
#            cur_sent = []
#        else:
#            int_data = int.from_bytes(data, byteorder='big', signed=False)
#            cur_sent.append(int_data)
#time.time() - start

class TokenIndexStream(regenerator.Stream):
    def __init__(self, filename):
        def load_token_indices():
            with open(filename, mode='r') as fh:
                for sent in fh:
                    yield th.tensor([int(word) for word in sent.strip().split()], dtype=th.int64)
            
        super().__init__(load_token_indices)
        
    def buffer_shuffle(self, size=8192, refill_frac=0.5):
        def generator_func():
            it = iter(self)
            split = int(size * refill_frac)

            def refill():
                return list(itertools.islice(it, split))
            
            buff = refill()
            while True:
                buff_new = refill()
                if not buff_new:
                    break

                buff.extend(buff_new)
                random.shuffle(buff)

                for item in buff[split:]:
                    yield item

                buff = buff[:split]
                    
            for item in buff:
                yield item
                    
        return self.from_func(generator_func)
    
    def buffer_sort(self, size=12_000, shuffle=True):
        def generator_func():
            it = iter(self)
            reverse = False

            while True:
                buff = list(itertools.islice(it, size))
                if not buff: break

                if shuffle:
                    random.shuffle(buff)

                buff.sort(key=len, reverse=reverse)
                reverse = not reverse

                for item in buff:
                    yield item

        return self.from_func(generator_func)
    
    def batch(self, size):                                                                     
        def generator_func():                                                                       
            it = iter(self)                                                                         
                                                                                                    
            while True:                                                                             
                batch = list(itertools.islice(it, size))                                            
                if not batch: break                                                                 
                                                                                                    
                yield th.nn.utils.rnn.pad_sequence(                                                 
                    batch,               
                    padding_value=Vocab.pad_idx,                                                    
                    batch_first=True)                                                               
                                                                                                    
        return self.from_func(generator_func)

class TorchDataset(th.utils.data.IterableDataset):
    def __init__(self, stream):
        super().__init__()
        self.stream = stream
        
    def __iter__(self):
        return iter(self.stream)

class ContinuousBagOfWords(th.nn.Module):
    def __init__(self, vocab_size, embed_dim=128, width=5, max_norm=None):
        super().__init__()

        self.vocab_size = vocab_size
        self.embed_dim = embed_dim
        self.width = width
        self.pad_width = self.width // 2

        self.embedding = th.nn.Embedding(
            num_embeddings=self.vocab_size,
            embedding_dim=self.embed_dim,
            padding_idx=Vocab.pad_idx,
            max_norm=max_norm)

        self.center_mask = th.nn.Parameter(
            th.ones(self.width, dtype=th.float32),
            requires_grad=False)
        self.center_mask[self.pad_width] = 0.0

        self.back_proj = th.nn.Linear(
            in_features=self.embed_dim,
            out_features=self.vocab_size)

    def forward(self, idxs):
        assert idxs.ndim == 2
        batch_size = idxs.shape[0]
        max_len = idxs.shape[1]

        word_vectors = self.embedding(idxs)
        assert word_vectors.shape == (batch_size, max_len, self.embed_dim)

        ## if self.pad_method == 'zero':
        padded_word_vectors = thfunc.pad(
            input=word_vectors,
            pad=(0, 0, self.pad_width, self.pad_width),
            value=0.0)

        folded_word_vectors = padded_word_vectors.unfold(1, self.width, 1)
        assert folded_word_vectors.shape == (batch_size, max_len, self.embed_dim, self.width)
        
        #summed_word_vectors = (folded_word_vectors * self.center_mask).sum(-1)
        summed_word_vectors = folded_word_vectors @ self.center_mask
        assert summed_word_vectors.shape == (batch_size, max_len, self.embed_dim)

        preds = self.back_proj(summed_word_vectors)
        assert preds.shape == (batch_size, max_len, self.vocab_size)

        return preds
    
    def save(self, filename):
        with open(filename, mode='wb') as fh:                                                  
            th.save(self.state_dict(), fh)
    
    @th.no_grad()
    def nearest(self, idxs, n=10):
        word_vectors = self.embedding(idxs)

        # FIXME: just pull the weights out of the embedding layer?
        full_embedding = self.embedding(th.arange(1, self.vocab_size))

        dist = th.cdist(word_vectors, full_embedding)

        # plus one because we start after `<pad>`
        nearest_idxs = th.argsort(dist, dim=1) + 1

        return nearest_idxs[:, :n]

    @th.no_grad()
    def nearest_cosine(self, idxs, n=10):
        word_vectors = self.embedding(idxs)

        # FIXME: just pull the weights out of the embedding layer?
        full_embedding = self.embedding(th.arange(1, self.vocab_size))

        dist = th.as_tensor(
            spdist.cdist(
                word_vectors.numpy(),
                full_embedding.numpy(),
                metric='cosine')
            ).abs()

        # plus one because we start after `<pad>`
        nearest_idxs = th.argsort(dist, dim=1) + 1

        return nearest_idxs[:, :n]
    
CBOW = ContinuousBagOfWords

def train(model, train_data, valid_data, max_epochs, learning_rate=0.01, learning_rate_decay=0.95,
          batch_size=64, stop_epochs=10, load_workers=20, use_gpu=True):
    
    def collate(data):
        return th.nn.utils.rnn.pad_sequence(                                                 
            data,               
            padding_value=Vocab.pad_idx,                                                    
            batch_first=True)
    
    train_loader = th.utils.data.DataLoader(
        train_data,
        batch_size=batch_size,
        pin_memory=use_gpu,
        num_workers=load_workers,
        drop_last=False,
        collate_fn=collate,
        shuffle=False)

    valid_loader = th.utils.data.DataLoader(
        valid_data,
        batch_size=batch_size,
        pin_memory=use_gpu,
        num_workers=load_workers,
        drop_last=False,
        collate_fn=collate,
        shuffle=False)

    if use_gpu:
        model.cuda()

    # all parameters with `requires_grad` set
    parameters = list(filter(lambda param: param.requires_grad, model.parameters()))
    
    # cross entropy loss, ignore padding tokens
    loss_func = th.nn.CrossEntropyLoss(
        ignore_index=Vocab.pad_idx,
        reduction='sum')

    # RProp optimizer
    optimizer = th.optim.Rprop(lr=learning_rate, params=parameters)
    #optimizer = th.optim.Adam(lr=learning_rate, params=parameters)
    
    # exponential learning rate decay
    #lr_scheduler = th.optim.lr_scheduler.ExponentialLR(optimizer, learning_rate_decay)

    # performance metrics
    train_losses = []
    valid_losses = []

    # best model state for early stopping
    best_train_loss = float('inf')
    best_valid_loss = float('inf')
    best_epoch = 0
    best_state = None

    # for each training epoch
    for epoch in range(max_epochs):                                                                                                                                                                     
        # enter training mode
        model.train()

        # zero the gradients
        # we accumulate over all batches, so this is per-epoch
        optimizer.zero_grad()

        # training pass
        epoch_train_losses = []
        for i, batch in enumerate(train_loader):
            # zero the gradients
            #optimizer.zero_grad()

            # move batch to requested device
            if use_gpu:
                batch = batch.cuda()

            preds = model(batch)
            loss = loss_func(                                                                       
                preds.reshape(-1, model.vocab_size),                                                
                batch.reshape(-1))

            loss_f = loss.item()
            epoch_train_losses.append(loss_f)

            loss.backward()

        optimizer.step()

        train_losses.append(sum(epoch_train_losses))

        # back to evaluation mode
        model.eval()

        # validation loss
        with th.no_grad():
            epoch_valid_losses = []
            for i, batch in enumerate(valid_loader):
                if use_gpu:
                    batch = batch.cuda()

                preds = model(batch)
                loss = loss_func(
                    preds.reshape(-1, model.vocab_size),
                    batch.reshape(-1))

                loss_f = loss.item()
                epoch_valid_losses.append(loss_f)

        valid_losses.append(sum(epoch_valid_losses))

        if valid_losses[-1] < best_valid_loss:
            best_train_loss = train_losses[-1]
            best_valid_loss = valid_losses[-1]
            best_epoch = epoch
            best_state = copy.deepcopy(model.state_dict())

        print('epoch:', epoch)
        print('best epoch:', best_epoch)
        print('learning rate:', next(iter(optimizer.param_groups))['lr'])
        print('train loss:', train_losses[-1])
        print('valid loss:', valid_losses[-1])
        print('=======')
                
        if (best_epoch + stop_epochs) <= epoch:
            print(f'terminating after {stop_epochs} of no improvement')
            break
            
        #lr_scheduler.step()

    if best_state is None:
        raise RuntimeError('best state not set!')

    model.load_state_dict(best_state)

    if use_gpu:
        model.cpu()

    return munch.munchify({
        'train_losses': train_losses,
        'valid_losses': valid_losses,
        'best_train_loss': best_train_loss,
        'best_valid_loss': best_valid_loss,
        'best_epoch': best_epoch,
    })

cbow = CBOW(vocab_size=len(vocab), embed_dim=32)#, max_norm=math.sqrt(32))

cbow.save('global_cbow.pkl')

emb_mat = cbow.embedding.weight.detach()

means = emb_mat.mean(0)
u, d, v = th.svd(emb_mat - means)

mags = d**2 / (emb_mat.shape[0] - 1.0)
w = v.T / d

pcomps = emb_mat @ w