AI21¶

本 Notebook 展示了如何在 LlamaIndex 中使用 AI21 的基础模型。默认模型是 jamba-1.5-mini。其他支持的模型包括 jamba-1.5-large 和 jamba-instruct。如果您想使用较旧的 Jurassic 模型，请指定模型名称 j2-mid 或 j2-ultra。

基本用法¶

如果您在 Colab 上打开此 Notebook，您可能需要安装 LlamaIndex 🦙。

In [ ]

%pip install llama-index-llms-ai21

%pip install llama-index-llms-ai21

In [ ]

!pip install llama-index

!pip install llama-index

设置 AI21 API 密钥¶

创建 AI21 实例时，可以将 API 密钥作为参数传递。如果未作为参数提供，则默认为环境变量 AI21_API_KEY 的值。

In [ ]

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)

import os from llama_index.llms.ai21 import AI21 # EITHER api_key =os.environ["AI21_API_KEY"] = api_key llm = AI21() # OR llm = AI21(api_key=api_key)

使用消息列表调用 chat¶

消息必须按从旧到新的顺序列出，以 user 角色的消息开始，并在 user 和 assistant 消息之间交替。

In [ ]

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"
)

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" )

In [ ]

print(resp)

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¶

In [ ]

from llama_index.llms.ai21 import AI21

api_key = "Your api key"
resp = AI21(api_key=api_key).complete("Paul Graham is ")

from llama_index.llms.ai21 import AI21 api_key = "Your api key" resp = AI21(api_key=api_key).complete("Paul Graham is ")

In [ ]

print(resp)

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.

调用异步方法¶

In [ ]

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)

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 设置为每次对同一问题生成相同的答案。

In [ ]

from llama_index.llms.ai21 import AI21

llm = AI21(
    model="jamba-1.5-mini", api_key=api_key, max_tokens=100, temperature=0.5
)

from llama_index.llms.ai21 import AI21 llm = AI21( model="jamba-1.5-mini", api_key=api_key, max_tokens=100, temperature=0.5 )

In [ ]

resp = llm.complete("Paul Graham is ")

resp = llm.complete("Paul Graham is ")

In [ ]

print(resp)

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.

流式传输¶

使用 stream_chat 方法按每条消息一个 token 的速度流式传输生成的响应。

In [ ]

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)

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)

In [ ]

for r in resp:
    print(r.delta, end="")

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.

分词器¶

不同的模型使用不同的分词器。

In [ ]

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)

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)

工具调用¶

In [ ]

from llama_index.core.agent import FunctionCallingAgentWorker
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_worker = FunctionCallingAgentWorker.from_tools(
    [multiply_tool, add_tool, subtract_tool, divide_tool],
    llm=llm,
    verbose=True,
    allow_parallel_tool_calls=True,
)
agent = agent_worker.as_agent()

response = agent.chat(
    "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?"
)

from llama_index.core.agent import FunctionCallingAgentWorker 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_worker = FunctionCallingAgentWorker.from_tools( [multiply_tool, add_tool, subtract_tool, divide_tool], llm=llm, verbose=True, allow_parallel_tool_calls=True, ) agent = agent_worker.as_agent() response = agent.chat( "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?" )