example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from vocab import Vocab\n",
    "vocab = Vocab(['__BOS__', '__EOS__', '__UNK__', 'Hello', ',', 'world', 'i', 'am', 'a', 'vocab'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0, 3, 4, 5, 1]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab.tokenize(\"Hello , world\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'__BOS__ Hello , world __EOS__'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab.detokenize([0, 3, 4, 5, 1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'__BOS__ Hello , world __EOS__'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab.detokenize([0, 3, 4, 5, 1, 1, 1, 1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['__BOS__', 'Hello', ',', 'world', '__EOS__']"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab.detokenize([0, 3, 4, 5, 1], sep=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "inp_voc = Vocab.from_file(\"1.voc\")\n",
    "out_voc = Vocab.from_file(\"2.voc\")\n",
    "\n",
    "#in case you want to merge vocabularies\n",
    "joint_voc = Vocab.merge(inp_voc, out_voc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(32733, 32772, 64199)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# len of joint voc is slightly less because of duplicate tokens\n",
    "len(inp_voc), len(out_voc), len(joint_voc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[    0 10299   117 16873 23655 25369  8083  2771 29363 28273 29689  8083\n",
      "  16867 20770  4623 29363 19190  1155 16910 14624 24234 15585 31058     1]\n",
      " [    0 21398 22529 26581 25369 20347  1047 26342 27240 31058     1     1\n",
      "      1     1     1     1     1     1     1     1     1     1     1     1]\n",
      " [    0 30248 20578 14482 28847 10803 27539   501 15870 17060  3392 31476\n",
      "   5387 19686     1     1     1     1     1     1     1     1     1     1]\n",
      " [    0 32758 26886 20481  3544 23655  9626     1     1     1     1     1\n",
      "      1     1     1     1     1     1     1     1     1     1     1     1]]\n"
     ]
    }
   ],
   "source": [
    "batch_ix = out_voc.tokenize_many(open(\"./bpe_corpus2.txt\").read(300).split('\\n'))\n",
    "print(batch_ix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['__BOS__ С заселением проблем не было , хотя и приехали раньше , нас угостили кофе и заселили сразу же после уборки номера . __EOS__',\n",
       " '__BOS__ Дополни@@ тельное одеяло было предоставлено по первой просьбе . __EOS__',\n",
       " '__BOS__ Основ@@ ным критерием выбора была закрытая парков@@ ка-@@ здесь она во дворе ! __EOS__',\n",
       " '__BOS__ На обратном пути остановку не планировали __EOS__']"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out_voc.detokenize_many(batch_ix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['С заселением проблем не было , хотя и приехали раньше , нас угостили кофе и заселили сразу же после уборки номера .',\n",
       " 'Дополнительное одеяло было предоставлено по первой просьбе .',\n",
       " 'Основным критерием выбора была закрытая парковка-здесь она во дворе !',\n",
       " 'На обратном пути остановку не планировали']"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out_voc.detokenize_many(batch_ix, deprocess=True, unbpe=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}

vocab.py

import numpy as np


class Vocab:
    """
    Vocab converts between strings, tokens and token indices.
    It should normally be treated as immutable
    It should be saved with pickle.dump
    """
    _default_tokens = ("__BOS__", "__EOS__", "__UNK__")
    remove_bpe = lambda s: s.replace('@@ ', '')

    def __init__(self, tokens):
        tokens = tuple(tokens)
        assert len(tokens) == len(set(tokens)), "tokens must be unique"
        for i, t in enumerate(self._default_tokens):
            assert t in tokens and tokens.index(t) == i, "token must have %s at index %i" % (t,i)

        self.tokens = tokens
        self.token2id = {token: i for i, token in enumerate(self.tokens)}
        self.BOS = 0
        self.EOS = 1
        self.UNK = 2

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

    def tokenize(self, sentence, separator=' '):
        """ Converts sentence into a sequence of ids """
        if isinstance(sentence, str):
            sentence = sentence.split(separator)

        sentence = list(sentence)
        if "__EOS__" not in sentence:
            sentence.append("__EOS__")
        if sentence[0] != "__BOS__":
            sentence.insert(0, "__BOS__")

        return [self.token2id.get(token, self.UNK) for token in sentence]

    def detokenize(self, indices, crop=True, sep=' ', unbpe=False, deprocess=False):
        """ converts indices to words. If separator is not None, joins them over it """
        indices = tuple(indices)
        if self.EOS in indices:
            indices = indices[:indices.index(self.EOS) + 1]

        tokens = [self.tokens[token] for token in indices]
        if deprocess:
            tokens = [t for t in tokens if t not in self._default_tokens]

        if sep is None:
            return tokens
        else:
            line = sep.join(tokens)
            if unbpe:
                line = self.remove_bpe(line)
            return line

    def tokenize_many(self, lines, max_len=None, sep=' '):
        """
        convert variable length token sequences into fixed size matrix
        pads short sequences with self.EOS
        example usage:
        >>>print(vocab.tokenize_many(sentences[:3]))
        [[15 22 21 28 27 13  1  1  1  1  1]
         [30 21 15 15 21 14 28 27 13  1  1]
         [25 37 31 34 21 20 37 21 28 19 13]]
        """
        max_len = max_len or max(map(lambda s: len(s.split(sep)), lines)) + 2  # 2 for bos and eos

        matrix = np.zeros((len(lines), max_len), dtype='int32') + self.EOS
        for i, seq in enumerate(lines):
            tokens = self.tokenize(seq)[:max_len]
            matrix[i, :len(tokens)] = tokens

        return matrix

    def detokenize_many(self, matrix, crop=True, sep=' ', unbpe=False, deprocess=False):
        """
        Convert matrix of token ids into strings
        :param matrix: matrix of tokens of int32, shape=[batch,time]
        :param crop: if True, crops BOS and EOS from line
        :param sep: if not None, joins tokens over that string
        :param unbpe: if True, merges BPE into words
        :param deprocess: if True, removes all unknowns
        :return: a list of strings of
        """
        return [self.detokenize(sent, crop, sep, unbpe, deprocess) for sent in matrix]

    @classmethod
    def from_file(cls, voc_path):
        """ Parses vocab from a .voc file """
        tokens = set()
        with open(voc_path, 'r') as f:
            for line in f:
                token = line.split(" ")[0]
                tokens.update([token])

        return Vocab(list(cls._default_tokens) + list(tokens))

    @classmethod
    def from_sequences(cls, sentences, separator=' '):
        """ Infers tokens from a corpora of sentences (tokens separated by separator) """
        tokens = set()
        for s in sentences:
            tokens.update(s.split(separator))
        return Vocab(list(cls._default_tokens) + sorted(tokens))

    @classmethod
    def merge(cls, *vocabs):
        for vocab in vocabs:
            assert isinstance(vocab, Vocab)

        # all tokens excluding special
        tokens = {token for vocab in vocabs
                  for token in vocab.tokens
                  if token not in cls._default_tokens}

        return Vocab(list(cls._default_tokens) + sorted(tokens))