word2vec/CBOW.py at master · limengning/word2vec · GitHub

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
60
61
62
63
64
65
66
67
import paddle
from paddle.fluid.layers.nn import embedding, pad
import paddle.nn
from paddle.nn import Embedding
import paddle.nn.functional as F
import paddle.optimizer
import numpy as np

import load


class CBOW(paddle.nn.Layer):
    def __init__(self, vocab_size, embedding_size, init_scale=0.5):
        super(CBOW, self).__init__()
        self.vocab_size = vocab_size
        self.embedding_size = embedding_size

        self.embedding = Embedding(num_embeddings=self.vocab_size,
                                   embedding_dim=self.embedding_size, weight_attr=paddle.nn.initializer.Uniform(
                                       low=-init_scale, high=init_scale
                                   ))
        self.embedding_out = Embedding(num_embeddings=self.vocab_size,
                                       embedding_dim=self.embedding_size, weight_attr=paddle.nn.initializer.Uniform(
                                           low=-init_scale, high=init_scale
                                       ))

    def forward(self, context_words, center_words, label):
        context_words_emb = self.embedding(context_words)
        center_words_emb = self.embedding_out(center_words)

        word_sim = paddle.multiply(context_words_emb, center_words_emb)
        word_sim = paddle.sum(word_sim, axis=-1)
        word_sim = paddle.reshape(word_sim, shape=[-1])
        pred = F.sigmoid(word_sim)

        loss = F.binary_cross_entropy_with_logits(word_sim, label)
        loss = paddle.mean(loss)
        return pred, loss


batch_size = 64
epoch_num = 1
embedding_size = 200
step = 0
learning_rate = 0.001


# paddle.set_device('gpu:0')

cbow_model = CBOW(load.vocab_size, embedding_size)

adam = paddle.optimizer.Adam(
    learning_rate=learning_rate, parameters=cbow_model.parameters())

for context_words, center_words, label in load.build_batch(load.cbow_dataset, batch_size, epoch_num):
    context_tensor = paddle.to_tensor(context_words)
    center_tensor = paddle.to_tensor(center_words)
    label_tensor = paddle.to_tensor(label)
    pre, loss = cbow_model(context_tensor, center_tensor, label_tensor)
    loss.backward()
    adam.step()
    adam.clear_grad()
    step += 1
    if step % 10000 == 0:
        load.get_similar_tokens('movie', 5, cbow_model.embedding.weight)
        load.get_similar_tokens('one', 5, cbow_model.embedding.weight)
        load.get_similar_tokens('chip', 5, cbow_model.embedding.weight)