tinyms-ai · leonwanghui · May 8, 2021 · Mar 31, 2021 · Mar 31, 2021 · Mar 31, 2021
diff --git a/tests/st/bert/bert_for_classfier.py → tests/st/bert/bert_for_classifier.py b/tests/st/bert/bert_for_classfier.py → tests/st/bert/bert_for_classifier.py
diff --git a/tests/st/bert/bert_for_pre_pretraining.py → tests/st/bert/bert_for_pretraining.py b/tests/st/bert/bert_for_pre_pretraining.py → tests/st/bert/bert_for_pretraining.py
diff --git a/tinyms/callbacks.py b/tinyms/callbacks.py
@@ -93,7 +93,6 @@ class BertLossCallBack(Callback):
         None
 
     Examples:
-        >>> from tinyms import Tensor
         >>> from tinyms.callbacks import BertLossCallBack
         >>>
         >>> BertLossCallBack(dataset_size=1)

diff --git a/tinyms/data/loader.py b/tinyms/data/loader.py
@@ -33,8 +33,10 @@
 from .utils import generate_image_list, load_img
 
 
-__all__ = ['UnalignedDataset', 'GanImageFolderDataset', 'ImdbDataset',
-           'DistributedSampler', 'BertDataset']
+common_dataset = ['UnalignedDataset', 'GanImageFolderDataset', 'ImdbDataset', 'BertDataset']
+common_sampler = ['DistributedSampler']
+
+__all__ = common_dataset + common_sampler
 __all__.extend(engine.__all__)
 
 random.seed(1)

diff --git a/tinyms/model/bert/__init__.py b/tinyms/model/bert/__init__.py
@@ -14,7 +14,7 @@
 # ============================================================================
 """Bert Init."""
 
-from .bert_for_pre_training import BertNetworkWithLoss, BertPreTraining, \
+from .bert_for_pretraining import BertNetworkWithLoss, BertPreTraining, \
     BertPretrainingLoss, GetMaskedLMOutput, GetNextSentenceOutput, \
     BertTrainOneStepCell, BertTrainOneStepWithLossScaleCell, \
     BertTrainAccumulationAllReduceEachWithLossScaleLayer, \

diff --git a/tinyms/model/bert/bert_for_finetune.py b/tinyms/model/bert/bert_for_finetune.py
@@ -22,7 +22,7 @@
 from tinyms import context
 from tinyms import layers
 from tinyms import primitives as P
-from .bert_for_pre_training import clip_grad
+from .bert_for_pretraining import clip_grad
 from .finetune_eval_model import BertCLSModel, BertNERModel, BertSquadModel
 
 

diff --git a/tinyms/model/bert/bert_for_pre_training.py → tinyms/model/bert/bert_for_pretraining.py b/tinyms/model/bert/bert_for_pre_training.py → tinyms/model/bert/bert_for_pretraining.py
diff --git a/tinyms/text/transforms.py b/tinyms/text/transforms.py
@@ -16,6 +16,7 @@
 
 from . import _transform_ops
 from ._transform_ops import *
+from ..data import *
 
 
 __all__ = [
@@ -33,6 +34,8 @@ def __init__(self):
 
     def apply_ds(self, data_set, batch_size):
 
+        assert isinstance(data_set, BertDataset), "For BertDatasetTransform, BertDataset is needed"
+
         type_cast_op = TypeCast(ts.int32)
         data_set = data_set.map(operations=type_cast_op, input_columns="masked_lm_ids")
         data_set = data_set.map(operations=type_cast_op, input_columns="masked_lm_positions")