From ee56ddb219699f27a51e19abca482f062f9b99d2 Mon Sep 17 00:00:00 2001 From: Arne Binder Date: Wed, 6 Nov 2024 15:54:10 +0100 Subject: [PATCH] add parameter use_key_value_cache to BartAsPointerNetworkConfig --- .../models/base_models/bart_as_pointer_network.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/src/pie_modules/models/base_models/bart_as_pointer_network.py b/src/pie_modules/models/base_models/bart_as_pointer_network.py index ffb1ce9bf..9d9a3cc60 100644 --- a/src/pie_modules/models/base_models/bart_as_pointer_network.py +++ b/src/pie_modules/models/base_models/bart_as_pointer_network.py @@ -7,7 +7,7 @@ from transformers import BartConfig, BartModel, BartPreTrainedModel from transformers.modeling_outputs import Seq2SeqLMOutput, Seq2SeqModelOutput from transformers.models.bart.modeling_bart import shift_tokens_right -from transformers.utils import logging +from transformers.utils import ModelOutput, logging from pie_modules.models.base_models.bart_with_decoder_position_ids import ( BartModelWithDecoderPositionIds, @@ -58,6 +58,7 @@ def __init__( shared_decay: Optional[float] = None, encoder_layer_norm_decay: Optional[float] = 0.001, decoder_layer_norm_decay: Optional[float] = None, + use_key_value_cache: bool = True, # other BartConfig parameters **kwargs, ): @@ -82,6 +83,8 @@ def __init__( self.encoder_layer_norm_decay = encoder_layer_norm_decay self.decoder_layer_norm_decay = decoder_layer_norm_decay + self.use_key_value_cache = use_key_value_cache + class BartAsPointerNetwork(BartPreTrainedModel): config_class = BartAsPointerNetworkConfig @@ -311,6 +314,14 @@ def forward( encoder_attentions=outputs.encoder_attentions, ) + def _extract_past_from_model_output( + self, outputs: ModelOutput, standardize_cache_format: bool = False + ): + if not self.config.use_key_value_cache: + return None + else: + return super()._extract_past_from_model_output(outputs, standardize_cache_format) + def prepare_inputs_for_generation( self, decoder_input_ids,