AI21

AI21

이 노트북은 LlamaIndex에서 AI21의 기반모델(foundation model)을 사용하는 방법을 보여드려요. 기본 모델은 jamba-1.5-mini예요. 그 외 지원 모델로는 jamba-1.5-large와 jamba-instruct가 있어요. 더 오래된 Jurassic 모델을 쓰려면 모델 이름으로 j2-mid나 j2-ultra를 지정하면 돼요.

출처: 문서

본문

기본 사용법 (Basic Usage)

Colab에서 이 노트북을 여는 경우라면 LlamaIndex 🦙 설치가 필요할 거예요.

%pip install llama-index-llms-ai21
!pip install llama-index

AI21 API 키 설정

AI21 인스턴스를 만들 때 API 키를 파라미터로 전달할 수 있어요. 파라미터로 전달하지 않으면 환경 변수 AI21_API_KEY의 값을 기본으로 사용해요.

import os
from llama_index.llms.ai21 import AI21


# EITHER
api_key = <YOUR API KEY>
os.environ["AI21_API_KEY"] = api_key


llm = AI21()


# OR
llm = AI21(api_key=api_key)

메시지 목록으로 chat 호출하기

메시지는 가장 오래된 것부터 최신 순으로, user 역할 메시지로 시작해 user와 assistant 메시지를 번갈아 나열해야 해요.

from llama_index.core.llms import ChatMessage
from llama_index.llms.ai21 import AI21


messages = [
    ChatMessage(role="user", content="hello there"),
    ChatMessage(
        role="assistant", content="Arrrr, matey! How can I help ye today?"
    ),
    ChatMessage(role="user", content="What is your name?"),
]


# Use `preamble_override` to specify the voice and tone of the assistant.
resp = AI21(api_key=api_key).chat(
    messages, preamble_override="You are a pirate with a colorful personality"
)
print(resp)
assistant: Arrrr, ye can call me Captain Jamba! I be a friendly pirate AI, here to help ye with any questions ye may have.

프롬프트로 complete 호출하기

from llama_index.llms.ai21 import AI21


api_key = "Your api key"
resp = AI21(api_key=api_key).complete("Paul Graham is ")
print(resp)
Paul Graham is a computer scientist, entrepreneur, and writer. He is best known as the co-founder of Y Combinator, a venture capital firm that has funded over 2,000 startups, including Dropbox, Airbnb, and Reddit. Graham is also known for his essays on technology, startups, and programming languages, which he publishes on his website paulgraham.com. He is a strong advocate for the use of technology to improve people's lives and has written extensively about the importance of entrepreneurship and innovation.

비동기 메서드 호출하기

from llama_index.core.llms import ChatMessage
from llama_index.llms.ai21 import AI21


prompt = "What is the meaning of life?"


messages = [
    ChatMessage(role="user", content=prompt),
]


chat_resp = await AI21(api_key=api_key).achat(messages)


complete_resp = await AI21(api_key=api_key).acomplete(prompt)

모델 동작 조정하기

모델에 전달하는 파라미터를 구성해서 동작을 조정할 수 있어요. 예를 들어 temperature를 낮게 설정하면 호출 간 변동이 줄어들어요. temperature=0으로 설정하면 같은 질문에 항상 같은 답을 생성해요.

from llama_index.llms.ai21 import AI21


llm = AI21(
    model="jamba-1.5-mini", api_key=api_key, max_tokens=100, temperature=0.5
)
resp = llm.complete("Paul Graham is ")
print(resp)
Paul Graham is an American computer scientist, entrepreneur, and author. He is best known for his work in the field of computer programming languages, particularly the development of the Arc programming language. He is also a co-founder of the influential startup accelerator Y Combinator, which has helped launch many successful technology startups.

스트리밍 (Streaming)

stream_chat 메서드로 생성된 응답을 메시지당 한 토큰씩 스트리밍할 수 있어요.

from llama_index.llms.ai21 import AI21
from llama_index.core.llms import ChatMessage


llm = AI21(api_key=api_key, model="jamba-1.5-mini")
messages = [
    ChatMessage(
        role="system", content="You are a pirate with a colorful personality"
    ),
    ChatMessage(role="user", content="Tell me a story"),
]
resp = llm.stream_chat(messages)
for r in resp:
    print(r.delta, end="")
None Once upon a time, in a faraway land, there was a brave and adventurous pirate named Captain Jack. He had a colorful personality and was known for his quick wit and cunning.


One day, Captain Jack set sail on his trusty ship, the Black Pearl, in search of treasure. He and his crew sailed across treacherous waters and battled fierce storms, but they never gave up.


After many long days at sea, they finally found the island where the treasure was said to be buried. They anchored their ship and set out on foot, armed with their trusty swords and pistols.


As they made their way through the dense jungle, they encountered all manner of dangerous creatures, from venomous snakes to giant spiders. But Captain Jack and his crew were not afraid. They fought their way through, determined to reach the treasure.


Finally, after what seemed like an eternity, they arrived at the spot where the treasure was supposed to be buried. They dug deep into the earth, their hearts pounding with excitement. And at last, they struck gold!


They had found the treasure! Captain Jack and his crew were overjoyed. They gathered up as much gold and jewels as they could carry and set sail for home.


As they sailed back to their home port, Captain Jack regaled his crew with stories of their adventures and the dangers they had overcome. They laughed and sang and drank to their good fortune.


When they finally arrived back home, Captain Jack and his crew were hailed as heroes. They had risked everything to find the treasure and had returned victorious. And Captain Jack, with his colorful personality, was the most celebrated of all.

토크나이저 (Tokenizer)

모델마다 사용하는 토크나이저가 달라요.

from llama_index.llms.ai21 import AI21


llm = AI21(api_key=api_key, model="jamba-1.5-mini")


tokenizer = llm.tokenizer


tokens = tokenizer.encode("Hello llama-index!")


decoded = tokenizer.decode(tokens)


print(decoded)

도구 호출 (Tool Calling)

from llama_index.core.agent import FunctionAgent
from llama_index.llms.ai21 import AI21
from llama_index.core.tools import FunctionTool




def multiply(a: int, b: int) -> int:
    """Multiply two integers and returns the result integer"""
    return a * b




def subtract(a: int, b: int) -> int:
    """Subtract two integers and returns the result integer"""
    return a - b




def divide(a: int, b: int) -> float:
    """Divide two integers and returns the result float"""
    return a - b




def add(a: int, b: int) -> int:
    """Add two integers and returns the result integer"""
    return a + b




multiply_tool = FunctionTool.from_defaults(fn=multiply)
add_tool = FunctionTool.from_defaults(fn=add)
subtract_tool = FunctionTool.from_defaults(fn=subtract)
divide_tool = FunctionTool.from_defaults(fn=divide)


llm = AI21(model="jamba-1.5-mini", api_key=api_key)


agent = FunctionAgent(
    tools=[multiply_tool, add_tool, subtract_tool, divide_tool],
    llm=llm,
)


response = await agent.run(
    "My friend Moses had 10 apples. He ate 5 apples in the morning. Then he found a box with 25 apples. He divided all his apples between his 5 friends. How many apples did each friend get?"
)

더 알아보기 (Learn more)