Cohere int8 和二进制嵌入检索评估¶

Cohere Embed 是首个原生支持 float、int8、binary 和 ubinary 嵌入的嵌入模型。有关 Cohere int8 & 二进制嵌入的更多详细信息，请参阅他们的主要博客文章。

这个 notebook 帮助您评估这些不同的嵌入类型，并为您的 RAG 流水线选择一个。它使用我们的 RetrieverEvaluator 来使用 LlamaIndex 的 Retriever 模块评估嵌入的质量。

观察到的指标

1. 命中率
2. MRR (平均倒数排名)

对于任何给定的问题，这些指标将比较从真实上下文检索到的结果的质量。评估数据集是使用我们的合成数据集生成模块创建的。我们将使用 GPT-4 进行数据集生成以避免偏差。

注意：notebook 末尾显示的结果非常取决于数据集以及考虑的各种其他参数。我们建议您使用此 notebook 作为参考，在您自己的数据集上进行实验，并评估不同嵌入类型在您的 RAG 流水线中的使用情况。¶

安装¶

In [ ]

%pip install llama-index-llms-openai
%pip install llama-index-embeddings-cohere

%pip install llama-index-llms-openai %pip install llama-index-embeddings-cohere

设置 API 密钥¶

In [ ]

import os

os.environ["OPENAI_API_KEY"] = "YOUR OPENAI KEY"
os.environ["COHERE_API_KEY"] = "YOUR COHEREAI API KEY"

import os os.environ["OPENAI_API_KEY"] = "YOUR OPENAI KEY" os.environ["COHERE_API_KEY"] = "YOUR COHEREAI API KEY"

设置¶

在这里，我们加载数据（PG 论文），解析成节点。然后，我们使用我们的简单向量索引对这些数据进行索引，并为以下不同的嵌入类型获取一个检索器。

1. float
2. int8
3. binary
4. ubinary

In [ ]

import nest_asyncio

nest_asyncio.apply()

import nest_asyncio nest_asyncio.apply()

In [ ]

from llama_index.core.evaluation import generate_question_context_pairs
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.node_parser import SentenceSplitter
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.cohere import CohereEmbedding

from llama_index.core.evaluation import generate_question_context_pairs from llama_index.core import VectorStoreIndex, SimpleDirectoryReader from llama_index.core.node_parser import SentenceSplitter from llama_index.llms.openai import OpenAI from llama_index.embeddings.cohere import CohereEmbedding

下载数据¶

In [ ]

!mkdir -p 'data/paul_graham/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'

!mkdir -p 'data/paul_graham/' !wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'

--2024-03-27 20:26:33--  https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt
Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.111.133, 185.199.108.133, 185.199.110.133, ...
Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.111.133|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 75042 (73K) [text/plain]
Saving to: ‘data/paul_graham/paul_graham_essay.txt’

data/paul_graham/pa 100%[===================>]  73.28K  --.-KB/s    in 0.03s   

2024-03-27 20:26:34 (2.18 MB/s) - ‘data/paul_graham/paul_graham_essay.txt’ saved [75042/75042]

加载数据¶

In [ ]

documents = SimpleDirectoryReader("./data/paul_graham/").load_data()

documents = SimpleDirectoryReader("./data/paul_graham/").load_data()

创建节点¶

In [ ]

node_parser = SentenceSplitter(chunk_size=512)
nodes = node_parser.get_nodes_from_documents(documents)

node_parser = SentenceSplitter(chunk_size=512) nodes = node_parser.get_nodes_from_documents(documents)

In [ ]

# by default, the node ids are set to random uuids. To ensure same id's per run, we manually set them.
for idx, node in enumerate(nodes):
    node.id_ = f"node_{idx}"

# by default, the node ids are set to random uuids. To ensure same id's per run, we manually set them. for idx, node in enumerate(nodes): node.id_ = f"node_{idx}"

为不同嵌入类型创建检索器¶

In [ ]

# llm for question generation
# Take any other llm other than from cohereAI to avoid bias.
llm = OpenAI(model="gpt-4")


# Function to return embedding model
def cohere_embedding(
    model_name: str, input_type: str, embedding_type: str
) -> CohereEmbedding:
    return CohereEmbedding(
        api_key=os.environ["COHERE_API_KEY"],
        model_name=model_name,
        input_type=input_type,
        embedding_type=embedding_type,
    )


