Upload model

modelling_cxrrg.py

@@ -201,7 +201,6 @@ class CXRRGModel(VisionEncoderDecoderModel):
         self,
         input_ids,
         special_token_ids,
-        token_type_id_sections=None,
         past_key_values=None,
         use_cache=None,
         encoder_outputs=None,
@@ -227,7 +226,7 @@ class CXRRGModel(VisionEncoderDecoderModel):
             # `inputs_embeds` are only to be used in the 1st generation step:
             inputs_embeds = torch.cat([encoder_outputs[0], self.decoder.get_input_embeddings()(input_ids)], dim=1)
 
-            decoder_token_type_ids = self.token_ids_to_token_type_ids(input_ids, special_token_ids, token_type_id_sections)
+            decoder_token_type_ids = self.token_ids_to_token_type_ids(input_ids, special_token_ids)
             decoder_token_type_ids = torch.cat(
                 [
                     torch.full(
@@ -256,7 +255,7 @@ class CXRRGModel(VisionEncoderDecoderModel):
             decoder_position_ids.masked_fill_(report_attention_mask == 0, 1)
             
             # Always place token_ids_to_token_type_ids_past before input_ids = input_ids[:, remove_prefix_length:]:
-            decoder_token_type_ids = self.token_ids_to_token_type_ids_past(input_ids, special_token_ids, token_type_id_sections)
+            decoder_token_type_ids = self.token_ids_to_token_type_ids_past(input_ids, special_token_ids)
             decoder_position_ids = decoder_position_ids[:, -1:]
 
             past_length = past_key_values[0][0].shape[2]
@@ -283,7 +282,7 @@ class CXRRGModel(VisionEncoderDecoderModel):
         )
         return input_dict
         
-    def token_ids_to_token_type_ids(self, token_ids, special_token_ids, token_type_id_sections=None):
+    def token_ids_to_token_type_ids(self, token_ids, special_token_ids):
         """
         Extract token type identifiers from the token identifiers.
 
@@ -296,10 +295,8 @@ class CXRRGModel(VisionEncoderDecoderModel):
             token_type_ids - token type identifiers.
         """
 
-        token_type_id_sections = token_type_id_sections if token_type_id_sections is not None else list(range(len(special_token_ids) + 1))
-
         mbatch_size, seq_len = token_ids.shape
-        token_type_ids = torch.full_like(token_ids, token_type_id_sections[0], dtype=torch.long, device=token_ids.device)
+        token_type_ids = torch.full_like(token_ids, self.config.section_ids[0], dtype=torch.long, device=token_ids.device)
 
         for i, j in enumerate(special_token_ids):
             # Find first occurrence of special tokens that indicate the boundary between sections:
@@ -322,11 +319,11 @@ class CXRRGModel(VisionEncoderDecoderModel):
                     )
                 ])
 
-                token_type_ids[ids[:, 0], ids[:, 1]] = token_type_id_sections[i + 1]
+                token_type_ids[ids[:, 0], ids[:, 1]] = self.config.section_ids[i + 1]
 
         return token_type_ids
 
-    def token_ids_to_token_type_ids_past(self, token_ids, special_token_ids, token_type_id_sections=None):
+    def token_ids_to_token_type_ids_past(self, token_ids, special_token_ids):
         """
         Extract token type identifiers from the token identifiers if past != None. Make sure to input all the
         token_ids (e.g., do not input input_ids = input_ids[:, remove_prefix_length:] from prepare_inputs_for_generation).
@@ -339,8 +336,7 @@ class CXRRGModel(VisionEncoderDecoderModel):
             token_type_ids - token type identifiers.
         """
 
-        token_type_id_sections = token_type_id_sections if token_type_id_sections is not None else list(range(len(special_token_ids) + 1))
-        token_type_ids = torch.full([token_ids.shape[0], 1], token_type_id_sections[0], dtype=torch.long, device=token_ids.device)
+        token_type_ids = torch.full([token_ids.shape[0], 1], self.config.section_ids[0], dtype=torch.long, device=token_ids.device)
 
         # https://huggingface.co/docs/transformers/model_doc/bert#transformers.BertTokenizer.create_token_type_ids_from_sequences.example
         token_ids = token_ids[:, :-1]
@@ -349,7 +345,7 @@ class CXRRGModel(VisionEncoderDecoderModel):
 
             # Find first occurrence of special token, which indicates the boundary between sections:
             exists = torch.any(token_ids == j, dim=1, keepdim=True)
-            token_type_ids[exists] = token_type_id_sections[i + 1]
+            token_type_ids[exists] = self.config.section_ids[i + 1]
 
         return token_type_ids