pytorchで作るword2vec

Tweet

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader

# テキストデータ
corpus = [
    'I like to play football',
    'I love chocolate',
    'Football is my favorite sport',
    'I enjoy playing video games'
]

# テキストを単語に分割する
tokenized_corpus = [sentence.lower().split() for sentence in corpus]

# ユニークな単語のリストを作成する
vocabulary = list(set([word for sentence in tokenized_corpus for word in sentence]))

# 単語をインデックスにマッピングする辞書を作成する
word_to_idx = {word: i for i, word in enumerate(vocabulary)}
idx_to_word = {i: word for i, word in enumerate(vocabulary)}

# ハイパーパラメータ
embedding_dim = 10
window_size = 2
learning_rate = 0.01
num_epochs = 100
batch_size = 4


# CBOW用のデータセットクラス
class CBOWDataset(Dataset):
    def __init__(self, corpus, word_to_idx, window_size):
        self.data = []
        for sentence in corpus:
            indices = [word_to_idx[word] for word in sentence]
            for i in range(window_size, len(indices) - window_size):
                context = indices[i - window_size:i] + indices[i + 1:i + window_size + 1]
                target = indices[i]
                self.data.append((context, target))

    def __getitem__(self, index):
        context, target = self.data[index]
        return torch.tensor(context), torch.tensor(target)

    def __len__(self):
        return len(self.data)


# CBOWモデルの定義
class CBOW(nn.Module):
    def __init__(self, vocab_size, embedding_dim):
        super(CBOW, self).__init__()
        self.embeddings = nn.Embedding(vocab_size, embedding_dim)
        self.linear = nn.Linear(embedding_dim, vocab_size)

    def forward(self, inputs):
        embeds = self.embeddings(inputs).sum(dim=1)
        out = self.linear(embeds)
        return out


# データローダーの作成
dataset = CBOWDataset(tokenized_corpus, word_to_idx, window_size)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# モデルの初期化
model = CBOW(len(vocabulary), embedding_dim)

# 損失関数と最適化手法の定義
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate)

# 学習ループ
for epoch in range(num_epochs):
    epoch_loss = 0
    for context, target in dataloader:
        optimizer.zero_grad()
        output = model(context)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
    print('Epoch {}, Loss: {:.4f}'.format(epoch + 1, epoch_loss))

# 単語の埋め込みを取得する
embedding_weights = model.embeddings.weight.detach().numpy()

# 埋め込みを表示する
for i, word in enumerate(vocabulary):
    print('Word: {}, Embedding: {}'.format(word, embedding_weights[i]))