如何获取Llama 2模型流式输出？SageMaker部署模型解析求助

本地运行场景（Hugging Face Transformers） #

使用transformers库时，通过设置streaming=True让文本生成管道返回生成器，迭代生成器即可逐段获取输出：

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

model_name = "meta-llama/Llama-2-7b-chat-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# 初始化流式生成管道
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    streaming=True,
    max_new_tokens=200
)

prompt = "请介绍流式输出的核心原理"
# 迭代生成器，即时打印增量内容
for chunk in pipe(prompt):
    generated_text = chunk['generated_text']
    # 提取本次新增的文本片段
    delta = generated_text[len(prompt):]
    print(delta, end='', flush=True)
    prompt = generated_text  # 更新prompt，避免重复输出前缀

本地运行场景（Llama.cpp） #

通过Python绑定启用stream=True参数，直接迭代生成结果：

from llama_cpp import Llama

llm = Llama(model_path="./llama-2-7b-chat.gguf")
prompt = "请介绍流式输出的核心原理"

for output in llm.create_completion(
    prompt,
    stream=True,
    max_tokens=200
):
    print(output['choices'][0]['text'], end='', flush=True)

一次性返回而非流式返回的问题，通常和端点部署配置、客户端调用方式有关，以下是针对性解决方案：

1. 确保端点推理脚本支持流式输出 #

修改SageMaker推理脚本（model.py），让predict_fn返回生成器而非一次性结果：

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

def model_fn(model_dir):
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    model = AutoModelForCausalLM.from_pretrained(model_dir).to("cuda")
    return {"model": model, "tokenizer": tokenizer}

def predict_fn(input_data, model):
    prompt = input_data["prompt"]
    tokenizer = model["tokenizer"]
    llm = model["model"]
    
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    # 开启流式生成模式
    outputs = llm.generate(
        **inputs,
        max_new_tokens=200,
        do_sample=True,
        stream=True
    )
    
    prev_text = prompt
    for output in outputs:
        decoded_text = tokenizer.decode(output[0], skip_special_tokens=True)
        # 仅返回本次新增的文本
        delta = decoded_text[len(prev_text):]
        yield {"generated_text": delta}
        prev_text = decoded_text

2. 客户端使用流式调用接口 #

必须使用boto3的invoke_endpoint_with_response_stream方法，而非普通的invoke_endpoint：

import boto3
import json

client = boto3.client("sagemaker-runtime")
endpoint_name = "your-llama-endpoint-name"

payload = json.dumps({"prompt": "请介绍流式输出的核心原理"})
response = client.invoke_endpoint_with_response_stream(
    EndpointName=endpoint_name,
    ContentType="application/json",
    Body=payload
)

# 逐块读取流式响应
event_stream = response["Body"]
for event in event_stream:
    if "PayloadPart" in event:
        chunk_data = json.loads(event["PayloadPart"]["Bytes"].decode("utf-8"))
        print(chunk_data["generated_text"], end='', flush=True)

3. 排查一次性返回的常见原因 #

• 检查推理脚本是否错误地将生成器结果收集为列表后一次性返回；
• 确认客户端未使用普通同步调用接口，必须切换为流式调用；
• 查看SageMaker端点日志，验证服务器端是否分块发送响应。

内容的提问来源于stack exchange，提问作者Irits Anna Mathew