allow Output subclasses in contrastive search

custom_generate/generate.py

@@ -1,3 +1,4 @@
+from dataclasses import replace
 import logging
 from typing import TYPE_CHECKING, Optional, Union
 
@@ -14,7 +15,6 @@ from transformers.generation.utils import (
     GenerateNonBeamOutput,
     GenerationMixin,
 )
-from transformers.modeling_outputs import CausalLMOutputWithPast, Seq2SeqLMOutput
 from transformers.utils import ModelOutput
 
 
@@ -414,7 +414,8 @@ def _contrastive_search(
                 for layer in outputs.decoder_attentions:
                     layer = torch.stack(torch.split(layer, top_k, dim=0))[range(batch_size), selected_idx, ...]
                     next_step_decoder_attentions += (layer,)
-            outputs = Seq2SeqLMOutput(
+            outputs = replace(
+                outputs,
                 past_key_values=next_past_key_values,
                 decoder_hidden_states=next_decoder_hidden_states,
                 decoder_attentions=next_step_decoder_attentions or None,
@@ -426,11 +427,13 @@ def _contrastive_search(
                 for layer in outputs.attentions:
                     layer = torch.stack(torch.split(layer, top_k, dim=0))[range(batch_size), selected_idx, ...]
                     next_step_attentions += (layer,)
-            outputs = CausalLMOutputWithPast(
+            outputs = replace(
+                outputs,
                 past_key_values=next_past_key_values,
                 hidden_states=next_decoder_hidden_states,
                 attentions=next_step_attentions or None,
             )
+
         # contrastive_search main logic end
 
         # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping