Merge pull request #99 from BiEchi/patch-1

Update lm.py to suit OPT-like models
jalammar · Aug 10, 2023 · e292773 · e292773
2 parents 78abb38 + aab5b7d
commit e292773
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(