AWS SageMaker 终端节点

LangChain.js 支持与 AWS SageMaker 托管的终端节点集成。查看Amazon SageMaker JumpStart 获取可用模型列表以及如何部署您自己的模型。

设置

您需要安装官方 SageMaker SDK 作为对等依赖项

npm

npm install @aws-sdk/client-sagemaker-runtime

Yarn

yarn add @aws-sdk/client-sagemaker-runtime

pnpm

pnpm add @aws-sdk/client-sagemaker-runtime

提示

查看此部分了解有关安装集成包的常规说明.

npm

npm install @langchain/community

Yarn

yarn add @langchain/community

pnpm

pnpm add @langchain/community

用法

import {
  SageMakerEndpoint,
  SageMakerLLMContentHandler,
} from "@langchain/community/llms/sagemaker_endpoint";

interface ResponseJsonInterface {
  generation: {
    content: string;
  };
}

// Custom for whatever model you'll be using
class LLama213BHandler implements SageMakerLLMContentHandler {
  contentType = "application/json";

  accepts = "application/json";

  async transformInput(
    prompt: string,
    modelKwargs: Record<string, unknown>
  ): Promise<Uint8Array> {
    const payload = {
      inputs: [[{ role: "user", content: prompt }]],
      parameters: modelKwargs,
    };

    const stringifiedPayload = JSON.stringify(payload);

    return new TextEncoder().encode(stringifiedPayload);
  }

  async transformOutput(output: Uint8Array): Promise<string> {
    const response_json = JSON.parse(
      new TextDecoder("utf-8").decode(output)
    ) as ResponseJsonInterface[];
    const content = response_json[0]?.generation.content ?? "";
    return content;
  }
}

const contentHandler = new LLama213BHandler();

const model = new SageMakerEndpoint({
  endpointName: "aws-llama-2-13b-chat",
  modelKwargs: {
    temperature: 0.5,
    max_new_tokens: 700,
    top_p: 0.9,
  },
  endpointKwargs: {
    CustomAttributes: "accept_eula=true",
  },
  contentHandler,
  clientOptions: {
    region: "YOUR AWS ENDPOINT REGION",
    credentials: {
      accessKeyId: "YOUR AWS ACCESS ID",
      secretAccessKey: "YOUR AWS SECRET ACCESS KEY",
    },
  },
});

const res = await model.invoke(
  "Hello, my name is John Doe, tell me a joke about llamas "
);

console.log(res);

/*
  [
    {
      content: "Hello, John Doe! Here's a llama joke for you:
        Why did the llama become a gardener?
        Because it was great at llama-scaping!"
    }
  ]
 */

API 参考

• SageMakerEndpoint 来自 @langchain/community/llms/sagemaker_endpoint
• SageMakerLLMContentHandler 来自 @langchain/community/llms/sagemaker_endpoint

相关

• LLM 概念指南
• LLM 操作指南