Set correct nhid for HF TF models

diagnnose/models/import_model.py

@@ -30,12 +30,11 @@ def import_model(*args, **kwargs) -> LanguageModel:
 
 def _import_transformer_lm(*args, **kwargs):
     """ Imports a Transformer LM. """
-    from .fairseq_lm import FairseqLM
-    from .huggingface_lm import HuggingfaceLM
-
     if kwargs["transformer_type"] == "fairseq":
+        from .fairseq_lm import FairseqLM
         return FairseqLM(*args, **kwargs)
 
+    from .huggingface_lm import HuggingfaceLM
     return HuggingfaceLM(*args, **kwargs)
 
 

diagnnose/models/transformer_lm.py

@@ -280,10 +280,14 @@ def top_layer(self) -> int:
         return -1
 
     def nhid(self, activation_name: ActivationName) -> int:
-        if activation_name[1] == "out":
-            return self.pretrained_model.config.vocab_size
+        model_config = self.pretrained_model.config
 
-        return self.pretrained_model.config.hidden_size
+        if activation_name[1] == "out":
+            return model_config.vocab_size
+        elif hasattr(model_config, "word_embed_proj_dim"):
+            return model_config.word_embed_proj_dim
+        else:
+            return self.pretrained_model.config.hidden_size
 
     @staticmethod
     def activation_names() -> ActivationNames:

requirements.txt

@@ -5,4 +5,3 @@ transformers>=4.0.0
 tqdm
 dill
 unidecode
-fairseq