OpenAI检索API（通过助理代理）的基准测试¶

%pip install llama-index-readers-file pymupdf
%pip install llama-index-agent-openai
%pip install llama-index-llms-openai

!pip install llama-index

import nest_asyncio

nest_asyncio.apply()

from pathlib import Path
from llama_index.core import Document, VectorStoreIndex
from llama_index.readers.file import PyMuPDFReader
from llama_index.core.node_parser import SimpleNodeParser
from llama_index.llms.openai import OpenAI

loader = PyMuPDFReader()
docs0 = loader.load(file_path=Path("./data/llama2.pdf"))

doc_text = "\n\n".join([d.get_content() for d in docs0])
docs = [Document(text=doc_text)]

node_parser = SimpleNodeParser.from_defaults()
nodes = node_parser.get_nodes_from_documents(docs)

from llama_index.core.evaluation import QueryResponseDataset# 可选eval_dataset = QueryResponseDataset.from_json(    "data/llama2_eval_qr_dataset.json")

from llama_index.core.evaluation import DatasetGenerator, QueryResponseDataset
from llama_index.llms.openai import OpenAI

# 注意：如果数据集尚未保存，请运行此代码# 注意：我们只从前20个节点生成，因为其余的是引用llm = OpenAI(model="gpt-4-1106-preview")dataset_generator = DatasetGenerator(    nodes[:20],    llm=llm,    show_progress=True,    num_questions_per_chunk=3,)eval_dataset = await dataset_generator.agenerate_dataset_from_nodes(num=60)eval_dataset.save_json("data/llama2_eval_qr_dataset.json")

# 可选eval_dataset = QueryResponseDataset.from_json(    "data/llama2_eval_qr_dataset.json")

from llama_index.core.evaluation.eval_utils import (
    get_responses,
    get_results_df,
)
from llama_index.core.evaluation import (
    CorrectnessEvaluator,
    SemanticSimilarityEvaluator,
    BatchEvalRunner,
)
from llama_index.llms.openai import OpenAI

eval_llm = OpenAI(model="gpt-4-1106-preview")
evaluator_c = CorrectnessEvaluator(llm=eval_llm)
evaluator_s = SemanticSimilarityEvaluator(llm=eval_llm)
evaluator_dict = {
    "correctness": evaluator_c,
    "semantic_similarity": evaluator_s,
}
batch_runner = BatchEvalRunner(evaluator_dict, workers=2, show_progress=True)

import numpy as npimport timeimport osimport picklefrom tqdm import tqdmdef get_responses_sync(    eval_qs, query_engine, show_progress=True, save_path=None):    if show_progress:        eval_qs_iter = tqdm(eval_qs)    else:        eval_qs_iter = eval_qs    pred_responses = []    start_time = time.time()    for eval_q in eval_qs_iter:        print(f"eval q: {eval_q}")        pred_response = agent.query(eval_q)        print(f"predicted response: {pred_response}")        pred_responses.append(pred_response)        if save_path is not None:            # save intermediate responses (to cache in case something breaks)            avg_time = (time.time() - start_time) / len(pred_responses)            pickle.dump(                {"pred_responses": pred_responses, "avg_time": avg_time},                open(save_path, "wb"),            )    return pred_responsesasync def run_evals(    query_engine,    eval_qa_pairs,    batch_runner,    disable_async_for_preds=False,    save_path=None,):    # then evaluate    # TODO: evaluate a sample of generated results    eval_qs = [q for q, _ in eval_qa_pairs]    eval_answers = [a for _, a in eval_qa_pairs]    if save_path is not None:        if not os.path.exists(save_path):            start_time = time.time()            if disable_async_for_preds:                pred_responses = get_responses_sync(                    eval_qs,                    query_engine,                    show_progress=True,                    save_path=save_path,                )            else:                pred_responses = get_responses(                    eval_qs, query_engine, show_progress=True                )            avg_time = (time.time() - start_time) / len(eval_qs)            pickle.dump(                {"pred_responses": pred_responses, "avg_time": avg_time},                open(save_path, "wb"),            )        else:            # [optional] load            pickled_dict = pickle.load(open(save_path, "rb"))            pred_responses = pickled_dict["pred_responses"]            avg_time = pickled_dict["avg_time"]    else:        start_time = time.time()        pred_responses = get_responses(            eval_qs, query_engine, show_progress=True        )        avg_time = (time.time() - start_time) / len(eval_qs)    eval_results = await batch_runner.aevaluate_responses(        eval_qs, responses=pred_responses, reference=eval_answers    )    return eval_results, {"avg_time": avg_time}

from llama_index.agent.openai import OpenAIAssistantAgent

agent = OpenAIAssistantAgent.from_new(
    name="SEC Analyst",
    instructions="You are a QA assistant designed to analyze sec filings.",
    openai_tools=[{"type": "retrieval"}],
    instructions_prefix="Please address the user as Jerry.",
    files=["data/llama2.pdf"],
    verbose=True,
)

response = agent.query(
    "What are the key differences between Llama 2 and Llama 2-Chat?"
)

llm = OpenAI(model="gpt-4-1106-preview")
base_index = VectorStoreIndex(nodes)
base_query_engine = base_index.as_query_engine(similarity_top_k=2, llm=llm)

base_eval_results, base_extra_info = await run_evals(
    base_query_engine,
    eval_dataset.qr_pairs,
    batch_runner,
    save_path="data/llama2_preds_base.pkl",
)

assistant_eval_results, assistant_extra_info = await run_evals(
    agent,
    eval_dataset.qr_pairs[:55],
    batch_runner,
    save_path="data/llama2_preds_assistant.pkl",
    disable_async_for_preds=True,
)