Generate: faster `can_generate` check on TF and Flax (#23398)

130e1542 · Joao Gante · GitHub · 2922e394 · 130e1542 · 130e1542
Unverified Commit 130e1542 authored May 16, 2023 by Joao Gante Committed by GitHub May 16, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

src/transformers/modeling_flax_utils.py src/transformers/modeling_flax_utils.py +1 -1

src/transformers/modeling_tf_utils.py src/transformers/modeling_tf_utils.py +1 -1

No files found.
--- a/src/transformers/modeling_flax_utils.py
+++ b/src/transformers/modeling_flax_utils.py
@@ -474,7 +474,7 @@ class FlaxPreTrainedModel(PushToHubMixin, FlaxGenerationMixin):
            `bool`: Whether this model can generate sequences with `.generate()`.
        """
        # Detects whether `prepare_inputs_for_generation` has been overwritten, which is a requirement for generation
-        if "GenerationMixin" in str(self.prepare_inputs_for_generation):
+        if "GenerationMixin" in str(self.prepare_inputs_for_generation.__func__):
            return False
        return True


--- a/src/transformers/modeling_tf_utils.py
+++ b/src/transformers/modeling_tf_utils.py
@@ -1243,7 +1243,7 @@ class TFPreTrainedModel(tf.keras.Model, TFModelUtilsMixin, TFGenerationMixin, Pu
            `bool`: Whether this model can generate sequences with `.generate()`.
        """
        # Detects whether `prepare_inputs_for_generation` has been overwritten, which is a requirement for generation
-        if "GenerationMixin" in str(self.prepare_inputs_for_generation):
+        if "GenerationMixin" in str(self.prepare_inputs_for_generation.__func__):
            return False
        return True