File size: 1,744 Bytes
1783e08 | 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 | import torch
import torch.nn as nn
import torch.optim as optim
from datasets import Dataset, load_dataset, concatenate_datasets
from tokenizers import Tokenizer
from huggingface_hub import PyTorchModelHubMixin, snapshot_download
tokenizer = Tokenizer.from_pretrained("LocalWisdom/PurpleGloVe")
# from train.ipynb
class GloVe(nn.Module, PyTorchModelHubMixin):
def __init__(self, config):
nn.Module.__init__(self)
self.context_window = config["context_window"]
self.x_max = config["x_max"]
self.alpha = config["alpha"]
self.vocab_size = config["vocab_size"]
self.embedding_dim = config["embedding_dim"]
self.wi = nn.Embedding(self.vocab_size, self.embedding_dim)
self.wj = nn.Embedding(self.vocab_size, self.embedding_dim)
self.bi = nn.Embedding(self.vocab_size, 1)
self.bj = nn.Embedding(self.vocab_size, 1)
def f(self, Xij):
return torch.clamp((Xij / self.x_max) ** self.alpha, max=1)
def forward(self, i, j, Xij):
wi = self.wi(i)
wj = self.wj(j)
bi = self.bi(i).squeeze(-1)
bj = self.bj(j).squeeze(-1)
weighting = self.f(Xij)
dot = (wi * wj).sum(dim=-1)
epsilon = 1e-8
loss = weighting * (dot + bi + bj - torch.log(Xij + epsilon)) ** 2
return torch.mean(loss)
# added new embed function
def embed(self, word_id):
# I had a mistake by naming the subscript of vector W and B with i and j
# which is misleading since both is subscript i but one is target word
# vector and the other is context word vector
with torch.no_grad():
word_id = torch.tensor([word_id])
return self.wi(word_id) + self.wj(word_id)
PurpleGloVe = GloVe.from_pretrained("LocalWisdom/PurpleGloVe")
word_str = "king"
word = tokenizer.encode(word_str).ids[0]
word_e = PurpleGloVe.embed(word)
print(word_e)
|