|
import abc |
|
import functools |
|
import os |
|
import time |
|
|
|
import bpemb |
|
|
|
import torch |
|
import torchtext |
|
|
|
|
|
from seq2struct.utils import registry |
|
|
|
|
|
class Embedder(metaclass=abc.ABCMeta): |
|
|
|
@abc.abstractmethod |
|
def tokenize(self, sentence): |
|
'''Given a string, return a list of tokens suitable for lookup.''' |
|
pass |
|
|
|
@abc.abstractmethod |
|
def untokenize(self, tokens): |
|
'''Undo tokenize.''' |
|
pass |
|
|
|
@abc.abstractmethod |
|
def lookup(self, token): |
|
'''Given a token, return a vector embedding if token is in vocabulary. |
|
|
|
If token is not in the vocabulary, then return None.''' |
|
pass |
|
|
|
@abc.abstractmethod |
|
def contains(self, token): |
|
pass |
|
|
|
@abc.abstractmethod |
|
def to(self, device): |
|
'''Transfer the pretrained embeddings to the given device.''' |
|
pass |
|
|
|
|
|
@registry.register('word_emb', 'glove') |
|
class GloVe(Embedder): |
|
|
|
def __init__(self, kind, lemmatize=False): |
|
cache = os.path.join(os.environ.get('CACHE_DIR', os.getcwd()), '.vector_cache') |
|
self.glove = torchtext.vocab.GloVe(name=kind, cache=cache) |
|
self.dim = self.glove.dim |
|
self.vectors = self.glove.vectors |
|
self.lemmatize = lemmatize |
|
self.corenlp_annotators = ['tokenize', 'ssplit'] |
|
if lemmatize: |
|
self.corenlp_annotators.append('lemma') |
|
|
|
@functools.lru_cache(maxsize=1024) |
|
def tokenize(self, text): |
|
ann = corenlp.annotate(text, self.corenlp_annotators) |
|
if self.lemmatize: |
|
return [tok.lemma.lower() for sent in ann.sentence for tok in sent.token] |
|
else: |
|
return [tok.word.lower() for sent in ann.sentence for tok in sent.token] |
|
|
|
@functools.lru_cache(maxsize=1024) |
|
def tokenize_for_copying(self, text): |
|
ann = corenlp.annotate(text, self.corenlp_annotators) |
|
text_for_copying = [tok.originalText.lower() for sent in ann.sentence for tok in sent.token] |
|
if self.lemmatize: |
|
text = [tok.lemma.lower() for sent in ann.sentence for tok in sent.token] |
|
else: |
|
text = [tok.word.lower() for sent in ann.sentence for tok in sent.token] |
|
return text, text_for_copying |
|
|
|
def untokenize(self, tokens): |
|
return ' '.join(tokens) |
|
|
|
def lookup(self, token): |
|
i = self.glove.stoi.get(token) |
|
if i is None: |
|
return None |
|
return self.vectors[i] |
|
|
|
def contains(self, token): |
|
return token in self.glove.stoi |
|
|
|
def to(self, device): |
|
self.vectors = self.vectors.to(device) |
|
|
|
|
|
@registry.register('word_emb', 'bpemb') |
|
class BPEmb(Embedder): |
|
def __init__(self, dim, vocab_size, lang='en'): |
|
self.bpemb = bpemb.BPEmb(lang=lang, dim=dim, vs=vocab_size) |
|
self.dim = dim |
|
self.vectors = torch.from_numpy(self.bpemb.vectors) |
|
|
|
def tokenize(self, text): |
|
return self.bpemb.encode(text) |
|
|
|
def untokenize(self, tokens): |
|
return self.bpemb.decode(tokens) |
|
|
|
def lookup(self, token): |
|
i = self.bpemb.spm.PieceToId(token) |
|
if i == self.bpemb.spm.unk_id(): |
|
return None |
|
return self.vectors[i] |
|
|
|
def contains(self, token): |
|
return self.lookup(token) is not None |
|
|
|
def to(self, device): |
|
self.vectors = self.vectors.to(device) |
|
|