# Function to return retriver for different embedding type embedding model
def retriver(nodes, embedding_type="float", model_name="embed-english-v3.0"):
    vector_index = VectorStoreIndex(
        nodes,
        embed_model=cohere_embedding(
            model_name, "search_document", embedding_type
        ),
    )
    retriever = vector_index.as_retriever(
        similarity_top_k=2,
        embed_model=cohere_embedding(
            model_name, "search_query", embedding_type
        ),
    )
    return retriever

# llm for question generation # Take any other llm other than from cohereAI to avoid bias. llm = OpenAI(model="gpt-4") # Function to return embedding model def cohere_embedding( model_name: str, input_type: str, embedding_type: str ) -> CohereEmbedding: return CohereEmbedding( api_key=os.environ["COHERE_API_KEY"], model_name=model_name, input_type=input_type, embedding_type=embedding_type, ) # Function to return retriver for different embedding type embedding model def retriver(nodes, embedding_type="float", model_name="embed-english-v3.0"): vector_index = VectorStoreIndex( nodes, embed_model=cohere_embedding( model_name, "search_document", embedding_type ), ) retriever = vector_index.as_retriever( similarity_top_k=2, embed_model=cohere_embedding( model_name, "search_query", embedding_type ), ) return retriever

In [ ]

# Build retriever for float embedding type
retriver_float = retriver(nodes)

# Build retriever for int8 embedding type
retriver_int8 = retriver(nodes, "int8")

# Build retriever for binary embedding type
retriver_binary = retriver(nodes, "binary")

# Build retriever for ubinary embedding type
retriver_ubinary = retriver(nodes, "ubinary")

# Build retriever for float embedding type retriver_float = retriver(nodes) # Build retriever for int8 embedding type retriver_int8 = retriver(nodes, "int8") # Build retriever for binary embedding type retriver_binary = retriver(nodes, "binary") # Build retriever for ubinary embedding type retriver_ubinary = retriver(nodes, "ubinary")

尝试检索¶

我们将使用 float 检索器尝试对一个示例查询进行检索。

In [ ]

retrieved_nodes = retriver_float.retrieve("What did the author do growing up?")

retrieved_nodes = retriver_float.retrieve("What did the author do growing up?")

In [ ]

from llama_index.core.response.notebook_utils import display_source_node

for node in retrieved_nodes:
    display_source_node(node, source_length=1000)

from llama_index.core.response.notebook_utils import display_source_node for node in retrieved_nodes: display_source_node(node, source_length=1000)

节点 ID: node_2
相似度 0.3641554823852197
文本: 我清晰地记得，看着他坐在电脑前，直接输入程序到电脑里，我感到多么印象深刻和羡慕。

那时候电脑很贵，我缠着父亲很多年才说服他买了一台，大约在 1980 年，是一台 TRS-80。那时的黄金标准是 Apple II，但 TRS-80 也够用了。那时我才真正开始编程。我写了简单的游戏，一个预测我的模型火箭能飞多高的程序，以及一个我父亲用来写书的文字处理器。内存里只能容纳大约 2 页文本，所以他一次写 2 页，然后打印出来，但这比打字机好多了。

虽然我喜欢编程，但我没打算在大学学习它。大学里我打算学习哲学，那听起来强大得多。对我天真的高中生来说，它似乎是研究终极真理的学问，与之相比，其他领域研究的东西不过是领域知识...

节点 ID: node_0
相似度 0.36283154406791923
文本: 我在做什么

2021 年 2 月

在大学之前，我在学校之外做的两件主要事情是写作和编程。我不写论文。我写的是当时（可能现在仍然是）初学者应该写的东西：短篇小说。我的故事很糟糕。它们几乎没有情节，只有情感强烈的角色，我以为这使它们显得深刻。

我尝试编写的第一个程序是在我们学区用于当时称为“数据处理”的 IBM 1401 上。那是在 9 年级，所以我当时 13 或 14 岁。学区的 1401 恰好在我们初中地下室，我和我的朋友 Rich Draves 获得了使用许可。那里就像迷你邦德反派的巢穴，所有这些看起来像是外星人的机器——CPU、磁盘驱动器、打印机、读卡器——都放在高架地板上，在明亮的荧光灯下。

