From 97e8995cf183b3fb5d404bf018489e2e29be5268 Mon Sep 17 00:00:00 2001 From: Joao Gante Date: Sun, 14 Jul 2024 16:39:42 +0100 Subject: [PATCH] Whisper: move to tensor cpu before converting to np array at decode time (#31954) --- src/transformers/models/whisper/tokenization_whisper.py | 7 +++++-- .../models/whisper/tokenization_whisper_fast.py | 7 +++++-- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/src/transformers/models/whisper/tokenization_whisper.py b/src/transformers/models/whisper/tokenization_whisper.py index aea9dd289e9ed0..f19f218569a293 100644 --- a/src/transformers/models/whisper/tokenization_whisper.py +++ b/src/transformers/models/whisper/tokenization_whisper.py @@ -872,8 +872,11 @@ def _strip_prompt(self, token_ids: List[int], prompt_token_id: int, decoder_star @staticmethod def _convert_to_list(token_ids): # convert type to ndarray if necessary - if "torch" in str(type(token_ids)) or "tensorflow" in str(type(token_ids)) and hasattr(token_ids, "numpy"): - token_ids = token_ids.numpy() + if hasattr(token_ids, "numpy"): + if "torch" in str(type(token_ids)): + token_ids = token_ids.cpu().numpy() + elif "tensorflow" in str(type(token_ids)): + token_ids = token_ids.numpy() # now the token ids are either a numpy array, or a list of lists if isinstance(token_ids, np.ndarray): token_ids = token_ids.tolist() diff --git a/src/transformers/models/whisper/tokenization_whisper_fast.py b/src/transformers/models/whisper/tokenization_whisper_fast.py index d1dee38263b14c..d1205d1a8ec01b 100644 --- a/src/transformers/models/whisper/tokenization_whisper_fast.py +++ b/src/transformers/models/whisper/tokenization_whisper_fast.py @@ -605,8 +605,11 @@ def _strip_prompt(self, token_ids: List[int], prompt_token_id: int, decoder_star # Copied from transformers.models.whisper.tokenization_whisper.WhisperTokenizer._convert_to_list def _convert_to_list(token_ids): # convert type to ndarray if necessary - if "torch" in str(type(token_ids)) or "tensorflow" in str(type(token_ids)) and hasattr(token_ids, "numpy"): - token_ids = token_ids.numpy() + if hasattr(token_ids, "numpy"): + if "torch" in str(type(token_ids)): + token_ids = token_ids.cpu().numpy() + elif "tensorflow" in str(type(token_ids)): + token_ids = token_ids.numpy() # now the token ids are either a numpy array, or a list of lists if isinstance(token_ids, np.ndarray): token_ids = token_ids.tolist()