Skip to main content
Amazon SageMaker 是一项全托管服务,用于快速轻松地构建、训练和部署机器学习(ML)模型。
Amazon SageMaker ExperimentsAmazon SageMaker 的一项功能,允许你组织、跟踪、比较和评估 ML 实验及模型版本。
本笔记本展示了如何使用 LangChain 回调将提示词和其他 LLM 超参数记录并跟踪到 SageMaker Experiments 中。我们通过不同场景来演示该功能:
  • 场景 1单一 LLM - 使用单个 LLM 模型根据给定提示词生成输出的情况。
  • 场景 2顺序链 - 使用两个 LLM 模型组成顺序链的情况。
  • 场景 3带工具的智能体(思维链) - 在 LLM 基础上使用多种工具(搜索和数学)的情况。
在本笔记本中,我们将创建一个实验来记录每个场景中的提示词。

安装与设置

pip install -qU  sagemaker
pip install -qU  langchain-openai
pip install -qU  google-search-results
首先,设置所需的 API 密钥
import os

## Add your API keys below
os.environ["OPENAI_API_KEY"] = "<ADD-KEY-HERE>"
os.environ["SERPAPI_API_KEY"] = "<ADD-KEY-HERE>"
from langchain_community.callbacks.sagemaker_callback import SageMakerCallbackHandler
from langchain.agents import create_agent, load_tools
from langchain_classic.chains import LLMChain, SimpleSequentialChain
from langchain_core.prompts import PromptTemplate
from langchain_openai import OpenAI
from sagemaker.analytics import ExperimentAnalytics
from sagemaker.experiments.run import Run
from sagemaker.session import Session

LLM 提示词跟踪

# LLM Hyperparameters
HPARAMS = {
    "temperature": 0.1,
    "model_name": "gpt-3.5-turbo-instruct",
}

# Bucket used to save prompt logs (Use `None` is used to save the default bucket or otherwise change it)
BUCKET_NAME = None

# Experiment name
EXPERIMENT_NAME = "langchain-sagemaker-tracker"

# Create SageMaker Session with the given bucket
session = Session(default_bucket=BUCKET_NAME)

场景 1 - LLM

RUN_NAME = "run-scenario-1"
PROMPT_TEMPLATE = "tell me a joke about {topic}"
INPUT_VARIABLES = {"topic": "fish"}
with Run(
    experiment_name=EXPERIMENT_NAME, run_name=RUN_NAME, sagemaker_session=session
) as run:
    # Create SageMaker Callback
    sagemaker_callback = SageMakerCallbackHandler(run)

    # Define LLM model with callback
    llm = OpenAI(callbacks=[sagemaker_callback], **HPARAMS)

    # Create prompt template
    prompt = PromptTemplate.from_template(template=PROMPT_TEMPLATE)

    # Create LLM Chain
    chain = LLMChain(llm=llm, prompt=prompt, callbacks=[sagemaker_callback])

    # Run chain
    chain.run(**INPUT_VARIABLES)

    # Reset the callback
    sagemaker_callback.flush_tracker()

场景 2 - 顺序链

RUN_NAME = "run-scenario-2"

PROMPT_TEMPLATE_1 = """You are a playwright. Given the title of play, it is your job to write a synopsis for that title.
Title: {title}
Playwright: This is a synopsis for the above play:"""
PROMPT_TEMPLATE_2 = """You are a play critic from the New York Times. Given the synopsis of play, it is your job to write a review for that play.
Play Synopsis: {synopsis}
Review from a New York Times play critic of the above play:"""

INPUT_VARIABLES = {
    "input": "documentary about good video games that push the boundary of game design"
}
with Run(
    experiment_name=EXPERIMENT_NAME, run_name=RUN_NAME, sagemaker_session=session
) as run:
    # Create SageMaker Callback
    sagemaker_callback = SageMakerCallbackHandler(run)

    # Create prompt templates for the chain
    prompt_template1 = PromptTemplate.from_template(template=PROMPT_TEMPLATE_1)
    prompt_template2 = PromptTemplate.from_template(template=PROMPT_TEMPLATE_2)

    # Define LLM model with callback
    llm = OpenAI(callbacks=[sagemaker_callback], **HPARAMS)

    # Create chain1
    chain1 = LLMChain(llm=llm, prompt=prompt_template1, callbacks=[sagemaker_callback])

    # Create chain2
    chain2 = LLMChain(llm=llm, prompt=prompt_template2, callbacks=[sagemaker_callback])

    # Create Sequential chain
    overall_chain = SimpleSequentialChain(
        chains=[chain1, chain2], callbacks=[sagemaker_callback]
    )

    # Run overall sequential chain
    overall_chain.run(**INPUT_VARIABLES)

    # Reset the callback
    sagemaker_callback.flush_tracker()

场景 3 - 带工具的智能体

RUN_NAME = "run-scenario-3"
PROMPT_TEMPLATE = "Who is the oldest person alive? And what is their current age raised to the power of 1.51?"

with Run(
    experiment_name=EXPERIMENT_NAME,
    run_name=RUN_NAME,
    sagemaker_session=session,
) as run:
    # Create SageMaker Callback
    sagemaker_callback = SageMakerCallbackHandler(run)

    # Define LLM model with callback
    llm = OpenAI(callbacks=[sagemaker_callback], **HPARAMS)

    # Define tools
    tools = load_tools(
        ["serpapi", "llm-math"],
        llm=llm,
        callbacks=[sagemaker_callback],
    )

    # Create agent
    agent = create_agent(
        model=llm,
        tools=tools,
        callbacks=[sagemaker_callback],
    )

    # Run agent
    agent.invoke(PROMPT_TEMPLATE)

    # Reset callback
    sagemaker_callback.flush_tracker()

加载日志数据

提示词记录完成后,我们可以按如下方式轻松加载并将其转换为 Pandas DataFrame。
# Load
logs = ExperimentAnalytics(experiment_name=EXPERIMENT_NAME)

# Convert as pandas dataframe
df = logs.dataframe(force_refresh=True)

print(df.shape)
df.head()
如上所示,实验中有三次运行(行),分别对应每个场景。每次运行都将提示词及相关 LLM 设置/超参数以 JSON 格式记录并保存在 S3 存储桶中。你可以自由地从每个 JSON 路径加载和探索日志数据。