我们使用的语言是 Fortran 的早期版本。你必须在穿孔卡上键入程序，然后堆叠它们...

评估数据集 - (查询, 上下文) 对的合成数据集生成¶

在这里，我们在现有文本语料库上构建一个简单的评估数据集。

我们使用 generate_question_context_pairs 在给定的非结构化文本语料库上生成一组（问题，上下文）对。这利用 LLM 从每个上下文块自动生成问题。

我们会得到一个 EmbeddingQAFinetuneDataset 对象。在高层，它包含一组映射到查询和相关文档块的 ID，以及语料库本身。

In [ ]

from llama_index.core.evaluation import (
    generate_question_context_pairs,
    EmbeddingQAFinetuneDataset,
)

from llama_index.core.evaluation import ( generate_question_context_pairs, EmbeddingQAFinetuneDataset, )

In [ ]

qa_dataset = generate_question_context_pairs(
    nodes, llm=llm, num_questions_per_chunk=2
)

qa_dataset = generate_question_context_pairs( nodes, llm=llm, num_questions_per_chunk=2 )

100%|██████████| 59/59 [04:10<00:00,  4.24s/it]

In [ ]

queries = qa_dataset.queries.values()
print(list(queries)[0])

queries = qa_dataset.queries.values() print(list(queries)[0])

"Describe the author's initial experiences with programming on the IBM 1401. What were some of the challenges he faced and how did these experiences shape his understanding of programming?"

In [ ]

# [optional] save
qa_dataset.save_json("pg_eval_dataset.json")

# [optional] save qa_dataset.save_json("pg_eval_dataset.json")

In [ ]

# [optional] load
qa_dataset = EmbeddingQAFinetuneDataset.from_json("pg_eval_dataset.json")

# [optional] load qa_dataset = EmbeddingQAFinetuneDataset.from_json("pg_eval_dataset.json")

使用 RetrieverEvaluator 进行检索评估¶

现在我们可以运行检索评估了。我们将在生成好的评估数据集上运行 RetrieverEvaluator。

为不同 embedding_types 定义 RetrieverEvaluator¶

In [ ]

from llama_index.core.evaluation import RetrieverEvaluator

metrics = ["mrr", "hit_rate"]

# Retrieval evaluator for float embedding type
retriever_evaluator_float = RetrieverEvaluator.from_metric_names(
    metrics, retriever=retriver_float
)

# Retrieval evaluator for int8 embedding type
retriever_evaluator_int8 = RetrieverEvaluator.from_metric_names(
    metrics, retriever=retriver_int8
)

# Retrieval evaluator for binary embedding type
retriever_evaluator_binary = RetrieverEvaluator.from_metric_names(
    metrics, retriever=retriver_binary
)

# Retrieval evaluator for ubinary embedding type
retriever_evaluator_ubinary = RetrieverEvaluator.from_metric_names(
    metrics, retriever=retriver_ubinary
)

from llama_index.core.evaluation import RetrieverEvaluator metrics = ["mrr", "hit_rate"] # Retrieval evaluator for float embedding type retriever_evaluator_float = RetrieverEvaluator.from_metric_names( metrics, retriever=retriver_float ) # Retrieval evaluator for int8 embedding type retriever_evaluator_int8 = RetrieverEvaluator.from_metric_names( metrics, retriever=retriver_int8 ) # Retrieval evaluator for binary embedding type retriever_evaluator_binary = RetrieverEvaluator.from_metric_names( metrics, retriever=retriver_binary ) # Retrieval evaluator for ubinary embedding type retriever_evaluator_ubinary = RetrieverEvaluator.from_metric_names( metrics, retriever=retriver_ubinary )

In [ ]

# try it out on a sample query
sample_id, sample_query = list(qa_dataset.queries.items())[0]
sample_expected = qa_dataset.relevant_docs[sample_id]

eval_result = retriever_evaluator_float.evaluate(sample_query, sample_expected)
print(eval_result)

# try it out on a sample query sample_id, sample_query = list(qa_dataset.queries.items())[0] sample_expected = qa_dataset.relevant_docs[sample_id] eval_result = retriever_evaluator_float.evaluate(sample_query, sample_expected) print(eval_result)

