Address review comments

gluonnlp/base.py

@@ -0,0 +1,30 @@
+# coding: utf-8
+
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# pylint: disable=abstract-method
+"""Helper functions."""
+
+__all__ = ['_str_types']
+
+import sys
+
+if sys.version_info[0] == 3:
+    _str_types = (str, )
+else:
+    _str_types = (str, unicode)

gluonnlp/data/candidate_sampler.py

@@ -21,7 +21,6 @@
 __all__ = ['CandidateSampler', 'UnigramCandidateSampler']
 
 import mxnet as mx
-import numpy as np
 
 
 class CandidateSampler(object):

gluonnlp/model/train/embedding.py

@@ -24,18 +24,13 @@
 
 import logging
 import struct
-import sys
 
 import numpy as np
 from mxnet import cpu, nd
 from mxnet.gluon import Block, HybridBlock, nn
 
 from ...vocab import create_subword_function
-
-if sys.version_info[0] == 3:
-    _str_types = (str, )
-else:
-    _str_types = (str, unicode)
+from ...base import _str_types
 
 
 class EmbeddingModel(Block):
@@ -110,23 +105,27 @@ class SimpleEmbeddingModel(EmbeddingModel, Block):
         Initializer for the embeddings matrix.
     sparse_grad : bool, default True
         Specifies mxnet.gluon.nn.Embedding sparse_grad argument.
+    dtype : str, default 'float32'
+        dtype argument passed to gluon.nn.Embedding
 
     """
 
     def __init__(self, token_to_idx, embedding_size, weight_initializer=None,
-                 sparse_grad=True, **kwargs):
+                 sparse_grad=True, dtype='float32', **kwargs):
         assert isinstance(token_to_idx, dict)
 
         super(SimpleEmbeddingModel,
               self).__init__(embedding_size=embedding_size, **kwargs)
         self.token_to_idx = token_to_idx
         self.weight_initializer = weight_initializer
         self.sparse_grad = sparse_grad
+        self.dtype = dtype
 
         with self.name_scope():
             self.embedding = nn.Embedding(
                 len(token_to_idx), embedding_size,
-                weight_initializer=weight_initializer, sparse_grad=sparse_grad)
+                weight_initializer=weight_initializer, sparse_grad=sparse_grad,
+                dtype=dtype)
 
     def __call__(self, words, wordsmask=None):
         return super(SimpleEmbeddingModel, self).__call__(words, wordsmask)
@@ -139,7 +138,7 @@ def forward(self, words, wordsmask=None):
         words : mx.nd.NDArray
             Array of token indices.
         wordsmask : mx.nd.NDArray
-            Mask for embeddings returend by the word level embedding operator.
+            Mask for embeddings returned by the word level embedding operator.
 
         """
         #pylint: disable=arguments-differ
@@ -186,20 +185,19 @@ def __getitem__(self, tokens):
 
 class _MaskedSumEmbedding(HybridBlock):
     def __init__(self, num_tokens, embedding_size, weight_initializer=None,
-                 sparse_grad=True, **kwargs):
+                 sparse_grad=True, dtype='float32', **kwargs):
         super(_MaskedSumEmbedding, self).__init__(**kwargs)
         self.num_tokens = num_tokens
         self.embedding_size = embedding_size
         self.weight_initializer = weight_initializer
         self.sparse_grad = sparse_grad
+        self.dtype = dtype
 
         with self.name_scope():
             self.embedding = nn.Embedding(
-                num_tokens,
-                embedding_size,
-                weight_initializer=weight_initializer,
-                sparse_grad=sparse_grad,
-            )
+                num_tokens, embedding_size,
+                weight_initializer=weight_initializer, sparse_grad=sparse_grad,
+                dtype=dtype)
 
     def hybrid_forward(self, F, x, mask):
         #pylint: disable=arguments-differ
@@ -208,7 +206,7 @@ def hybrid_forward(self, F, x, mask):
         return F.sum(masked_embeddings, axis=-2)
 
 
-class FasttextEmbeddingModel(EmbeddingModel, Block):
+class FasttextEmbeddingModel(EmbeddingModel):
     """FastText embedding model.
 
     The FasttextEmbeddingModel combines a word level embedding matrix and a
@@ -234,32 +232,32 @@ class FasttextEmbeddingModel(EmbeddingModel, Block):
         Initializer for the embeddings and subword embeddings matrix.
     sparse_grad : bool, default True
         Specifies mxnet.gluon.nn.Embedding sparse_grad argument.
+    dtype : str, default 'float32'
+        dtype argument passed to gluon.nn.Embedding
 
     """
     FASTTEXT_FILEFORMAT_MAGIC = 793712314
 
     def __init__(self, token_to_idx, subword_function, embedding_size,
-                 weight_initializer=None, sparse_grad=True, **kwargs):
+                 weight_initializer=None, sparse_grad=True, dtype='float32',
+                 **kwargs):
         super(FasttextEmbeddingModel,
               self).__init__(embedding_size=embedding_size, **kwargs)
         self.token_to_idx = token_to_idx
         self.subword_function = subword_function
         self.weight_initializer = weight_initializer
         self.sparse_grad = sparse_grad
+        self.dtype = dtype
 
         with self.name_scope():
             self.embedding = nn.Embedding(
-                len(token_to_idx),
-                embedding_size,
-                weight_initializer=weight_initializer,
-                sparse_grad=sparse_grad,
-            )
+                len(token_to_idx), embedding_size,
+                weight_initializer=weight_initializer, sparse_grad=sparse_grad,
+                dtype=dtype)
             self.subword_embedding = _MaskedSumEmbedding(
-                len(subword_function),
-                embedding_size,
-                weight_initializer=weight_initializer,
-                sparse_grad=sparse_grad,
-            )
+                len(subword_function), embedding_size,
+                weight_initializer=weight_initializer, sparse_grad=sparse_grad,
+                dtype=dtype)
 
     @classmethod
     def load_fasttext_format(cls, path, ctx=cpu(), **kwargs):
@@ -451,7 +449,7 @@ def forward(self, words, subwords, wordsmask=None, subwordsmask=None,
             subwords of the unique tokens in `words` with
             `words_to_unique_subwords_indices` containing the reverse mapping.
         wordsmask : mx.nd.NDArray, optional
-            Mask for embeddings returend by the word level embedding operator.
+            Mask for embeddings returned by the word level embedding operator.
         subwordsmask : mx.nd.NDArray, optional
             A mask for the subword embeddings looked up from `subwords`.
             Applied before sum reducing the subword embeddings.

scripts/word_embeddings/candidate_sampler.py

@@ -18,13 +18,11 @@
 # under the License.
 """Candidate samplers"""
 
-__all__ = ['UnigramCandidateSampler']
+__all__ = ['remove_accidental_hits']
 
 import mxnet as mx
 import numpy as np
 
-import gluonnlp as nlp
-
 try:
     from numba import njit
     numba_njit = njit(nogil=True)
@@ -70,5 +68,3 @@ def remove_accidental_hits(candidates, true_samples, true_samples_mask=None):
         _candidates_mask(candidates_np, true_samples_np, true_samples_mask_np))
 
     return candidates, candidates_mask.as_in_context(candidates.context)
-
-