Upload folder using huggingface_hub

custom_generate/generate.py

@@ -1,7 +1,7 @@
 from typing import Union
 import torch
 from transformers import LogitsProcessorList, StoppingCriteriaList, GenerationConfig
-from transformers.generation.utils import GenerateNonBeamOutput, GenerateDecoderOnlyOutput
+from transformers.generation.utils import GenerationMixin, GenerateNonBeamOutput, GenerateDecoderOnlyOutput
 import torch.nn as nn
 import torch.nn.functional as F
 import numpy as np
@@ -332,5 +332,7 @@ def generate(model, *args, **kwargs):
             `'high'` to improve short-answer tasks. Check the [documentation](https://huggingface.co/transformers-community/dola)
             or [the paper](https://huggingface.co/papers/2309.03883) for more details.
     """
-    generation_outputs = model.generate(*args, custom_generate=_dola_decoding, **kwargs)
+    generation_outputs = GenerationMixin.generate(
+        model, *args, custom_generate=_dola_decoding, **kwargs
+    )
     return generation_outputs