Query: "Describe the author's initial experiences with programming on the IBM 1401. What were some of the challenges he faced and how did these experiences shape his understanding of programming?"
Metrics: {'mrr': 0.5, 'hit_rate': 1.0}

In [ ]

# Evaluation on the entire dataset

# float embedding type
eval_results_float = await retriever_evaluator_float.aevaluate_dataset(
    qa_dataset
)

# int8 embedding type
eval_results_int8 = await retriever_evaluator_int8.aevaluate_dataset(
    qa_dataset
)

# binary embedding type
eval_results_binary = await retriever_evaluator_binary.aevaluate_dataset(
    qa_dataset
)

# ubinary embedding type
eval_results_ubinary = await retriever_evaluator_ubinary.aevaluate_dataset(
    qa_dataset
)

# Evaluation on the entire dataset # float embedding type eval_results_float = await retriever_evaluator_float.aevaluate_dataset( qa_dataset ) # int8 embedding type eval_results_int8 = await retriever_evaluator_int8.aevaluate_dataset( qa_dataset ) # binary embedding type eval_results_binary = await retriever_evaluator_binary.aevaluate_dataset( qa_dataset ) # ubinary embedding type eval_results_ubinary = await retriever_evaluator_ubinary.aevaluate_dataset( qa_dataset )

定义 display_results 以在 dataframe 中显示每个检索器的结果。¶

In [ ]

import pandas as pd


def display_results(name, eval_results):
    """Display results from evaluate."""

    metric_dicts = []
    for eval_result in eval_results:
        metric_dict = eval_result.metric_vals_dict
        metric_dicts.append(metric_dict)

    full_df = pd.DataFrame(metric_dicts)

    hit_rate = full_df["hit_rate"].mean()
    mrr = full_df["mrr"].mean()
    columns = {"Embedding Type": [name], "hit_rate": [hit_rate], "mrr": [mrr]}

    metric_df = pd.DataFrame(columns)

    return metric_df

import pandas as pd def display_results(name, eval_results): """Display results from evaluate.""" metric_dicts = [] for eval_result in eval_results: metric_dict = eval_result.metric_vals_dict metric_dicts.append(metric_dict) full_df = pd.DataFrame(metric_dicts) hit_rate = full_df["hit_rate"].mean() mrr = full_df["mrr"].mean() columns = {"Embedding Type": [name], "hit_rate": [hit_rate], "mrr": [mrr]} metric_df = pd.DataFrame(columns) return metric_df

评估结果¶

In [ ]

# metrics for float embedding type
metrics_float = display_results("float", eval_results_float)

# metrics for int8 embedding type
metrics_int8 = display_results("int8", eval_results_int8)

# metrics for binary embedding type
metrics_binary = display_results("binary", eval_results_binary)

# metrics for ubinary embedding type
metrics_ubinary = display_results("ubinary", eval_results_ubinary)

# metrics for float embedding type metrics_float = display_results("float", eval_results_float) # metrics for int8 embedding type metrics_int8 = display_results("int8", eval_results_int8) # metrics for binary embedding type metrics_binary = display_results("binary", eval_results_binary) # metrics for ubinary embedding type metrics_ubinary = display_results("ubinary", eval_results_ubinary)

In [ ]

combined_metrics = pd.concat(
    [metrics_float, metrics_int8, metrics_binary, metrics_ubinary]
)
combined_metrics.set_index(["Embedding Type"], append=True, inplace=True)

combined_metrics = pd.concat( [metrics_float, metrics_int8, metrics_binary, metrics_ubinary] ) combined_metrics.set_index(["Embedding Type"], append=True, inplace=True)

In [ ]

combined_metrics

combined_metrics

Out[ ]

		hit_rate	mrr
	Embedding Type
0	float	0.805085	0.665254
	int8	0.813559	0.673729
	binary	0.491525	0.394068
	ubinary	0.449153	0.377119

注意：上面显示的结果非常取决于数据集以及考虑的各种其他参数。我们建议您使用此 notebook 作为参考，在您自己的数据集上进行实验，并评估不同嵌入类型在您的 RAG 流水线中的使用情况。¶