Merge branch 'master' into type_docs_tests

examples/nlp/glue_benchmark/glue_benchmark_with_bert.py

@@ -61,9 +61,10 @@
 """
 
 import argparse
-import json
 import os
 
+from transformers import BertConfig
+
 import nemo.collections.nlp as nemo_nlp
 import nemo.core as nemo_core
 from nemo import logging
@@ -187,6 +188,7 @@
     add_time_to_log_dir=True,
 )
 
+
 if args.bert_checkpoint is None:
     """ Use this if you're using a standard BERT model.
     To see the list of pretrained models, call:
@@ -205,14 +207,31 @@
         tokenizer = NemoBertTokenizer(args.pretrained_bert_model)
     else:
         raise ValueError(f"received unexpected tokenizer '{args.tokenizer}'")
+
     if args.bert_config is not None:
-        with open(args.bert_config) as json_file:
-            config = json.load(json_file)
-        model = nemo_nlp.nm.trainables.huggingface.BERT(**config)
+        config = BertConfig.from_json_file(args.bert_config).to_dict()
+        args.vocab_size = config['vocab_size']
+        args.hidden_size = config['hidden_size']
+        args.num_hidden_layers = config['num_hidden_layers']
+        args.num_attention_heads = config['num_attention_heads']
+        args.intermediate_size = config['intermediate_size']
+        args.hidden_act = config['hidden_act']
+        args.max_seq_length = config['max_position_embeddings']
+
+        model = nemo_nlp.nm.trainables.huggingface.BERT(
+            vocab_size=args.vocab_size,
+            num_hidden_layers=args.num_hidden_layers,
+            hidden_size=args.hidden_size,
+            num_attention_heads=args.num_attention_heads,
+            intermediate_size=args.intermediate_size,
+            max_position_embeddings=args.max_seq_length,
+            hidden_act=args.hidden_act,
+        )
+        logging.info(f"using {args.bert_config}")
     else:
         model = nemo_nlp.nm.trainables.huggingface.BERT(pretrained_model_name=args.pretrained_bert_model)
-
     model.restore_from(args.bert_checkpoint)
+    logging.info(f"model resotred from {args.bert_checkpoint}")
 
 hidden_size = model.hidden_size
 

nemo/core/neural_types/axes.py

@@ -34,7 +34,7 @@ class AxisKind(AxisKindAbstract):
     """This Enum represents what does varying axis dimension mean.
     For example, does this dimension correspond to width, batch, time, etc.
     The "Dimension" and "Channel" kinds are the same and used to represent
-    a general axis.
+    a general axis. "Any" axis will accept any axis kind fed to it.
     """
 
     Batch = 0
@@ -43,6 +43,7 @@ class AxisKind(AxisKindAbstract):
     Channel = 2
     Width = 3
     Height = 4
+    Any = 5
 
     def __repr__(self):
         return self.__str__()
@@ -64,6 +65,8 @@ def from_str(label):
             return AxisKind.Width
         elif _label == "h" or _label == "height":
             return AxisKind.Height
+        elif _label == "any":
+            return AxisKind.Any
         else:
             raise ValueError(f"Can't create AxisKind from {label}")
 

nemo/core/neural_types/neural_type.py

@@ -186,7 +186,9 @@ def __compare_axes(axes_a, axes_b) -> int:
         for axis_a, axis_b in zip(axes_a, axes_b):
             kinds_a[axis_a.kind] = axis_a.size
             kinds_b[axis_b.kind] = axis_b.size
-            if (
+            if axis_a.kind == AxisKind.Any:
+                same = True
+            elif (
                 axis_a.kind != axis_b.kind
                 or axis_a.is_list != axis_b.is_list
                 or (axis_a.size != axis_b.size and axis_a.size is not None)

tests/core/test_neural_types.py

@@ -176,3 +176,13 @@ def test_unspecified_dimensions(self):
         t1 = NeuralType(('B', 'T', 'C'), SpectrogramType())
         self.assertEqual(t1.compare(t0), NeuralTypeComparisonResult.SAME)
         self.assertEqual(t0.compare(t1), NeuralTypeComparisonResult.DIM_INCOMPATIBLE)
+
+    def test_any_axis(self):
+        t0 = NeuralType(('B', 'Any', 'Any'), VoidType())
+        t1 = NeuralType(('B', 'Any', 'Any'), SpectrogramType())
+        t2 = NeuralType(('B', 'T', 'C'), SpectrogramType())
+        self.assertEqual(t0.compare(t1), NeuralTypeComparisonResult.SAME)
+        self.assertEqual(t0.compare(t2), NeuralTypeComparisonResult.SAME)
+        self.assertEqual(t1.compare(t2), NeuralTypeComparisonResult.SAME)
+        self.assertEqual(t2.compare(t1), NeuralTypeComparisonResult.INCOMPATIBLE)
+        self.assertEqual(t1.compare(t0), NeuralTypeComparisonResult.INCOMPATIBLE)