File size: 3,342 Bytes
d758c99 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
import abc
import functools
import os
import time
import bpemb
#import corenlp
import torch
import torchtext
#from seq2struct.resources import corenlp
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)
|