From 634054601d7f27116690d47111096b0a3a63fc4d Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Thu, 9 Apr 2026 10:50:17 -0700 Subject: [PATCH 1/4] Fix compute_hidden_states_hf.py: handle BatchEncoding in apply_chat_template In transformers 4.46+, apply_chat_template() with return_tensors="pt" returns a BatchEncoding object that no longer subclasses dict. The previous isinstance(tokenized, dict) guard evaluated to False and fell through to tokenized (the BatchEncoding), causing input_ids.shape[1] to call BatchEncoding.__getattr__("shape") and raise AttributeError. Fix by checking isinstance(tokenized, torch.Tensor) instead, which correctly handles both old transformers (plain tensor return) and new transformers (BatchEncoding return). Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Ye Yu --- .../collect_hidden_states/compute_hidden_states_hf.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py index 1060df45622..046f7c4d14e 100644 --- a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py +++ b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py @@ -209,7 +209,14 @@ async def submit_generates(): tokenized = tokenizer.apply_chat_template( conversations, return_tensors="pt", add_generation_template=False ) - input_ids = tokenized["input_ids"] if isinstance(tokenized, dict) else tokenized + # apply_chat_template return type varies by transformers version: + # - older versions return a plain torch.Tensor (input_ids directly) + # - newer versions (4.46+) return a BatchEncoding which no longer + # subclasses dict, so isinstance(tokenized, dict) is False + if isinstance(tokenized, torch.Tensor): + input_ids = tokenized + else: + input_ids = tokenized["input_ids"] num_input_tokens = input_ids.shape[1] if num_input_tokens <= 10 or num_input_tokens > args.max_seq_len: num_skipped_too_long += 1 From ea48682536d52896ee91c3f79957459801f519df Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Thu, 9 Apr 2026 10:53:07 -0700 Subject: [PATCH 2/4] Simplify apply_chat_template fix: use ["input_ids"] directly Since transformers 4.46+ is required, apply_chat_template always returns a BatchEncoding. Drop the torch.Tensor fallback and just index ["input_ids"] directly. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Ye Yu --- .../compute_hidden_states_hf.py | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py index 046f7c4d14e..c779dbef6c5 100644 --- a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py +++ b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py @@ -206,17 +206,9 @@ async def submit_generates(): continue # Tokenize and check length - tokenized = tokenizer.apply_chat_template( + input_ids = tokenizer.apply_chat_template( conversations, return_tensors="pt", add_generation_template=False - ) - # apply_chat_template return type varies by transformers version: - # - older versions return a plain torch.Tensor (input_ids directly) - # - newer versions (4.46+) return a BatchEncoding which no longer - # subclasses dict, so isinstance(tokenized, dict) is False - if isinstance(tokenized, torch.Tensor): - input_ids = tokenized - else: - input_ids = tokenized["input_ids"] + )["input_ids"] num_input_tokens = input_ids.shape[1] if num_input_tokens <= 10 or num_input_tokens > args.max_seq_len: num_skipped_too_long += 1 From f762336ae2a8e809e45af56f60bda8bd9176909a Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Thu, 9 Apr 2026 11:20:24 -0700 Subject: [PATCH 3/4] Fix apply_chat_template: add return_dict=True for cross-version compat Without return_dict=True, apply_chat_template returns a raw torch.Tensor on transformers <5.0 (default return_dict=False) and a BatchEncoding on transformers >=5.0 (default changed to True). Subscripting a Tensor with ["input_ids"] raises TypeError on <5.0. Passing return_dict=True explicitly forces BatchEncoding on all versions (verified locally on 4.57.1 and 5.0.0). Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Ye Yu --- .../collect_hidden_states/compute_hidden_states_hf.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py index c779dbef6c5..1f604e44c66 100644 --- a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py +++ b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py @@ -206,8 +206,12 @@ async def submit_generates(): continue # Tokenize and check length + # return_dict=True ensures BatchEncoding is returned on all transformers + # versions: in <5.0 the default is False (returns raw tensor), in 5.0+ + # the default changed to True (returns BatchEncoding). input_ids = tokenizer.apply_chat_template( - conversations, return_tensors="pt", add_generation_template=False + conversations, return_tensors="pt", return_dict=True, + add_generation_template=False )["input_ids"] num_input_tokens = input_ids.shape[1] if num_input_tokens <= 10 or num_input_tokens > args.max_seq_len: From 4be7f6f2fabdb87cfb65fc9109ac576673e2a129 Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Thu, 9 Apr 2026 11:29:57 -0700 Subject: [PATCH 4/4] Fix ruff format: collapse apply_chat_template kwargs to single line Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Ye Yu --- .../collect_hidden_states/compute_hidden_states_hf.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py index 1f604e44c66..713b2326fcc 100644 --- a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py +++ b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py @@ -210,8 +210,7 @@ async def submit_generates(): # versions: in <5.0 the default is False (returns raw tensor), in 5.0+ # the default changed to True (returns BatchEncoding). input_ids = tokenizer.apply_chat_template( - conversations, return_tensors="pt", return_dict=True, - add_generation_template=False + conversations, return_tensors="pt", return_dict=True, add_generation_template=False )["input_ids"] num_input_tokens = input_ids.shape[1] if num_input_tokens <= 10 or num_input_tokens > args.max_seq_len: