Update installation constraints for streaming (#1661)

* Update installation constraints for streaming * Add type:ignores for pyright issues
mosaicml · Oct 22, 2022 · f5f9571 · f5f9571
1 parent 44fbd62
commit f5f9571
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 14 deletions.
diff --git a/composer/algorithms/alibi/attention_surgery_functions/_gpt2.py b/composer/algorithms/alibi/attention_surgery_functions/_gpt2.py
@@ -27,10 +27,11 @@ def gpt2_attention_converter(module: torch.nn.Module, module_index: int, max_seq
 
     assert isinstance(module, GPT2Attention)
     del module_index  # unused
-    module = register_alibi(module=module,
-                            n_heads=int(module.num_heads),
-                            max_token_length=max_sequence_length,
-                            causal=True)
+    module = register_alibi(
+        module=module,
+        n_heads=int(module.num_heads),  #type: ignore num_heads member of GPT2Attention
+        max_token_length=max_sequence_length,
+        causal=True)
     setattr(module, '_attn', MethodType(_attn, module))
 
     module = enlarge_mask(module, max_sequence_length)

diff --git a/composer/algorithms/gated_linear_units/gated_linear_units.py b/composer/algorithms/gated_linear_units/gated_linear_units.py
@@ -38,13 +38,14 @@ def from_BertOutput(layer: torch.nn.Module,
     assert isinstance(
         layer, BertOutput
     ), 'The replacement policy requires an instance of transformers.models.bert.modeling_bert.BertOutput for the necessary fields to be defined.'
-    return BERTGatedFFOutput(d_embed=layer.dense.out_features,
-                             d_ff=layer.dense.in_features,
-                             dropout_rate=layer.dropout.p,
-                             act_fn=act_fn,
-                             layernorm_eps=layer.LayerNorm.eps,
-                             gated_layer_bias=gated_layer_bias,
-                             non_gated_layer_bias=non_gated_layer_bias)
+    return BERTGatedFFOutput(
+        d_embed=layer.dense.out_features,  #type: ignore dense.out_features member of BertOutput
+        d_ff=layer.dense.in_features,  #type: ignore dense.in_features member of BertOutput
+        dropout_rate=layer.dropout.p,  #type: ignore dropout.p member of BertOutput
+        act_fn=act_fn,
+        layernorm_eps=layer.LayerNorm.eps,  #type: ignore LayerNorm.eps member of BertOutput
+        gated_layer_bias=gated_layer_bias,
+        non_gated_layer_bias=non_gated_layer_bias)
 
 
 def from_BertIntermediate(layer: torch.nn.Module, module_index: int) -> torch.nn.Identity:
@@ -108,7 +109,7 @@ def apply_gated_linear_units(model: torch.nn.Module,
                              'specify `act_fn` to manually override activation functions.')
 
         # since our set is of 1, let's extract the only activation function remaining.
-        (act_fn,) = act_fns
+        (act_fn,) = act_fns  #type: ignore will fail below if None
 
         if act_fn is None:
             raise ValueError(

diff --git a/composer/datasets/synthetic_lm.py b/composer/datasets/synthetic_lm.py
@@ -185,7 +185,8 @@ def generate_synthetic_tokenizer(tokenizer_family: str,
             raise ValueError(f'{tokenizer_params.tokenizer_cls} should sub-class transformers.PreTrainedTokenizer.')
         # print("Temporary path:", tmp_path)
         # input("Waiting for input..")
-        tokenizer = tokenizer_params.tokenizer_cls.from_pretrained(tmp_tokenizer_dir)
+        tokenizer = tokenizer_params.tokenizer_cls.from_pretrained(  #type: ignore
+            tmp_tokenizer_dir)
 
     return tokenizer
 

diff --git a/setup.py b/setup.py
@@ -181,7 +181,7 @@ def package_files(prefix: str, directory: str, extension: str):
 ]
 
 extra_deps['streaming'] = [
-    'mosaicml-streaming<=0.1.*',
+    'mosaicml-streaming<0.2.*',
     'boto3>=1.21.45,<2',
     'paramiko>=2.11.0,<3',
 ]