add OPT embedding attribute

i-machine-think · Mar 30, 2023 · f5d5f07 · f5d5f07
1 parent b06d692
commit f5d5f07
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/diagnnose/models/huggingface_lm.py b/diagnnose/models/huggingface_lm.py
@@ -74,6 +74,9 @@ def embeddings(self) -> Callable[[Tensor], Tensor]:
             elif hasattr(base_model, "word_emb"):
                 # Transformer-XL
                 return base_model.word_emb
+            elif hasattr(base_model, "decoder"):
+                # OPT
+                return base_model.decoder.embed_tokens
             else:
                 raise AttributeError("word embedding attribute not found")