Function train_model

nlp_class2/word2vec_theano.py:83–274

(savedir)

def train_model(savedir):
  # get the data
  sentences, word2idx = get_wiki() #get_text8()


  # number of unique words
  vocab_size = len(word2idx)


  # config
  window_size = 5
  learning_rate = 0.025*128
  final_learning_rate = 0.0001*128
  num_negatives = 5 # number of negative samples to draw per input word
  samples_per_epoch = int(1e5)
  epochs = 1
  D = 50 # word embedding size


  # learning rate decay
  learning_rate_delta = (learning_rate - final_learning_rate) / epochs
  # learning_rate_delta = 0


  # params
  W = np.random.randn(vocab_size, D) / np.sqrt(D + vocab_size) # input-to-hidden
  V = np.random.randn(D, vocab_size) / np.sqrt(D + vocab_size) # hidden-to-output


  # theano variables
  thW = theano.shared(W)
  thV = theano.shared(V)

  # theano placeholders
  th_pos_word   = T.ivector('pos_word')
  th_neg_word   = T.ivector('neg_word')
  th_context    = T.ivector('context')
  th_lr         = T.scalar('learning_rate')

  # get the output and loss
  input_words = T.concatenate([th_pos_word, th_neg_word])
  W_subset = thW[input_words]
  dbl_context = T.concatenate([th_context, th_context])
  V_subset = thV[:, dbl_context]
  logits = W_subset.dot(V_subset)
  out = T.nnet.sigmoid(logits)

  n = th_pos_word.shape[0]
  th_cost = -T.log(out[:n]).mean() - T.log(1 - out[n:]).mean()


  # specify the updates
  gW = T.grad(th_cost, W_subset)
  gV = T.grad(th_cost, V_subset)
  W_update = T.inc_subtensor(W_subset, -th_lr*gW)
  V_update = T.inc_subtensor(V_subset, -th_lr*gV)
  updates = [(thW, W_update), (thV, V_update)]