Fix FunPackAdvisorLLM tokenize: handle BatchEncoding from apply_chat_template
apply_chat_template(..., tokenize=True, return_tensors="pt") returns a BatchEncoding in newer transformers versions instead of a plain tensor. Accessing .shape[1] on a BatchEncoding raises AttributeError because it looks up "shape" in the data dict and finds nothing. Extract input_ids from the result with hasattr check before calling .to(device) and computing prompt_length. Applied to both the normal chat-template path and the skip_template encode() path. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Sonnet 4.6
parent
a04bb95df7
commit
27637d99df
+11
-6
@@ -11169,9 +11169,10 @@ class _FunPackAdvisorLLMWrapper:
|
||||
thinking = bool(kwargs.get("thinking", True))
|
||||
if skip_template:
|
||||
# Text already has role tokens baked in (e.g. from TextGenerateLTX2Prompt).
|
||||
input_ids = self._tokenizer.encode(
|
||||
result = self._tokenizer.encode(
|
||||
str(text), return_tensors="pt", add_special_tokens=False
|
||||
).to(self._device)
|
||||
)
|
||||
input_ids = (result.input_ids if hasattr(result, "input_ids") else result).to(self._device)
|
||||
else:
|
||||
if system_prompt:
|
||||
messages = [
|
||||
@@ -11181,14 +11182,18 @@ class _FunPackAdvisorLLMWrapper:
|
||||
else:
|
||||
messages = [{"role": "user", "content": str(text)}]
|
||||
try:
|
||||
input_ids = self._tokenizer.apply_chat_template(
|
||||
result = self._tokenizer.apply_chat_template(
|
||||
messages, tokenize=True, add_generation_prompt=True, return_tensors="pt",
|
||||
enable_thinking=thinking,
|
||||
).to(self._device)
|
||||
)
|
||||
except TypeError:
|
||||
input_ids = self._tokenizer.apply_chat_template(
|
||||
result = self._tokenizer.apply_chat_template(
|
||||
messages, tokenize=True, add_generation_prompt=True, return_tensors="pt"
|
||||
).to(self._device)
|
||||
)
|
||||
# apply_chat_template may return a BatchEncoding or a plain tensor depending
|
||||
# on the transformers version and tokenizer. Extract input_ids from either.
|
||||
input_ids = result.input_ids if hasattr(result, "input_ids") else result
|
||||
input_ids = input_ids.to(self._device)
|
||||
return {"input_ids": input_ids, "prompt_length": int(input_ids.shape[1]), "thinking": thinking}
|
||||
|
||||
def generate(self, tokens, do_sample=True, max_length=800, temperature=0.7,
|
||||
|
||||
Reference in New Issue
Block a user