text: Title: Exporting probabilities over the learned vocabulary
  Body: Currently, the `extract_features.py` file supports extracting representations before the last output layer but I want to extract the final probabilities over the vocabulary. I modified the code with the following addition but it doesn't seem to work (the probabilities I'm getting are all very small).
  
  ```
  model_output = model.get_sequence_output()
  
  batch_size = tf.shape(model_output)[0]
  length = tf.shape(model_output)[1]
  hidden_size = tf.shape(model_output)[2]
  
  model_output = tf.reshape(inputs, [-1, hidden_size])
  logits = tf.matmul(model_output, model.get_embedding_table(), transpose_b=True)
  probs = tf.nn.softmax(logits, axis=-1)
  probs = tf.reshape(probs, [batch_size, length, bert_config.vocab_size])
  
  ```