Source code for torchtext.vocab

from collections import defaultdict
from functools import partial
import logging
import os
import zipfile
import gzip

from urllib.request import urlretrieve
import torch
from tqdm import tqdm
import tarfile

from .utils import reporthook

from collections import Counter

logger = logging.getLogger(__name__)

[docs]class Vocab(object): """Defines a vocabulary object that will be used to numericalize a field. Attributes: freqs: A collections.Counter object holding the frequencies of tokens in the data used to build the Vocab. stoi: A collections.defaultdict instance mapping token strings to numerical identifiers. itos: A list of token strings indexed by their numerical identifiers. """ # TODO (@mttk): Populate classs with default values of special symbols UNK = '<unk>'
[docs] def __init__(self, counter, max_size=None, min_freq=1, specials=('<unk>', '<pad>'), vectors=None, unk_init=None, vectors_cache=None, specials_first=True): """Create a Vocab object from a collections.Counter. Args: counter: collections.Counter object holding the frequencies of each value found in the data. max_size: The maximum size of the vocabulary, or None for no maximum. Default: None. min_freq: The minimum frequency needed to include a token in the vocabulary. Values less than 1 will be set to 1. Default: 1. specials: The list of special tokens (e.g., padding or eos) that will be prepended to the vocabulary. Default: ['<unk'>, '<pad>'] vectors: One of either the available pretrained vectors or custom pretrained vectors (see Vocab.load_vectors); or a list of aforementioned vectors unk_init (callback): by default, initialize out-of-vocabulary word vectors to zero vectors; can be any function that takes in a Tensor and returns a Tensor of the same size. Default: 'torch.zeros' vectors_cache: directory for cached vectors. Default: '.vector_cache' specials_first: Whether to add special tokens into the vocabulary at first. If it is False, they are added into the vocabulary at last. Default: True. """ self.freqs = counter counter = counter.copy() min_freq = max(min_freq, 1) self.itos = list() self.unk_index = None if specials_first: self.itos = list(specials) # only extend max size if specials are prepended max_size = None if max_size is None else max_size + len(specials) # frequencies of special tokens are not counted when building vocabulary # in frequency order for tok in specials: del counter[tok] # sort by frequency, then alphabetically words_and_frequencies = sorted(counter.items(), key=lambda tup: tup[0]) words_and_frequencies.sort(key=lambda tup: tup[1], reverse=True) for word, freq in words_and_frequencies: if freq < min_freq or len(self.itos) == max_size: break self.itos.append(word) if Vocab.UNK in specials: # hard-coded for now unk_index = specials.index(Vocab.UNK) # position in list # account for ordering of specials, set variable self.unk_index = unk_index if specials_first else len(self.itos) + unk_index self.stoi = defaultdict(self._default_unk_index) else: self.stoi = defaultdict() if not specials_first: self.itos.extend(list(specials)) # stoi is simply a reverse dict for itos self.stoi.update({tok: i for i, tok in enumerate(self.itos)}) self.vectors = None if vectors is not None: self.load_vectors(vectors, unk_init=unk_init, cache=vectors_cache) else: assert unk_init is None and vectors_cache is None
def _default_unk_index(self): return self.unk_index def __getitem__(self, token): return self.stoi.get(token, self.stoi.get(Vocab.UNK)) def __getstate__(self): # avoid picking defaultdict attrs = dict(self.__dict__) # cast to regular dict attrs['stoi'] = dict(self.stoi) return attrs def __setstate__(self, state): if state.get("unk_index", None) is None: stoi = defaultdict() else: stoi = defaultdict(self._default_unk_index) stoi.update(state['stoi']) state['stoi'] = stoi self.__dict__.update(state) def __eq__(self, other): if self.freqs != other.freqs: return False if self.stoi != other.stoi: return False if self.itos != other.itos: return False if self.vectors != other.vectors: return False return True def __len__(self): return len(self.itos) def lookup_indices(self, tokens): indices = [self.__getitem__(token) for token in tokens] return indices def extend(self, v, sort=False): words = sorted(v.itos) if sort else v.itos for w in words: if w not in self.stoi: self.itos.append(w) self.stoi[w] = len(self.itos) - 1
[docs] def load_vectors(self, vectors, **kwargs): """ Args: vectors: one of or a list containing instantiations of the GloVe, CharNGram, or Vectors classes. Alternatively, one of or a list of available pretrained vectors: charngram.100d fasttext.en.300d fasttext.simple.300d glove.42B.300d glove.840B.300d glove.twitter.27B.25d glove.twitter.27B.50d glove.twitter.27B.100d glove.twitter.27B.200d glove.6B.50d glove.6B.100d glove.6B.200d glove.6B.300d Remaining keyword arguments: Passed to the constructor of Vectors classes. """ if not isinstance(vectors, list): vectors = [vectors] for idx, vector in enumerate(vectors): if isinstance(vector, str): # Convert the string pretrained vector identifier # to a Vectors object if vector not in pretrained_aliases: raise ValueError( "Got string input vector {}, but allowed pretrained " "vectors are {}".format( vector, list(pretrained_aliases.keys()))) vectors[idx] = pretrained_aliases[vector](**kwargs) elif not isinstance(vector, Vectors): raise ValueError( "Got input vectors of type {}, expected str or " "Vectors object".format(type(vector))) tot_dim = sum(v.dim for v in vectors) self.vectors = torch.Tensor(len(self), tot_dim) for i, token in enumerate(self.itos): start_dim = 0 for v in vectors: end_dim = start_dim + v.dim self.vectors[i][start_dim:end_dim] = v[token.strip()] start_dim = end_dim assert(start_dim == tot_dim)
[docs] def set_vectors(self, stoi, vectors, dim, unk_init=torch.Tensor.zero_): """ Set the vectors for the Vocab instance from a collection of Tensors. Args: stoi: A dictionary of string to the index of the associated vector in the `vectors` input argument. vectors: An indexed iterable (or other structure supporting __getitem__) that given an input index, returns a FloatTensor representing the vector for the token associated with the index. For example, vector[stoi["string"]] should return the vector for "string". dim: The dimensionality of the vectors. unk_init (callback): by default, initialize out-of-vocabulary word vectors to zero vectors; can be any function that takes in a Tensor and returns a Tensor of the same size. Default: 'torch.zeros' """ self.vectors = torch.Tensor(len(self), dim) for i, token in enumerate(self.itos): wv_index = stoi.get(token, None) if wv_index is not None: self.vectors[i] = vectors[wv_index] else: self.vectors[i] = unk_init(self.vectors[i])
[docs]class SubwordVocab(Vocab):
[docs] def __init__(self, counter, max_size=None, specials=('<pad>'), vectors=None, unk_init=torch.Tensor.zero_): """Create a revtok subword vocabulary from a collections.Counter. Args: counter: collections.Counter object holding the frequencies of each word found in the data. max_size: The maximum size of the subword vocabulary, or None for no maximum. Default: None. specials: The list of special tokens (e.g., padding or eos) that will be prepended to the vocabulary in addition to an <unk> token. vectors: One of either the available pretrained vectors or custom pretrained vectors (see Vocab.load_vectors); or a list of aforementioned vectors unk_init (callback): by default, initialize out-of-vocabulary word vectors to zero vectors; can be any function that takes in a Tensor and returns a Tensor of the same size. Default: 'torch.zeros """ try: import revtok except ImportError: print("Please install revtok.") raise # Hardcode unk_index as subword_vocab has no specials_first argument self.unk_index = (specials.index(SubwordVocab.UNK) if SubwordVocab.UNK in specials else None) if self.unk_index is None: self.stoi = defaultdict() else: self.stoi = defaultdict(self._default_unk_index) self.stoi.update({tok: i for i, tok in enumerate(specials)}) self.itos = specials.copy() self.segment = revtok.SubwordSegmenter(counter, max_size) max_size = None if max_size is None else max_size + len(self.itos) # sort by frequency/entropy, then alphabetically toks = sorted(self.segment.vocab.items(), key=lambda tup: (len(tup[0]) != 1, -tup[1], tup[0])) for tok, _ in toks: if len(self.itos) == max_size: break self.itos.append(tok) self.stoi[tok] = len(self.itos) - 1 if vectors is not None: self.load_vectors(vectors, unk_init=unk_init)
def _infer_shape(f): num_lines, vector_dim = 0, None for line in f: if vector_dim is None: row = line.rstrip().split(b" ") vector = row[1:] # Assuming word, [vector] format if len(vector) > 2: # The header present in some (w2v) formats contains two elements. vector_dim = len(vector) num_lines += 1 # First element read else: num_lines += 1 return num_lines, vector_dim
[docs]class Vectors(object):
[docs] def __init__(self, name, cache=None, url=None, unk_init=None, max_vectors=None): """ Args: name: name of the file that contains the vectors cache: directory for cached vectors url: url for download if vectors not found in cache unk_init (callback): by default, initialize out-of-vocabulary word vectors to zero vectors; can be any function that takes in a Tensor and returns a Tensor of the same size max_vectors (int): this can be used to limit the number of pre-trained vectors loaded. Most pre-trained vector sets are sorted in the descending order of word frequency. Thus, in situations where the entire set doesn't fit in memory, or is not needed for another reason, passing `max_vectors` can limit the size of the loaded set. """ cache = '.vector_cache' if cache is None else cache self.itos = None self.stoi = None self.vectors = None self.dim = None self.unk_init = torch.Tensor.zero_ if unk_init is None else unk_init self.cache(name, cache, url=url, max_vectors=max_vectors)
def __getitem__(self, token): if token in self.stoi: return self.vectors[self.stoi[token]] else: return self.unk_init(torch.Tensor(self.dim)) def cache(self, name, cache, url=None, max_vectors=None): import ssl ssl._create_default_https_context = ssl._create_unverified_context if os.path.isfile(name): path = name if max_vectors: file_suffix = '_{}.pt'.format(max_vectors) else: file_suffix = '.pt' path_pt = os.path.join(cache, os.path.basename(name)) + file_suffix else: path = os.path.join(cache, name) if max_vectors: file_suffix = '_{}.pt'.format(max_vectors) else: file_suffix = '.pt' path_pt = path + file_suffix if not os.path.isfile(path_pt): if not os.path.isfile(path) and url:'Downloading vectors from {}'.format(url)) if not os.path.exists(cache): os.makedirs(cache) dest = os.path.join(cache, os.path.basename(url)) if not os.path.isfile(dest): with tqdm(unit='B', unit_scale=True, miniters=1, desc=dest) as t: try: urlretrieve(url, dest, reporthook=reporthook(t)) except KeyboardInterrupt as e: # remove the partial zip file os.remove(dest) raise e'Extracting vectors into {}'.format(cache)) ext = os.path.splitext(dest)[1][1:] if ext == 'zip': with zipfile.ZipFile(dest, "r") as zf: zf.extractall(cache) elif ext == 'gz': if dest.endswith('.tar.gz'): with, 'r:gz') as tar: tar.extractall(path=cache) if not os.path.isfile(path): raise RuntimeError('no vectors found at {}'.format(path))"Loading vectors from {}".format(path)) ext = os.path.splitext(path)[1][1:] if ext == 'gz': open_file = else: open_file = open vectors_loaded = 0 with open_file(path, 'rb') as f: num_lines, dim = _infer_shape(f) if not max_vectors or max_vectors > num_lines: max_vectors = num_lines itos, vectors, dim = [], torch.zeros((max_vectors, dim)), None for line in tqdm(f, total=max_vectors): # Explicitly splitting on " " is important, so we don't # get rid of Unicode non-breaking spaces in the vectors. entries = line.rstrip().split(b" ") word, entries = entries[0], entries[1:] if dim is None and len(entries) > 1: dim = len(entries) elif len(entries) == 1: logger.warning("Skipping token {} with 1-dimensional " "vector {}; likely a header".format(word, entries)) continue elif dim != len(entries): raise RuntimeError( "Vector for token {} has {} dimensions, but previously " "read vectors have {} dimensions. All vectors must have " "the same number of dimensions.".format(word, len(entries), dim)) try: if isinstance(word, bytes): word = word.decode('utf-8') except UnicodeDecodeError:"Skipping non-UTF8 token {}".format(repr(word))) continue vectors[vectors_loaded] = torch.tensor([float(x) for x in entries]) vectors_loaded += 1 itos.append(word) if vectors_loaded == max_vectors: break self.itos = itos self.stoi = {word: i for i, word in enumerate(itos)} self.vectors = torch.Tensor(vectors).view(-1, dim) self.dim = dim'Saving vectors to {}'.format(path_pt)) if not os.path.exists(cache): os.makedirs(cache), self.stoi, self.vectors, self.dim), path_pt) else:'Loading vectors from {}'.format(path_pt)) self.itos, self.stoi, self.vectors, self.dim = torch.load(path_pt) def __len__(self): return len(self.vectors)
[docs] def get_vecs_by_tokens(self, tokens, lower_case_backup=False): """Look up embedding vectors of tokens. Args: tokens: a token or a list of tokens. if `tokens` is a string, returns a 1-D tensor of shape `self.dim`; if `tokens` is a list of strings, returns a 2-D tensor of shape=(len(tokens), self.dim). lower_case_backup : Whether to look up the token in the lower case. If False, each token in the original case will be looked up; if True, each token in the original case will be looked up first, if not found in the keys of the property `stoi`, the token in the lower case will be looked up. Default: False. Examples: >>> examples = ['chip', 'baby', 'Beautiful'] >>> vec = text.vocab.GloVe(name='6B', dim=50) >>> ret = vec.get_vecs_by_tokens(tokens, lower_case_backup=True) """ to_reduce = False if not isinstance(tokens, list): tokens = [tokens] to_reduce = True if not lower_case_backup: indices = [self[token] for token in tokens] else: indices = [self[token] if token in self.stoi else self[token.lower()] for token in tokens] vecs = torch.stack(indices) return vecs[0] if to_reduce else vecs
[docs]class GloVe(Vectors): url = { '42B': '', '840B': '', 'twitter.27B': '', '6B': '', } def __init__(self, name='840B', dim=300, **kwargs): url = self.url[name] name = 'glove.{}.{}d.txt'.format(name, str(dim)) super(GloVe, self).__init__(name, url=url, **kwargs)
[docs]class FastText(Vectors): url_base = '{}.vec' def __init__(self, language="en", **kwargs): url = self.url_base.format(language) name = os.path.basename(url) super(FastText, self).__init__(name, url=url, **kwargs)
[docs]class CharNGram(Vectors): name = 'charNgram.txt' url = ('' 'jmt_pre-trained_embeddings.tar.gz') def __init__(self, **kwargs): super(CharNGram, self).__init__(, url=self.url, **kwargs) def __getitem__(self, token): vector = torch.Tensor(1, self.dim).zero_() if token == "<unk>": return self.unk_init(vector) chars = ['#BEGIN#'] + list(token) + ['#END#'] num_vectors = 0 for n in [2, 3, 4]: end = len(chars) - n + 1 grams = [chars[i:(i + n)] for i in range(end)] for gram in grams: gram_key = '{}gram-{}'.format(n, ''.join(gram)) if gram_key in self.stoi: vector += self.vectors[self.stoi[gram_key]] num_vectors += 1 if num_vectors > 0: vector /= num_vectors else: vector = self.unk_init(vector) return vector
pretrained_aliases = { "charngram.100d": partial(CharNGram), "fasttext.en.300d": partial(FastText, language="en"), "fasttext.simple.300d": partial(FastText, language="simple"), "glove.42B.300d": partial(GloVe, name="42B", dim="300"), "glove.840B.300d": partial(GloVe, name="840B", dim="300"), "glove.twitter.27B.25d": partial(GloVe, name="twitter.27B", dim="25"), "glove.twitter.27B.50d": partial(GloVe, name="twitter.27B", dim="50"), "glove.twitter.27B.100d": partial(GloVe, name="twitter.27B", dim="100"), "glove.twitter.27B.200d": partial(GloVe, name="twitter.27B", dim="200"), "glove.6B.50d": partial(GloVe, name="6B", dim="50"), "glove.6B.100d": partial(GloVe, name="6B", dim="100"), "glove.6B.200d": partial(GloVe, name="6B", dim="200"), "glove.6B.300d": partial(GloVe, name="6B", dim="300") } """Mapping from string name to factory function"""
[docs]def build_vocab_from_iterator(iterator, num_lines=None): """ Build a Vocab from an iterator. Args: iterator: Iterator used to build Vocab. Must yield list or iterator of tokens. num_lines: The expected number of elements returned by the iterator. (Default: None) Optionally, if known, the expected number of elements can be passed to this factory function for improved progress reporting. """ counter = Counter() with tqdm(unit_scale=0, unit='lines', total=num_lines) as t: for tokens in iterator: counter.update(tokens) t.update(1) word_vocab = Vocab(counter) return word_vocab


Access comprehensive developer documentation for PyTorch

View Docs


Get in-depth tutorials for beginners and advanced developers

View Tutorials


Find development resources and get your questions answered

View Resources