Pushing type fix and adding more information in case of connection error

Author: shubham-s-agarwal

medcat/meta_cat.py

@@ -384,7 +384,7 @@ def save(self, save_dir_path: str) -> None:
 
         if self.config.model.model_name == 'bert':
             model_config_save_path = os.path.join(save_dir_path, 'bert_config.json')
-            self.model.bert_config.to_json_file(model_config_save_path)
+            self.model.bert_config.to_json_file(model_config_save_path) # type: ignore
         # This is everything we need to save from the class, we do not
         # save the class itself.
 

medcat/utils/meta_cat/models.py

@@ -2,7 +2,7 @@
 from collections import OrderedDict
 from typing import Optional, Any, List, Iterable
 from torch import nn, Tensor
-from transformers import BertModel, AutoConfig, BertConfig
+from transformers import BertModel, AutoConfig
 from medcat.meta_cat import ConfigMetaCAT
 import logging
 logger = logging.getLogger(__name__)
@@ -87,30 +87,35 @@ def forward(self,
 class BertForMetaAnnotation(nn.Module):
     _keys_to_ignore_on_load_unexpected: List[str] = [r"pooler"]  # type: ignore
 
-    def __init__(self, config, save_dir_path = None):
+    def __init__(self, config, save_dir_path=None):
         super(BertForMetaAnnotation, self).__init__()
         if save_dir_path:
             try:
                 _bertconfig = AutoConfig.from_pretrained(save_dir_path + "/bert_config.json",
                                                      num_hidden_layers=config.model['num_layers'])
-            except:
+            except Exception:
                 _bertconfig = AutoConfig.from_pretrained(config.model.model_variant,
                                                          num_hidden_layers=config.model['num_layers'])
+                logger.info("BERT config not found locally — downloaded successfully from Hugging Face.")
+
         else:
             _bertconfig = AutoConfig.from_pretrained(config.model.model_variant,num_hidden_layers=config.model['num_layers'])
 
         if config.model['input_size'] != _bertconfig.hidden_size:
             logger.warning("Input size for %s model should be %d, provided input size is %d. Input size changed to %d",config.model.model_variant,_bertconfig.hidden_size,config.model['input_size'],_bertconfig.hidden_size)
 
         if config.model['load_bert_pretrained_weights']:
-            bert = BertModel.from_pretrained(config.model.model_variant, config=_bertconfig)
+            try:
+                bert = BertModel.from_pretrained(config.model.model_variant, config=_bertconfig)
+            except Exception:
+                raise Exception("Could not load BERT pretrained weights from Hugging Face. \nIf you're seeing a connection error, set `config.model.load_bert_pretrained_weights=False` and make sure to load the model pack from disk instead.")
         else:
             bert = BertModel(_bertconfig)
 
         self.config = config
         self.config.use_return_dict = False
         self.bert = bert
-        self.bert_config: BertConfig = _bertconfig
+        self.bert_config = _bertconfig
         self.num_labels = config.model["nclasses"]
         for param in self.bert.parameters():
             param.requires_grad = not config.model.model_freeze_layers