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

最新推荐文章于 2024-09-02 23:14:47 发布

qq_37836323

最新推荐文章于 2024-09-02 23:14:47 发布

阅读量244

点赞数 1

文章标签： python 人工智能深度学习

本文链接：https://blog.csdn.net/qq_29929123/article/details/139797579

版权

在本教程中，我们将介绍如何使用LlamaIndex抽象和Cohere进行定制排序器的构建。完成后，您将能够创建一个定制排序器并利用它进行增强的数据检索。

环境设置

首先，我们需要安装必要的软件包。

%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_api_key = "YOUR OPENAI API KEY"
cohere_api_key = "YOUR COHEREAI API KEY"

import os

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

导入所需模块。

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.node_parser import SimpleNodeParser
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.cohere import CohereEmbedding
from llama_index.core.retrievers import BaseRetriever, VectorIndexRetriever
from llama_index.core import QueryBundle
from llama_index.core.indices.query.schema import QueryType
from llama_index.core.schema import NodeWithScore
from llama_index.postprocessor.cohere_rerank import CohereRerank
from llama_index.core.evaluation import EmbeddingQAFinetuneDataset
from llama_index.finetuning import generate_cohere_reranker_finetuning_dataset
from llama_index.core.evaluation import generate_question_context_pairs, RetrieverEvaluator
from llama_index.finetuning import CohereRerankerFinetuneEngine

from typing import List
import pandas as pd
import nest_asyncio

nest_asyncio.apply()

下载和加载数据

我们将使用Lyft和Uber 2021年的10K 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'

加载数据。

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

数据处理与节点创建

将数据分块，每块限制为400个tokens。

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生成问题

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 to the context information provided.\
"""

qa_dataset_lyft_train = generate_question_context_pairs(
    lyft_nodes[:256],
    llm=llm,
    num_questions_per_chunk=1,
    qa_generate_prompt_tmpl=qa_generate_prompt_tmpl,
)

qa_dataset_lyft_train.save_json("lyft_train_dataset.json")

qa_dataset_lyft_val = generate_question_context_pairs(
    lyft_nodes[257:321],
    llm=llm,
    num_questions_per_chunk=1,
    qa_generate_prompt_tmpl=qa_generate_prompt_tmpl,
)

qa_dataset_lyft_val.save_json("lyft_val_dataset.json")

qa_dataset_uber_val = generate_question_context_pairs(
    uber_nodes[:150],
    llm=llm,
    num_questions_per_chunk=1,
    qa_generate_prompt_tmpl=qa_generate_prompt_tmpl,
)

qa_dataset_uber_val.save_json("uber_val_dataset.json")

生成定制排序器训练数据集

embed_model = CohereEmbedding(
    cohere_api_key=cohere_api_key,
    model_name="embed-english-v3.0",
    input_type="search_document",
)

generate_cohere_reranker_finetuning_dataset(
    qa_dataset_lyft_train, finetune_dataset_file_name="train.jsonl"
)

generate_cohere_reranker_finetuning_dataset(
    qa_dataset_lyft_val, finetune_dataset_file_name="val.jsonl"
)

generate_cohere_reranker_finetuning_dataset(
    qa_dataset_lyft_train,
    num_negatives=5,
    hard_negatives_gen_method="random",
    finetune_dataset_file_name="train_5_random.jsonl",
    embed_model=embed_model,
)

generate_cohere_reranker_finetuning_dataset(
    qa_dataset_lyft_val,
    num_negatives=5,
    hard_negatives_gen_method="random",
    finetune_dataset_file_name="val_5_random.jsonl",
    embed_model=embed_model,
)

generate_cohere_reranker_finetuning_dataset(
    qa_dataset_lyft_train,
    num_negatives=5,
    hard_negatives_gen_method="cosine_similarity",
    finetune_dataset_file_name="train_5_cosine_similarity.jsonl",
    embed_model=embed_model,
)

generate_cohere_reranker_finetuning_dataset(
    qa_dataset_lyft_val,
    num_negatives=5,
    hard_negatives_gen_method="cosine_similarity",
    finetune_dataset_file_name="val_5_cosine_similarity.jsonl",
    embed_model=embed_model,
)

训练定制排序器

finetune_model_no_hard_negatives = CohereRerankerFinetuneEngine(
    train_file_name="train.jsonl",
    val_file_name="val.jsonl",
    model_name="lyft_reranker_0_hard_negatives",
    model_type="RERANK",
    base_model="english",
)
finetune_model_no_hard_negatives.finetune()  # 中转API

finetune_model_random_hard_negatives = CohereRerankerFinetuneEngine(
    train_file_name="train_5_random.jsonl",
    val_file_name="val_5_random.jsonl",
    model_name="lyft_reranker_5_random_hard_negatives",
    model_type="RERANK",
    base_model="english",
)
finetune_model_random_hard_negatives.finetune()  # 中转API

finetune_model_cosine_hard_negatives = CohereRerankerFinetuneEngine(
    train_file_name="train_5_cosine_similarity.jsonl",
    val_file_name="val_5_cosine_similarity.jsonl",
    model_name="lyft_reranker_5_cosine_hard_negatives",
    model_type="RERANK",
    base_model="english",
)
finetune_model_cosine_hard_negatives.finetune()  # 中转API

测试定制排序器

reranker_base = CohereRerank(top_n=5)
reranker_model_0 = finetune_model_no_hard_negatives.get_finetuned_model(top_n=5)
reranker_model_5_random = finetune_model_random_hard_negatives.get_finetuned_model(top_n=5)
reranker_model_5_cosine = finetune_model_cosine_hard_negatives.get_finetuned_model(top_n=5)

RERANKERS = {
    "WithoutReranker": "None",
    "CohereRerank": reranker_base,
    "CohereRerank_0": reranker_model_0,
    "CohereRerank_5_random": reranker_model_5_random,
    "CohereRerank_5_cosine": reranker_model_5_cosine,
}

def display_results(embedding_name, reranker_name, eval_results):
    metric_dicts = [eval_result.metric_vals_dict for eval_result in eval_results]
    full_df = pd.DataFrame(metric_dicts)
    hit_rate = full_df["hit_rate"].mean()
    mrr = full_df["mrr"].mean()
    metric_df = pd.DataFrame(
        {
            "Embedding": [embedding_name],
            "Reranker": [reranker_name],
            "hit_rate": [hit_rate],
            "mrr": [mrr],
        }
    )
    return metric_df

index_embed_model = CohereEmbedding(
    cohere_api_key=cohere_api_key,
    model_name="embed-english-v3.0",
    input_type="search_document",
)

query_embed_model = CohereEmbedding(
    cohere_api_key=cohere_api_key,
    model_name="embed-english-v3.0",
    input_type="query",
)

index = VectorStoreIndex.from_documents(
    lyft_docs,
    index_embed_model=index_embed_model,
    query_embed_model=query_embed_model,
    chunk_size=400,
    max_seq_len=512,
    batch_size=4,
)

for reranker_name, reranker in RERANKERS.items():
    if reranker_name == "WithoutReranker":
        retriever = VectorIndexRetriever(index)
    else:
        retriever = VectorIndexRetriever(index, retriever_postprocessor=reranker)
    evaluator = RetrieverEvaluator(retriever=retriever)
    eval_results = evaluator.evaluate(qa_dataset_lyft_val, metric="hit_rate", top_k=5)
    metric_df = display_results("Cohere", reranker_name, eval_results)
    print(metric_df)