SagemakerGenerator — Amazon SageMaker 텍스트 생성

이 컴포넌트는 Amazon SageMaker에 배포된 대규모 언어 모델(LLM)을 이용해 텍스트 생성을 할 수 있게 해 줘요. SageMaker에 올려둔 나만의 모델 엔드포인트를 Haystack 파이프라인에서 호출하는 한 조각이라고 보시면 돼요.

출처: 공식 문서 — SagemakerGenerator

SagemakerGenerator는 AWS SageMaker에 배포된 모델을 사용할 수 있게 해 줘요. SageMaker에서 배포해 둔 모델을 Haystack에서 바로 호출하는 거죠.

파라미터 개요 (Parameters Overview)

SagemakerGenerator는 동작하려면 AWS 자격 증명이 필요해요. AWS_ACCESS_KEY_ID와 AWS_SECRET_ACCESS_KEY 환경 변수를 설정하세요.

컴포넌트가 동작하려면 초기화할 때 SageMaker 엔드포인트도 지정해야 해요. 엔드포인트 이름을 model 파라미터로 넘기면 됩니다.

generator = SagemakerGenerator(model="jumpstart-dft-hf-llm-falcon-7b-instruct-bf16")

추가로, 특정 모델에 유효한 텍스트 생성 파라미터는 generation_kwargs 파라미터를 통해 SagemakerGenerator로 직접 전달할 수 있어요. 초기화할 때도, run() 메서드를 호출할 때도 둘 다 가능하죠.

모델이 커스텀 속성(custom attributes)도 요구한다면, 초기화할 때 aws_custom_attributes 파라미터에 딕셔너리로 넘기면 돼요.

이런 커스텀 파라미터가 필요한 대표적인 모델 계열이 Llama2인데, {"accept_eula": True}로 초기화해야 해요.

generator = SagemakerGenerator(
    model="jumpstart-dft-meta-textgenerationneuron-llama-2-7b",
    aws_custom_attributes={"accept_eula": True},
)

사용법 (Usage)

SagemakerGenerator를 쓰려면 amazon-sagemaker-haystack 패키지를 설치해야 해요.

pip install amazon-sagemaker-haystack

단독 사용 (On its own)

기본 사용법:

from haystack_integrations.components.generators.amazon_sagemaker import (
    SagemakerGenerator,
)

client = SagemakerGenerator(model="jumpstart-dft-hf-llm-falcon-7b-instruct-bf16")
response = client.run("Briefly explain what NLP is in one sentence.")
print(response)
# >> {'replies': ["Natural Language Processing (NLP) is a subfield of artificial intelligence and computational linguistics that focuses on the interaction between computers and human languages..."],
# >>  'metadata': [{}]}

파이프라인에서 사용 (In a pipeline)

RAG 파이프라인에서 사용:

from haystack_integrations.components.generators.amazon_sagemaker import (
    SagemakerGenerator,
)
from haystack import Pipeline
from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
from haystack.components.builders import PromptBuilder

template = """
Given the following information, answer the question.

Context:
{% for document in documents %}
    {{ document.content }}
{% endfor %}

Question: What's the official language of {{ country }}?
"""
pipe = Pipeline()

pipe.add_component("retriever", InMemoryBM25Retriever(document_store=docstore))
pipe.add_component("prompt_builder", PromptBuilder(template=template))
pipe.add_component(
    "llm",
    SagemakerGenerator(model="jumpstart-dft-hf-llm-falcon-7b-instruct-bf16"),
)
pipe.connect("retriever", "prompt_builder.documents")
pipe.connect("prompt_builder", "llm")

pipe.run({"prompt_builder": {"country": "France"}})