Update lm.py to suit OPT-like models

Some newer models like OPT treat embedding layers as a separate layer 'Embedding' instead of pure tensor.
jalammar · Jul 22, 2023 · aab5b7d · aab5b7d
1 parent 0ef61c5
commit aab5b7d
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/ecco/lm.py b/src/ecco/lm.py
@@ -86,7 +86,10 @@ def __init__(self,
             self.model_type = self.model_config['type']
             embeddings_layer_name = self.model_config['embedding']
             embed_retriever = attrgetter(embeddings_layer_name)
-            self.model_embeddings = embed_retriever(self.model)
+            if type(embed_retriever(self.model)) == torch.nn.Embedding:
+                self.model_embeddings = embed_retriever(self.model).weight
+            else:
+                self.model_embeddings = embed_retriever(self.model)
             self.collect_activations_layer_name_sig = self.model_config['activations'][0]
         except KeyError:
             raise ValueError(