Jamba-Instruct 在 LlamaIndex 上的 256k 上下文窗口

利用 LlamaIndex 市场领先的 RAG 策略与 AI21 Labs 的长上下文基础模型 Jamba-Instruct，为企业构建最先进的 RAG 应用。

我们 AI21 Labs 很高兴地宣布，我们开创性的 Jamba-Instruct 基础模型现已通过领先的数据框架 LlamaIndex 提供。通过此次集成，开发者现在可以构建功能强大的企业 RAG 应用，借助 Jamba-Instruct 令人印象深刻的 256K 上下文窗口和 LlamaIndex 为 RAG 提供的复杂端到端解决方案，提高准确性和成本效益。

尽管许多模型声称拥有长上下文窗口，但英伟达的研究人员发现，大多数在评估时表现不佳，这揭示了其宣称的上下文窗口长度与有效长度之间的差异。Jamba-Instruct 是市场上少数不仅能实现宣称长度与有效长度一致，而且其上下文窗口长度远超同类其他任何模型的模型之一。

通过提供一个 256K 的上下文窗口（大致相当于 800 页文本），Jamba-Instruct 增加了检索到的块（chunk）数量，并且可以极大地改进整个 RAG 系统，而不是仅仅尝试改进搜索机制或加入额外的重新排名组件。使用像 Jamba-Instruct 这样的长上下文基础模型使得使用 RAG 查询私有企业数据更加可靠和容易。

在以下 Notebook (也可直接在 colab 上获取) 中，我们将演示一个查询金融文档集合的示例，展示 Jamba-Instruct 的 256K 上下文窗口如何让 RAG 管道一次检索更多块，从而提供准确的答案。

金融文档上的 RAG 问答

要开始使用，您需要安装这些软件包。您还需要 API 密钥来设置 OpenAI 的嵌入和 AI21 的 Jamba-Instruct。

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

import os
from llama_index.core.llama_dataset import download_llama_dataset
from llama_index.core.llama_pack import download_llama_pack
from llama_index.core import VectorStoreIndex
from llama_index.core import SimpleDirectoryReader
from llama_index.llms.ai21 import AI21

os.environ['OPENAI_API_KEY'] = 'YOUR_OPENAI_API_KEY' # For embeddings
os.environ['AI21_API_KEY'] = 'YOUR_AI21_API_KEY' # For the generation

# Setup jamba instruct as the llm
llm = AI21(
    model='jamba-instruct',
    temperature=0,
    max_tokens=2000
)

接下来，从 亚马逊投资者关系页面 下载 5 份亚马逊的 10-K 表格。

# Get the data - download 10k forms from AMZN from the last five years
os.mkdir("data")
!wget 'https://d18rn0p25nwr6d.cloudfront.net/CIK-0001018724/c7c14359-36fa-40c3-b3ca-5bf7f3fa0b96.pdf' -O 'data/amazon_2023.pdf'
!wget 'https://d18rn0p25nwr6d.cloudfront.net/CIK-0001018724/d2fde7ee-05f7-419d-9ce8-186de4c96e25.pdf' -O 'data/amazon_2022.pdf'
!wget 'https://d18rn0p25nwr6d.cloudfront.net/CIK-0001018724/f965e5c3-fded-45d3-bbdb-f750f156dcc9.pdf' -O 'data/amazon_2021.pdf'
!wget 'https://d18rn0p25nwr6d.cloudfront.net/CIK-0001018724/336d8745-ea82-40a5-9acc-1a89df23d0f3.pdf' -O 'data/amazon_2020.pdf'
!wget 'https://d18rn0p25nwr6d.cloudfront.net/CIK-0001018724/4d39f579-19d8-4119-b087-ee618abf82d6.pdf' -O 'data/amazon_2019.pdf'

设置您的索引和查询引擎，以创建 RAG 系统的检索和生成组件。

# Setup the index
file_list = [os.path.join("data", f) for f in os.listdir("data")]

amzn_10k_docs = SimpleDirectoryReader(input_files=file_list).load_data()
index = VectorStoreIndex.from_documents(documents=amzn_10k_docs)

# Build a query engine
default_query_engine = index.as_query_engine(llm)

让我们输入一个查询来确保我们的 RAG 系统正在工作。

answer = default_query_engine.query("What was the company's revenue in 2021?")
print(answer.response)

The company's revenue in 2021 was $469,822 million.

太棒了！它奏效了。现在让我们尝试一个类似的查询来继续验证。

answer = default_query_engine.query("What was the company's revenue in 2023?")
print(answer.response)

The company's revenue in 2023 was not explicitly mentioned in the provided context. However, it is mentioned that the company's operating income increased to $36.9 billion in 2023, compared to $12.2 billion in 2022.

我们可以看到存在一个问题——我们知道问题的答案肯定包含在我们的文档中，但我们的 RAG 系统却声称找不到答案。这是因为默认检索到的块（chunk）数量相当少（只有几个块）。这使得整个系统容易出错，并且无法捕获文档中确实存在的信息。

然而，借助 Jamba-Instruct 这个能够有效处理 256K 上下文窗口的模型，我们可以将检索到的块数量从几个（默认值）增加到 100，从而极大地改进整个 RAG 系统。

让我们在我们现有索引的基础上构建一个新的查询引擎，并尝试之前失败的查询。

# Large amount of chunks in the retrieval process
extended_query_engine = index.as_query_engine(llm,
                                              similarity_top_k=100)

answer = extended_query_engine.query("What was the company's revenue in 2023?")
print(answer.response)

The company's revenue in 2023 was $574.785 million.

我们看到 RAG 系统在 Jamba-Instruct 的 256K 上下文窗口的帮助下，现在能够生成准确的答案了。

让我们再尝试一个答案来验证我们的新 RAG 系统。

answer = default_query_engine.query("Was there a stock split in the last five years?")
print(answer.response)

No, there was no stock split in the last five years.

answer = extended_query_engine.query("Was there a stock split in the last five years?")
print(answer.response)

Yes, there was a stock split in the last five years. On May 27, 2022, Amazon.com, Inc. effected a 20-for-1 stock split of its common stock.

上下文为王

通常，争论被框定为“RAG 与长上下文”之争。我们 AI21 Labs 认为这种看法是错误的。相反，应该是长上下文 加上 RAG。当在 AI 系统中结合使用时，长上下文模型可以提高 RAG 系统的质量和准确性，这在涉及冗长文档或海量信息数据库的企业环境中尤为有用。

未来，随着 RAG 系统不断扩展，文档数量和块（chunk）长度将大幅增加。只有能够真正实现其上下文长度的长上下文模型才能处理如此大量的文本。