使用LlamaIndex和Cohere进行自定义重排序器的构建指南

在当今的AI技术领域，构建高效的数据检索系统是一个关键任务。本篇文章将介绍如何使用LlamaIndex和Cohere构建一个自定义的重排序器，提升数据检索的效果。本文将详细介绍所需的步骤，并提供示例代码供参考。

安装必要的包

首先，我们需要安装一些必要的Python包：

%pip install llama-index-postprocessor-cohere-rerank
%pip install llama-index-llms-openai
%pip install llama-index-finetuning
%pip install llama-index-embeddings-cohere
!pip install llama-index cohere pypdf

初始化API密钥

接下来，我们需要初始化OpenAI和Cohere的API密钥：

import os

openai_api_key = "YOUR_OPENAI_API_KEY"
cohere_api_key = "YOUR_COHERE_API_KEY"

os.environ["OPENAI_API_KEY"] = openai_api_key
os.environ["COHERE_API_KEY"] = cohere_api_key

下载数据

我们将使用Lyft 2021和Uber 2021的10-K SEC文件进行训练和评估：

!mkdir -p 'data/10k/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10k/uber_2021.pdf' -O 'data/10k/uber_2021.pdf'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10k/lyft_2021.pdf' -O 'data/10k/lyft_2021.pdf'

加载数据

使用SimpleDirectoryReader加载数据：

from llama_index.core import SimpleDirectoryReader

lyft_docs = SimpleDirectoryReader(input_files=["./data/10k/lyft_2021.pdf"]).load_data()
uber_docs = SimpleDirectoryReader(input_files=["./data/10k/uber_2021.pdf"]).load_data()

创建节点

我们将文档拆分成小块以创建训练数据集：

from llama_index.core.node_parser import SimpleNodeParser

node_parser = SimpleNodeParser.from_defaults(chunk_size=400)
lyft_nodes = node_parser.get_nodes_from_documents(lyft_docs)
uber_nodes = node_parser.get_nodes_from_documents(uber_docs)

使用GPT-4生成问题

使用OpenAI的GPT-4模型从每个节点生成问题：

from llama_index.llms.openai import OpenAI

llm = OpenAI(temperature=0, model="gpt-4")
qa_generate_prompt_tmpl = """\
Context information is below.

---------------------
{context_str}
---------------------

Given the context information and not prior knowledge, generate only questions based on the below query.

You are a Professor. Your task is to setup {num_questions_per_chunk} questions for an upcoming quiz/examination. The questions should be diverse in nature across the document. The questions should not contain options, not start with Q1/ Q2. Restrict the questions