Upload model

config.json

@@ -21,6 +21,10 @@
     "do_sample": false,
     "early_stopping": false,
     "encoder_no_repeat_ngram_size": 0,
+    "end_of_section_token_ids": [
+      3,
+      2
+    ],
     "eos_token_id": 2,
     "exponential_decay_length_penalty": null,
     "finetuning_task": null,

modelling_cxrrg.py

@@ -1,13 +1,11 @@
-import functools
-import os
 from typing import Optional, Tuple, Union
 
 import torch
 import transformers
-from torch.nn import CrossEntropyLoss, Linear
+from torch.nn import CrossEntropyLoss
 from transformers import PreTrainedTokenizerFast, VisionEncoderDecoderModel
 from transformers.configuration_utils import PretrainedConfig
-from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput
+from transformers.modeling_outputs import Seq2SeqLMOutput
 from transformers.modeling_utils import PreTrainedModel
 from transformers.models.vision_encoder_decoder.configuration_vision_encoder_decoder import (
     VisionEncoderDecoderConfig,
@@ -457,13 +455,13 @@ class CXRRGModel(VisionEncoderDecoderModel):
         _, seq_len = token_ids.shape
 
         # The number of sections is the same as the number of separator_token_ids:
-        num_sections = len(self.config.decoder.separator_token_ids)
+        num_sections = len(self.config.decoder.end_of_section_token_ids)
 
         sections = {k: [] for k in range(num_sections)}
 
         for i in token_ids:
             prev_col = 0
-            for j, k in enumerate(self.config.decoder.separator_token_ids):
+            for j, k in enumerate(self.config.decoder.end_of_section_token_ids):
 
                 # The maximum sequence length was exceeded, thus no more tokens:
                 if prev_col >= seq_len: