単語の埋め込み

CONTEXT_SIZE = 2  # 2 words to the left, 2 to the right
raw_text = """We are about to study the idea of a computational process.
Computational processes are abstract beings that inhabit computers.
As they evolve, processes manipulate other abstract things called data.
The evolution of a process is directed by a pattern of rules
called a program. People create programs to direct processes. In effect,
we conjure the spirits of the computer with our spells.""".split()

vocab = set(raw_text)
vocab_size = len(vocab)

word_to_ix = {word: i for i, word in enumerate(vocab)}
ix_to_word = {i: word for i, word in enumerate(vocab)}

data = []
for i in range(CONTEXT_SIZE, len(raw_text) - CONTEXT_SIZE):
    context = (
          [raw_text[i - j - 1] for j in range(CONTEXT_SIZE)]
        + [raw_text[i + j + 1] for j in range(CONTEXT_SIZE)]
    )
    target = raw_text[i]
    data.append((context, target))

print(data[:3])

import torch
import torch.nn as nn

def make_context_vector(context, word_to_ix):
    idxs = [word_to_ix[w] for w in context]
    return torch.tensor(idxs, dtype=torch.long)

class CBOW(nn.Module):
    def __init__(self, vocab_size, embedding_dim):
        super().__init__()
        # out: 1 x emdedding_dim
        self.embeddings = nn.Embedding(vocab_size, embedding_dim)
        self.linear1 = nn.Linear(embedding_dim, 128)
        self.activation_function1 = nn.ReLU()
        # out: 1 x vocab_size
        self.linear2 = nn.Linear(128, vocab_size)
        self.activation_function2 = nn.LogSoftmax(dim = -1)

    def forward(self, inputs):
        embeds = sum(self.embeddings(inputs)).view(1,-1)
        out = self.linear1(embeds)
        out = self.activation_function1(out)
        out = self.linear2(out)
        out = self.activation_function2(out)
        return out

    def get_word_emdedding(self, word):
        word = torch.tensor([word_to_ix[word]])
        return self.embeddings(word).view(1,-1)
    

# set model
model = CBOW(vocab_size, embedding_dim=100)
loss_function = nn.NLLLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

# training
for epoch in range(50):
    total_loss = 0

    for context, target in data:
        context_vector = make_context_vector(context, word_to_ix)  

        log_probs = model(context_vector)

        total_loss += loss_function(log_probs, torch.tensor([word_to_ix[target]]))

    #optimize at the end of each epoch
    optimizer.zero_grad()
    total_loss.backward()
    optimizer.step()

# testing
context = ['People','create', 'to', 'direct']
context_vector = make_context_vector(context, word_to_ix)
a = model(context_vector)

# result
print(f'Raw text: {" ".join(raw_text)}\n')
print(f'Context: {context}\n')
print(f'Prediction: {ix_to_word[torch.argmax(a[0]).item()]}')

model.get_word_emdedding("programs")

# トークンに分割した文章の集合の例
sentences = [
    ["king", "male", "ruler"],
    ["queen", "female", "ruler"],
    ["man", "male"],
    ["woman", "female"],
]

from gensim.models import Word2Vec
model = Word2Vec(sentences=sentences, vector_size=100, window=5, min_count=1, workers=4)

# kingに意味が近い上位3個（similarの基準はコサイン類似度）
sims = model.wv.most_similar('king', topn=3)
sims

# king + woman - man ≒ queen というアレ
model.wv.most_similar(positive=['woman', 'king'], negative=['man'], topn=1)

# 自前で足し引きして類似度計算してみる
import numpy as np
def cosine_sim(a, b):
    return a @ b / (np.linalg.norm(a) * np.linalg.norm(b))

v = model.wv["king"] - model.wv["man"] + model.wv["woman"]
words = set(sum(sentences, [])) - {"king", "man", "woman"} # 計算に使ったものは必然的に類似度が高くなっちゃうので除く
for word in words:
    print(f"{word}: {cosine_sim(v, model.wv[word]):.3g}")

# queenが一番近くなった