在Colab中调用HuggingFace模型时出现StopIteration错误

问题场景 #

在Colab中已将HuggingFace Token存入Colab Secrets，成功获取token后，调用ChatHuggingFace.invoke()时触发StopIteration错误。

1. Token获取代码及输出 #

from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
from dotenv import load_dotenv
from google.colab import userdata
import os

print("Before try block:", 'hf_token' in locals())
# Initialize hf_token with a default value (None)
hf_token = None
try:
    hf_token = userdata.get('HF_TOKEN')
    print("Inside try block, after userdata.get:", 'hf_token' in locals())
    if hf_token:
      print("Hugging Face token successfully retrieved from Colab Secrets.")
    else:
      print("Warning: Hugging Face token not found in Colab Secrets. Ensure you've added it.")
except Exception as e:
    print(f"An error occurred while accessing Colab Secrets: {e}")

print("After try block:", 'hf_token' in locals())

输出：

Before try block: False

Inside try block, after userdata.get: True

Hugging Face token successfully retrieved from Colab Secrets.

After try block: True

2. 模型调用代码及错误输出 #

if hf_token:
    print("Inside the 'if hf_token' block, hf_token:", hf_token)
    llm = HuggingFaceEndpoint(
        repo_id="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
        task="text-generation",
        huggingfacehub_api_token=hf_token,
        max_new_tokens=150  # Increased max_new_tokens
    )
    # Initialize ChatHuggingFace directly with the HuggingFaceEndpoint instance
    model = ChatHuggingFace(llm=llm)
    try:
        result = model.invoke('Who is Lincoln')
        print(result)
    except Exception as e:
        print(f"Error during model invocation: Type: {type(e)}, Message: {e}")
else:
    print("Hugging Face model could not be initialized due to missing token.")

输出：

Inside the 'if hf_token' block, hf_token: [已隐藏的token]

Error during model invocation: Type: <class 'StopIteration'>, Message: 

解决方法 #

方法1：适配模型的聊天格式 #

TinyLlama-1.1B-Chat-v1.0需要特定对话格式，直接传入普通文本会导致输出异常。通过model_kwargs指定聊天模板，并使用标准消息格式调用：

if hf_token:
    print("Inside the 'if hf_token' block, hf_token:", hf_token)
    llm = HuggingFaceEndpoint(
        repo_id="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
        task="text-generation",
        huggingfacehub_api_token=hf_token,
        max_new_tokens=150,
        model_kwargs={
            "temperature": 0.7,
            "top_p": 0.9,
            "chat_template": "{% for message in messages %}{% if message['role'] == 'user' %}### User: {{ message['content'] }} {% elif message['role'] == 'assistant' %}### Assistant: {{ message['content'] }} {% endif %}{% endfor %}### Assistant:"
        }
    )
    # 使用标准HumanMessage格式传入问题
    from langchain_core.messages import HumanMessage
    result = model.invoke([HumanMessage(content="Who is Lincoln")])
    print(result.content)

方法2：本地加载模型替代API调用 #

改用transformers本地加载模型，避免HuggingFace Endpoint API的迭代器异常问题：

!pip install transformers accelerate

from langchain_huggingface import ChatHuggingFace, HuggingFacePipeline
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from google.colab import userdata

hf_token = userdata.get('HF_TOKEN')

# 加载tokenizer和模型
tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0", token=hf_token)
model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0", token=hf_token, device_map="auto")

# 创建文本生成pipeline
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=150,
    temperature=0.7,
    top_p=0.9
)

# 初始化ChatHuggingFace并调用
chat_model = ChatHuggingFace(pipeline=pipe)
from langchain_core.messages import HumanMessage
result = chat_model.invoke([HumanMessage(content="Who is Lincoln")])
print(result.content)

方法3：捕获异常并重试 #

如果暂时无法调整模型配置，可直接捕获StopIteration异常并尝试重试：

try:
    result = model.invoke('Who is Lincoln')
    print(result)
except StopIteration:
    print("模型调用出现迭代器终止错误，正在重试...")
    result = model.invoke('Who is Lincoln')
    print(result)

内容的提问来源于stack exchange，提问作者Sam2021