1 line
15 KiB
JSON
1 line
15 KiB
JSON
{"repo": "chaoql/rag-best-practices", "files": {"FTEmbed.py": {"cbe12e21df644ecac67b1ec9ff03014161f1f78e": {"t": 1785674170.0730977, "n": 5787, "h": "0d2365f678b56e05", "c": "import json\nimport os\n\nfrom llama_index.core import SimpleDirectoryReader\nfrom llama_index.core.node_parser import SentenceSplitter\nfrom llama_index.core.schema import MetadataMode\nfrom llama_index.finetuning import generate_qa_embedding_pairs\nfrom llama_index.core.evaluation import EmbeddingQAFinetuneDataset\nfrom llama_index.llms.openai import OpenAI\nfrom dotenv import load_dotenv, find_dotenv\nfrom dotenv import dotenv_values\nfrom custom.glmfz import ChatGLM\nfrom llama_index.core import Settings\nfrom llama_index.finetuning import SentenceTransformersFinetuneEngine\nfrom llama_index.core import ServiceContext, VectorStoreIndex\nfrom llama_index.core.schema import TextNode\nfrom tqdm.notebook import tqdm\nfrom sentence_transformers.evaluation import InformationRetrievalEvaluator\nfrom sentence_transformers import SentenceTransformer\nfrom pathlib import Path\nimport pandas as pd\nfrom llama_index.embeddings.huggingface import HuggingFaceEmbedding\n\n\ndef load_corpus(files, verbose=False):\n if verbose:\n print(f\"Loading files {files}\")\n\n reader = SimpleDirectoryReader(input_files=files)\n docs = reader.load_data()\n if verbose:\n print(f\"Loaded {len(docs)} docs\")\n\n parser = SentenceSplitter()\n nodes = parser.get_nodes_from_documents(docs, show_progress=verbose)\n\n if verbose:\n print(f\"Parsed {len(nodes)} nodes\")\n\n return nodes\n\n\ndef evaluate(\n dataset,\n embed_model,\n top_k=5,\n verbose=False,\n):\n corpus = dataset.corpus\n queries = dataset.queries\n relevant_docs = dataset.relevant_docs\n\n nodes = [TextNode(id_=id_, text=text) for id_, text in corpus.items()]\n index = VectorStoreIndex(\n nodes, embed_model=embed_model, show_progress=True\n )\n retriever = index.as_retriever(similarity_top_k=top_k)\n\n eval_results = []\n for query_id, query in tqdm(queries.items()):\n retrieved_nodes = retriever.retrieve(query)\n retrieved_ids = [node.node.node_id for node in retrieved_nodes]\n expected_id = relevant_docs[query_id][0]\n is_hit = expected_id in retrieved_ids # assume 1 relevant doc\n\n eval_result = {\n \"is_hit\": is_hit,\n \"retrieved\": retrieved_ids,\n \"expected\": expected_id,\n \"query\": query_id,\n }\n eval_results.append(eval_result)\n return eval_results\n\n\ndef evaluate_st(\n dataset,\n model_id,\n name,\n):\n corpus = dataset.corpus\n queries = dataset.queries\n relevant_docs = dataset.relevant_docs\n\n evaluator = InformationRetrievalEvaluator(\n queries, corpus, relevant_docs, name=name\n )\n model = SentenceTransformer(model_id)\n output_path = \"results/\"\n Path(output_path).mkdir(exist_ok=True, parents=True)\n return evaluator(model, output_path=output_path)\n\n\n_ = load_dotenv(find_dotenv()) # 导入环境\nconfig = dotenv_values(\".env\")\nSettings.llm = ChatGLM(\n api_key=config[\"GLM_KEY\"],\n model=\"glm-4\",\n api_base=\"https://open.bigmodel.cn/api/paas/v4/\",\n is_chat_model=True,\n)\nSettings.embed_model = HuggingFaceEmbedding(\n model_name=\"BAAI/bge-large-en-v1.5\",\n cache_folder=\"./BAAI/\",\n embed_batch_size=128,\n local_files_only=True, # 仅加载本地模型，不尝试下载\n device=\"cuda\",\n)\n\n\ndef finetuning_data_preparation(all_data: list, llm, verbose: bool = False,\n train_dataset_dir: str = \"ft_data/train_dataset.json\",\n val_dataset_dir: str = \"ft_data/val_dataset.json\",\n qa_finetune_train_dataset_dir: str = \"ft_data/qa_finetune_train_dataset.json\",\n qa_finetune_val_dataset_dir: str = \"ft_data/qa_finetune_val_dataset.json\"):\n docs_nodes = load_corpus(all_data, verbose=verbose)\n train_nodes = docs_nodes[:int(2 * len(docs_nodes) / 3)]\n val_nodes = docs_nodes[int(2 * len(docs_nodes) / 3 + 1):]\n\n # 自带保存\n train_dataset = generate_qa_embedding_pairs(llm=llm, nodes=train_nodes, verbose=verbose,\n output_path=\"ft_data/qa_finetune_train_dataset.json\")\n val_dataset = generate_qa_embedding_pairs(llm=llm, nodes=val_nodes, verbose=verbose,\n output_path=\"ft_data/qa_finetune_val_dataset.json\")\n train_dataset.save_json(train_dataset_dir)\n val_dataset.save_json(val_dataset_dir)\n\n\ndef finetuning_embedding(train_dataset_dir: str = \"ft_data/train_dataset.json\",\n val_dataset_dir: str = \"ft_data/val_dataset.json\",\n model_name: str = \"BAAI/bge-large-en-v1.5\",\n model_output_path: str = \"BAAI/ft-bge-large-en-v1.5\"):\n # [Optional] Load\n train_dataset = EmbeddingQAFinetuneDataset.from_json(train_dataset_dir)\n val_dataset = EmbeddingQAFinetuneDataset.from_json(val_dataset_dir)\n finetune_engine = SentenceTransformersFinetuneEngine(\n train_dataset,\n model_id=model_name,\n model_output_path=model_output_path,\n val_dataset=val_dataset,\n )\n if not os.path.exists(model_output_path): # 如果已经存在微调好的模型就不用重新微调了\n finetune_engine.finetune()\n embed_model = finetune_engine.get_finetuned_model()\n print(embed_model)\n return embed_model\n\n\ndef eval_finetuning_embedding(embed_model, val_dataset_dir: str = \"ft_data/val_dataset.json\",\n model_name: str = \"bge-large-en-v1.5\"):\n val_dataset = EmbeddingQAFinetuneDataset.from_json(val_dataset_dir)\n bge_val_results = evaluate(val_dataset, embed_model)\n df_bge = pd.DataFrame(bge_val_results)\n hit_rate_bge = df_bge['is_hit'].mean()\n print(f\"{model_name}模型的准确率为：{hit_rate_bge}\")\n # bge-large-en-v1.5模型的准确率为：0.6155913978494624\n # 微调后的bge-large-en-v1.5模型的准确率为：0.7093023255813954\n"}}, "main.py": {"cbe12e21df644ecac67b1ec9ff03014161f1f78e": {"t": 1785674170.634535, "n": 8318, "h": "0dc2df015d936c4f", "c": "from llama_index.embeddings.huggingface import HuggingFaceEmbedding\nfrom llama_index.llms.ollama import Ollama\nfrom llama_index.core.node_parser import SentenceWindowNodeParser\nfrom llama_index.core import PromptTemplate, get_response_synthesizer, StorageContext, VectorStoreIndex, \\\n SimpleDirectoryReader, Settings\nfrom FTEmbed import finetuning_data_preparation, finetuning_embedding, eval_finetuning_embedding\nfrom llama_index.core.indices.query.query_transform import HyDEQueryTransform\nfrom llama_index.core.query_engine import TransformQueryEngine\nfrom llama_index.core.response_synthesizers.type import ResponseMode\nfrom custom.glmfz import ChatGLM\nfrom llama_index.core.tools import QueryEngineTool\nfrom llama_index.core.query_engine import RouterQueryEngine\nfrom llama_index.core.selectors import LLMSingleSelector\nfrom dotenv import load_dotenv, find_dotenv\nfrom dotenv import dotenv_values\nfrom custom.query import build_query_engine\nfrom custom.prompt import qa_prompt_tmpl_str, simple_qa_prompt_tmpl_str, rag_description, norag_rag_description\nfrom utils import load_hybrid_data, load_txt_data\nimport warnings\n\nwarnings.filterwarnings('ignore')\n# 导入环境\n_ = load_dotenv(find_dotenv())\nconfig = dotenv_values(\".env\")\n\n# 设置参数\nwith_hyde = False # 是否采用假设文档\npersist_dir = \"store\" # 向量存储地址\nwith_hybrid_search = True # 是否采用混合检索\n# 5选2\ntop_k = 5\ntop_k_rerank = 2\nwith_sliding_window = True # 是否采用滑动窗口\nresponse_mode = ResponseMode.SIMPLE_SUMMARIZE # RAG架构，最佳实践为为TREE_SUMMARIZE\nwith_query_classification = False # 是否对输入的问题进行分类\nwith_rerank = True # 是否采用重排序\nwith_local_llm = False # 是否采用本地基于Ollama的大模型\nwith_Finetuning_embedding = False # 是否微调嵌入模型\nwith_Finetuning_embedding_eval = False # 是否测评微调嵌入模型的命中率\n# 提问\n# query_str = \"Did Fang Hung-chien kiss Miss Bao?\"\nquery_str = \"In the text, which lady did Fang Hongjian kiss on the ship, and under what circumstances did it happen?\"\n\n# 加载嵌入模型\nif with_Finetuning_embedding:\n # 微调需要开启VPN\n finetuning_data_preparation(all_data=[\"data/testdata.txt\"], llm=Settings.llm, verbose=False,\n train_dataset_dir=\"ft_data/train_dataset.json\",\n val_dataset_dir=\"ft_data/val_dataset.json\",\n qa_finetune_train_dataset_dir=\"ft_data/qa_finetune_train_dataset.json\",\n qa_finetune_val_dataset_dir=\"ft_data/qa_finetune_val_dataset.json\")\n Settings.embed_model = finetuning_embedding(train_dataset_dir=\"ft_data/train_dataset.json\",\n val_dataset_dir=\"ft_data/val_dataset.json\",\n model_name=\"BAAI/bge-large-en-v1.5\",\n model_output_path=\"BAAI/ft-bge-large-en-v1.5/\")\n # 测评微调后的嵌入模型命中率\n if with_Finetuning_embedding_eval:\n eval_finetuning_embedding(embed_model=Settings.embed_model, val_dataset_dir=\"ft_data/val_dataset.json\",\n model_name=\"ft-bge-large-en-v1.5\")\nelse:\n Settings.embed_model = HuggingFaceEmbedding(\n model_name=\"BAAI/bge-large-en-v1.5\",\n cache_folder=\"./BAAI/\",\n embed_batch_size=128,\n local_files_only=True, # 仅加载本地模型，不尝试下载\n device=\"cuda\",\n )\n\n# 加载大模型\nif with_local_llm:\n Settings.llm = Ollama(model=\"qwen2:1.5b\", request_timeout=30.0, temperature=0)\nelse:\n Settings.llm = ChatGLM(\n api_key=config[\"GLM_KEY\"],\n model=\"glm-4\",\n api_base=\"https://open.bigmodel.cn/api/paas/v4/\",\n is_chat_model=True,\n )\n\n# load data(是否采用混合检索)\nif with_hybrid_search:\n index, nodes = load_hybrid_data(input_file=[\"data/testdata.txt\"], persist_dir=\"hybrid_Store\")\nelse:\n index, nodes = load_txt_data(input_file=[\"data/testdata.txt\"], persist_dir=\"hybrid_Store\",\n with_sliding_window=with_sliding_window, chunk_size=512, chunk_overlap=128)\n\n# prompt\nqa_prompt_tmpl = PromptTemplate(qa_prompt_tmpl_str)\nsimple_qa_prompt_tmpl = PromptTemplate(simple_qa_prompt_tmpl_str) # norag\n\n# Build query_engine\nrag_query_engine = build_query_engine(index, response_mode, qa_prompt_tmpl, with_hybrid_search, top_k,\n top_k_rerank, with_rerank, nodes)\nsimple_query_engine = index.as_query_engine(similarity_top_k=top_k,\n text_qa_template=simple_qa_prompt_tmpl,\n response_synthesizer=get_response_synthesizer(\n response_mode=ResponseMode.GENERATION),\n )\n\n# HyDE(当问题较为简单时，不需要该模块参与)\nif with_hyde:\n hyde = HyDEQueryTransform(include_original=True)\n rag_query_engine = TransformQueryEngine(rag_query_engine, hyde)\n\n# response\n# Router Query Engine(Query Classification)\nif with_query_classification:\n rag_tool = QueryEngineTool.from_defaults(\n query_engine=rag_query_engine,\n description=rag_description,\n )\n simple_tool = QueryEngineTool.from_defaults(\n query_engine=simple_query_engine,\n description=norag_rag_description,\n )\n query_engine = RouterQueryEngine(\n selector=LLMSingleSelector.from_defaults(),\n query_engine_tools=[\n rag_tool,\n simple_tool,\n ],\n )\n response = query_engine.query(query_str)\nelse:\n response = rag_query_engine.query(query_str)\n\nprint(f\"Question: {str(query_str)}\")\nprint(\"------------------\")\nprint(f\"Response: {str(response)}\")\nprint(\"------------------\")\nif not with_query_classification or response.metadata['selector_result'].ind == 0:\n window = response.source_nodes[0].node.metadata[\"window\"] # 长度为3的窗口，包含了文本两侧的上下文。\n sentence = response.source_nodes[0].node.metadata[\"original_sentence\"] # 检索到的文本\n print(f\"Window: {window}\")\n print(\"------------------\")\n print(f\"Original Sentence: {sentence}\")\n print(\"------------------\")\n\n\"\"\"\n示例1：\nQuestion: Did Fang Hung-chien kiss Miss Bao?\n------------------\nResponse: Yes, Fang Hung-chien kissed Miss Bao.\n------------------\nWindow: A big wave shook the hull badly, and Miss Bao could not stand steadily. Fang hung-chien hooked her waist and stayed by the railing, kissing her greedily. Miss Bao's lips suggested that the body followed, and this hasty and rude kiss gradually stabilized and grew properly and densely. Miss Bao deftly pushed off Fang Hung-chien's arm, took a deep breath in her mouth and said, \"I'm suffocated by you! I have a cold and I can't breathe in my nose-it's too cheap for you, and you haven't asked me to love you! \"\n \"I beg you now, ok?\" It seems that all men who have never been in love, Fang Hung-chien regards the word \"love\" too noble and serious and refuses to apply it to women casually; He only felt that he wanted Miss Bao and didn't love her, so he was so evasive.\n------------------\nOriginal Sentence: Miss Bao deftly pushed off Fang Hung-chien's arm, took a deep breath in her mouth and said, \"I'm suffocated by you! \n------------------\n\n示例2：\nQuestion: In the text, which lady did Fang Hongjian kiss on the ship, and under what circumstances did it happen?\n------------------\nResponse: In the text, Fang Hongjian kissed Miss Bao on the ship. It happened in the dark shadows of the deck, after a big wave shook the hull, causing Miss Bao to lose her balance. Fang Hongjian hooked her waist to steady her and then kissed her.\n------------------\nWindow: After ten o'clock, there were only three or five pairs of men and women on the deck, all hiding in the dark shadows where the lights could not shine. Fang Hung-chien and Miss Bao walked side by side without talking. A big wave shook the hull badly, and Miss Bao could not stand steadily. Fang hung-chien hooked her waist and stayed by the railing, kissing her greedily. Miss Bao's lips suggested that the body followed, and this hasty and rude kiss gradually stabilized and grew properly and densely. Miss Bao deftly pushed off Fang Hung-chien's arm, took a deep breath in her mouth and said, \"I'm suffocated by you! I have a cold and I can't breathe in my nose-it's too cheap for you, and you haven't asked me to love you! \"\n------------------\nOriginal Sentence: Fang hung-chien hooked her waist and stayed by the railing, kissing her greedily. \n------------------\n\"\"\"\n"}}}} |