Fix handling of input and output embeddings in TFXGLM

Author: damianoamatruda

src/transformers/modeling_tf_utils.py

@@ -2069,6 +2069,12 @@ def _get_word_embedding_weight(model, embedding_layer):
             return embedding_layer
         # Otherwise, try to get them from the layer's attributes
 
+        embeds = getattr(embedding_layer, "kernel", None)
+        if embeds is not None:
+            # in a dense layer the kernel has a shape (last_dim, units), for us (dim, num_tokens)
+            # value has a shape (num_tokens, dim) then needs to be transposed
+            return tf.Variable(tf.transpose(embeds))
+
         embeds = getattr(embedding_layer, "weight", None)
         if embeds is not None:
             return embeds
@@ -2082,6 +2088,12 @@ def _get_word_embedding_weight(model, embedding_layer):
         # the argument after building the model
         model.build_in_name_scope()
 
+        embeds = getattr(embedding_layer, "kernel", None)
+        if embeds is not None:
+            # in a dense layer the kernel has a shape (last_dim, units), for us (dim, num_tokens)
+            # value has a shape (num_tokens, dim) then needs to be transposed
+            return tf.Variable(tf.transpose(embeds))
+
         embeds = getattr(embedding_layer, "weight", None)
         if embeds is not None:
             return embeds

src/transformers/models/xglm/modeling_tf_xglm.py

@@ -490,8 +490,9 @@ def __init__(
     def get_input_embeddings(self) -> TFSharedEmbeddings:
         return self.embed_tokens
 
-    def set_input_embeddings(self, value: TFSharedEmbeddings) -> None:
-        self.embed_tokens = value
+    def set_input_embeddings(self, value) -> None:
+        self.embed_tokens.vocab_size = value.shape[0]
+        self.embed_tokens.weight = value
 
     def _prepare_decoder_attention_mask(
         self,
@@ -888,8 +889,17 @@ def __init__(
     def get_output_embeddings(self):
         return self.lm_head
 
-    def set_output_embeddings(self, new_embeddings):
-        self.lm_head = new_embeddings
+    def set_output_embeddings(self, value):
+        self.lm_head = keras.layers.Dense(
+            shape_list(value)[0],
+            use_bias=False,
+            kernel_initializer=get_initializer(self.config.init_std),
+            name="lm_head",
+        )
+        # in a dense layer the kernel has a shape (last_dim, units), for us (dim, num_tokens)
+        # value has a shape (num_tokens, dim) then needs to be transposed
+        transposed_value = tf.transpose(value)
+        self.lm_head.kernel = transposed_value
 
     def prepare_inputs_for_generation(self, inputs, past_key_values=None, use_cache=None, **kwargs):
         # only last token for inputs_ids if past is defined in kwargs