RAG(Retrieval-Augmented Generation)是一种结合了信息检索和生成模型的技术,旨在提高生成模型的知识获取和生成能力。它通过在生成的过程中引入外部知识库或文档(如数据库、搜索引擎或文档存储),帮助生成更为准确和丰富的答案。 RAG 在自然语言处理(NLP)领域,特别是在对话生成、问答系统和文本摘要等任务中,具有非常重要的应用。它的核心思想是,生成模型不仅依赖于模型内部的知识,还可以在生成过程中动态地检索外部信息,从而增强生成内容的准确性和信息丰富度。
示例
下面这段代码实现了一个完整的工作流,涉及到以下几个关键步骤:
1. 创建 Elasticsearch 索引
- 首先,代码检查是否已经存在名为
my_index4
的 Elasticsearch 索引。如果存在,则先删除该索引,确保创建的是新的索引。 - 接着,创建一个新的索引
my_index4
,并定义其映射(mappings):id
字段是文本类型。content
字段是text
类型,并指定了ik_max_word
分词器,用于中文分词。content_vector
是dense_vector
类型,表示用于存储文档的嵌入向量,并使用余弦相似度进行向量检索。
2. 加载和处理 PDF 文件
- 使用
pdfjsLib
库加载本地的 PDF 文件(t.pdf
),并从中提取文本内容。文本按页面分批读取,直到所有页面的内容都被提取出来。 - 提取的文本内容会经过
sentence-splitter
库进行分句处理,将文档拆分成句子节点,便于后续的处理。
3. 嵌入向量化
- 为了让文本内容可以进行语义检索,代码使用了 OpenAI 的 Embeddings API(调用
openai.embeddings.create
)来生成文本的嵌入向量(embedding)。 - 该嵌入向量将被存储在 Elasticsearch 中,以便后续进行基于向量的检索。
- 将每个句子的文本与对应的嵌入向量一同批量插入 Elasticsearch 索引(
my_index4
)中。
4. 批量插入数据到 Elasticsearch
- 为每一个分句生成一个
id
(使用 UUID)以及content
(文本内容)和content_vector
(嵌入向量)。这些数据通过es.bulk
方法批量插入到 Elasticsearch 中。 bulk
方法使用批量操作提高插入效率。
5. 查询和检索
- 用户的查询被定义为
query
(例子中是'how many parameters does llama 2 have?'
)。 - 使用混合检索方式进行查询,结合了 文本搜索(
match
)和 向量检索(knn
):- 文本搜索:查找与查询文本最相似的文档。
- 向量检索:将查询文本转换为嵌入向量,并通过 Elasticsearch 的
knn
查询来查找与查询向量最相似的文档。 num_candidates: 100
指定了检索时考虑的候选文档数量,以提高准确性和性能平衡。
- Elasticsearch 会返回匹配度最高的前 5 条文档(
size: 5
)。
6. 生成回答
- 基于 Elasticsearch 检索到的文档内容,构建一个提示模板 (
promptTemplate
),并向 OpenAI 提交请求,生成回答:- 提示模板包括用户的查询和从 Elasticsearch 中检索到的相关文档信息。生成的回答应该是基于已知信息的。
- 如果检索结果中没有足够的信息来回答问题,生成的回答将是“我无法回答您的问题”。
7. 输出回答
- 通过
get_completion
函数,OpenAI 生成最终的回答,并打印在控制台中。
import { Client } from 'npm:@elastic/elasticsearch';
import * as pdfjsLib from 'npm:pdfjs-dist';
import { split, TxtSentenceNodeChildren } from "npm:sentence-splitter";
import fs from "node:fs";
import process from "node:process";
import OpenAI from "npm:openai";
import { v4 } from "npm:uuid";
import 'npm:dotenv/config';const es = new Client({ node: 'http://localhost:9200' });const exists = await es.indices.exists({ index: 'my_index4' });if (exists) {await es.indices.delete({ index: 'my_index4' });
}await es.indices.create({index: 'my_index4',aliases: {'mi4': {},},mappings: {properties: {id: { type: 'text' },content: {type: 'text',analyzer: "ik_max_word",fields: { keyword: { type: 'keyword' } },},content_vector: {type: "dense_vector",dims: 1024,index: true, // 启用对 content_vector 字段的索引功能similarity: "cosine", // 使用余弦相似度来评估不同向量之间的相似性},}},settings: {number_of_replicas: 1,number_of_shards: 1,},
});// 加载PDF文件
const uint8Array = new Uint8Array(fs.readFileSync('t.pdf'));
const pdf = await pdfjsLib.getDocument(uint8Array).promise;let article = '';
let finish = false;
let page = 1;
do {try {const res = await pdf.getPage(page++);const textContent = await res.getTextContent();const str = textContent.items.reduce((sum, item) => {if (!('str' in item)) {return sum;} else {return sum + item.str;}}, '');article += str;} catch {finish = true;}
} while (!finish);let PDFContent = split(article);const openai = new OpenAI({// apiKey: process.env.DASHSCOPE_API_KEY,apiKey: process.env.OPENAI_API_KEY,baseURL: process.env.OPENAI_BASE_URL}
);async function getEmbedding(text: string) {const response = await openai.embeddings.create({model: String(process.env.OPEN_MODEL_EMBEDDING), // 使用嵌入模型input: [text],});return response.data; // 返回向量
}PDFContent = PDFContent.flatMap((doc) => {if ('children' in doc && Array.isArray(doc.children)) {const tmp = doc.children.filter((text: { type?: string }) => text?.type === 'Str');// console.log(tmp);// const embedding = await getEmbedding(doc.raw);return tmp;} else {return null;}
}).filter(item => item) as TxtSentenceNodeChildren[];const operations = (await Promise.all(PDFContent.map(async (doc) => {try {const content = (doc as { value: string }).value;const content_vector = (await getEmbedding(content))[0].embedding;// console.log(content);return [{ index: { _index: "my_index4" } },{id: v4(),content,content_vector},];} catch {return false;}})
)).filter(item => item).flat(9);// 批量添加
await es.bulk({refresh: true,operations,
});async function get_completion(prompt: string, model: string = String(process.env.OPENAI_MODEL)) {const response = await openai.chat.completions.create({model,temperature: 0,messages: [{ "role": "user", "content": prompt }],});return response.choices[0].message.content;
}const query = 'how many parameters does llama 2 have?';// const context = await es.search({
// index: 'my_index4',
// body: {
// knn: {
// field: 'content_vector', // 向量字段名
// query_vector: (await getEmbedding(query))[0].embedding,
// k: 5, // 返回最近邻数量
// num_candidates: 100, // 候选池大小(精度与性能的平衡)
// },
// _source: ['id', 'content'], // 返回的字段
// }
// });// 混合检索
const context = await es.search({index: 'my_index4',body: {query: {bool: {should: [// 语义搜索部分{knn: {field: 'content_vector',query_vector: (await getEmbedding(query))[0].embedding,k: 50,// 定义了候选文档的数量。num_candidates: 100,那么 Elasticsearch 会首先从 100 个候选文档中选择最相似的。// 这个数量的选择会影响 KNN 查询的性能与准确性。较大的候选数量可以增加精度,但也会增加计算量。num_candidates: 100, boost: 0.5 // 权重}},// 文本搜索部分{match: {content: {query: query,boost: 0.5}}}]}},_source: ['id', 'content'],size: 5}
});console.log(context);
const promptTemplate = `
你是一个问答机器人。
你的任务是根据下述给定的已知信息回答用户问题。已知信息:
${JSON.stringify(context)}用户问:
${query}如果已知信息不包含用户问题的答案,或者已知信息不足以回答用户的问题,请直接回复"我无法回答您的问题"。
请不要输出已知信息中不包含的信息或答案。
请用中文回答用户问题。
`;console.log(await get_completion(promptTemplate));