bugfix: Add attention mask for generation

README.md

@@ -76,7 +76,7 @@ model_generation = T5ForConditionalGeneration.from_pretrained(modelcard_generati
 del model_generation.encoder  # we only need the decoder for generation. Deleting the encoder is optional, but saves memory.
 model = AutoModel.from_pretrained(modelcard, trust_remote_code=True, revision='main')
 tokenizer = AutoTokenizer.from_pretrained(modelcard)
-
+model_generation.shared = model.shared  # share embeddings between encoder and decoder. This mimics the T5 architecture.
 
 print('get dummy input (2 instances to show batching)')
 graph_1 = [
@@ -100,17 +100,17 @@ how = 'global'  # can be 'global' or 'local', depending on whether the local or
 data_1 = model.data_processor.encode_graph(tokenizer=tokenizer, g=graph_1, text=text_1, how=how)
 data_2 = model.data_processor.encode_graph(tokenizer=tokenizer, g=graph_2, text=text_2, how=how)
 datas = [data_1, data_2]
-model_inputs = model.data_processor.to_batch(data_instances=datas, tokenizer=tokenizer, max_seq_len=None, device='cpu')
+model_inputs, attention_mask = model.data_processor.to_batch(data_instances=datas, tokenizer=tokenizer, max_seq_len=None, device='cpu', return_attention_mask=True)
 
 print('compute token encodings')
 outputs = model(**model_inputs)
 
 print('generate conditional on encoded graph and text')
-outputs = model_generation.generate(encoder_outputs=outputs, max_new_tokens=10)
+outputs = model_generation.generate(encoder_outputs=outputs, max_new_tokens=10, attention_mask=attention_mask)
 print('generation 1:', tokenizer.decode(outputs[0], skip_special_tokens=True))
 print('generation 2:', tokenizer.decode(outputs[1], skip_special_tokens=False)) 
 ```
-
+Note that the embedding to map from the vocabulary to T5's hidden dimension is shared by the encoder and the decoder in T5. To mimic the T5 architecture, we run `model_generation.shared = model.shared` after loading the models. For inference this has no effect, since the embeddings are not updated during inference. However, during training / finetuning, the embeddings can become different for the encoder and decoder if they are not shared. 
 
 ## Contact
 More information can be found in our paper [Graph Language Models](https://arxiv.org/abs/2401.07105) or our [GitHub repository](https://github.com/Heidelberg-NLP/GraphLanguageModels). 
@@ -129,4 +129,5 @@ If this model is helpful for your work, please consider citing the paper:
 }
 ```
 
-
+## Acknowledgments
+Many thanks to Moritz Blum for his help on the generation part. 

wrapper_functions.py

@@ -416,13 +416,14 @@ class DataProcessor():
         return data
     
     @staticmethod
-    def to_batch(data_instances:list[Data], tokenizer, max_seq_len:Optional[int]=None, device:str='cpu', **kwargs)->dict: 
+    def to_batch(data_instances:list[Data], tokenizer, max_seq_len:Optional[int]=None, device:str='cpu', return_attention_mask:bool=False, **kwargs)->dict: 
         """
         converts list of data instances to batched inputs for GLM forward call. 
         :param datas: list of Data instances
-        :param max_seq_len: maximum sequence length
         :param tokenizer: tokenizer
+        :param max_seq_len: maximum sequence length
         :param device: device
+        :param return_attention_mask: whether to return attention mask. The attention mask is not used by the GLM encoder, but the decoder needs it to mask out padding tokens in cross attention. 
         :return: dictionary with keys 'input_ids', 'relative_position', 'sparsity_mask', and 'use_additional_bucket'
         """
         current_max_seq_len = max([data.input_ids.shape[1] for data in data_instances])
@@ -451,6 +452,9 @@ class DataProcessor():
             sparsity_mask = torch.zeros((len(data_instances), max_seq_len, max_seq_len), dtype=torch.bool, device=device)
             use_additional_bucket = torch.zeros((len(data_instances), max_seq_len, max_seq_len), dtype=torch.bool, device=device)
 
+        if return_attention_mask:
+            attention_mask = torch.zeros((len(data_instances), max_seq_len), dtype=torch.bool, device=device)
+
         # fill tensors
         for i, data in enumerate(data_instances):
             instance_len = min(data.input_ids.shape[1], max_seq_len)
@@ -459,6 +463,8 @@ class DataProcessor():
                 relative_position[i, :instance_len, :instance_len] = data.relative_position[:, :instance_len, :instance_len]
                 sparsity_mask[i, :instance_len, :instance_len] = data.sparsity_mask[:, :instance_len, :instance_len]
                 use_additional_bucket[i, :instance_len, :instance_len] = data.use_additional_bucket[:, :instance_len, :instance_len]
+            if return_attention_mask:
+                attention_mask[i, :instance_len] = 1
 
         model_input = {
             'input_ids': input_ids,
@@ -467,6 +473,8 @@ class DataProcessor():
             'use_additional_bucket': use_additional_bucket,
             **kwargs
         }
+        if return_attention_mask:
+            return model_input, attention_mask
         return model_input
 
     @staticmethod