20 KiB
20 KiB
In [1]:
import os
import torch
from torch import nn
from d2l import torch as d2lIn [2]:
#@save
d2l.DATA_HUB['glove.6b.50d'] = (d2l.DATA_URL + 'glove.6B.50d.zip',
'0b8703943ccdb6eb788e6f091b8946e82231bc4d')
#@save
d2l.DATA_HUB['glove.6b.100d'] = (d2l.DATA_URL + 'glove.6B.100d.zip',
'cd43bfb07e44e6f27cbcc7bc9ae3d80284fdaf5a')
#@save
d2l.DATA_HUB['glove.42b.300d'] = (d2l.DATA_URL + 'glove.42B.300d.zip',
'b5116e234e9eb9076672cfeabf5469f3eec904fa')
#@save
d2l.DATA_HUB['wiki.en'] = (d2l.DATA_URL + 'wiki.en.zip',
'c1816da3821ae9f43899be655002f6c723e91b88')In [3]:
#@save
class TokenEmbedding:
"""Token Embedding."""
def __init__(self, embedding_name):
self.idx_to_token, self.idx_to_vec = self._load_embedding(
embedding_name)
self.unknown_idx = 0
self.token_to_idx = {token: idx for idx, token in
enumerate(self.idx_to_token)}
def _load_embedding(self, embedding_name):
idx_to_token, idx_to_vec = ['<unk>'], []
data_dir = d2l.download_extract(embedding_name)
# GloVe website: https://nlp.stanford.edu/projects/glove/
# fastText website: https://fasttext.cc/
with open(os.path.join(data_dir, 'vec.txt'), 'r') as f:
for line in f:
elems = line.rstrip().split(' ')
token, elems = elems[0], [float(elem) for elem in elems[1:]]
# Skip header information, such as the top row in fastText
if len(elems) > 1:
idx_to_token.append(token)
idx_to_vec.append(elems)
idx_to_vec = [[0] * len(idx_to_vec[0])] + idx_to_vec
return idx_to_token, torch.tensor(idx_to_vec)
def __getitem__(self, tokens):
indices = [self.token_to_idx.get(token, self.unknown_idx)
for token in tokens]
vecs = self.idx_to_vec[torch.tensor(indices)]
return vecs
def __len__(self):
return len(self.idx_to_token)In [4]:
glove_6b50d = TokenEmbedding('glove.6b.50d')Downloading ../data/glove.6B.50d.zip from http://d2l-data.s3-accelerate.amazonaws.com/glove.6B.50d.zip...
In [5]:
len(glove_6b50d)Out [5]:
400001
In [6]:
glove_6b50d.token_to_idx['beautiful'], glove_6b50d.idx_to_token[3367]Out [6]:
(3367, 'beautiful')
In [7]:
def knn(W, x, k):
# Add 1e-9 for numerical stability
cos = torch.mv(W, x.reshape(-1,)) / (
torch.sqrt(torch.sum(W * W, axis=1) + 1e-9) *
torch.sqrt((x * x).sum()))
_, topk = torch.topk(cos, k=k)
return topk, [cos[int(i)] for i in topk]In [8]:
def get_similar_tokens(query_token, k, embed):
topk, cos = knn(embed.idx_to_vec, embed[[query_token]], k + 1)
for i, c in zip(topk[1:], cos[1:]): # Exclude the input word
print(f'cosine sim={float(c):.3f}: {embed.idx_to_token[int(i)]}')In [9]:
get_similar_tokens('chip', 3, glove_6b50d)cosine sim=0.856: chips cosine sim=0.749: intel cosine sim=0.749: electronics
In [10]:
get_similar_tokens('baby', 3, glove_6b50d)cosine sim=0.839: babies cosine sim=0.800: boy cosine sim=0.792: girl
In [11]:
get_similar_tokens('beautiful', 3, glove_6b50d)cosine sim=0.921: lovely cosine sim=0.893: gorgeous cosine sim=0.830: wonderful
In [12]:
def get_analogy(token_a, token_b, token_c, embed):
vecs = embed[[token_a, token_b, token_c]]
x = vecs[1] - vecs[0] + vecs[2]
topk, cos = knn(embed.idx_to_vec, x, 1)
return embed.idx_to_token[int(topk[0])] # Remove unknown wordsIn [13]:
get_analogy('man', 'woman', 'son', glove_6b50d)Out [13]:
'daughter'
In [14]:
get_analogy('beijing', 'china', 'tokyo', glove_6b50d)Out [14]:
'japan'
In [15]:
get_analogy('bad', 'worst', 'big', glove_6b50d)Out [15]:
'biggest'
In [16]:
get_analogy('do', 'did', 'go', glove_6b50d)Out [16]:
'went'