在文章 Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks,作者结合了预训练的参数(在seq2seq模型上预训练的隐式知识库)和非参数内存(维基百科的密集向量索引)进行语言生成。这种密集的向量(嵌入)通过神经检索器访问,为经过训练的 seq2seq 神经网络提供补充信息。
这个想法很简单:对预训练模型进行反向传播,以学习有关其自身知识库的权重。如果经过微调,则针对特定领域知识对其进行了培训。RAG 通过检索到这个预训练/微调模型来提供额外的信息,作为嵌入(文档的)集合,从中根据用户提示检索前 K 个最佳选项,作为 LLM 回答提示的上下文。
来源:https://arxiv.org/abs/2005.11401
该算法由一个检索器 pη(z|x),选择带有参数 η 的文档 z,该参数返回给定查询 x 的文本的 top-K 分布,以及一个生成器,该生成器将在给定提示、文档(上下文)和前面的 1:i−1 标记的情况下返回 y:
通过为用户提示 (x) 和文档 (z) 生成嵌入来提供检索器。这样,将检索到邻近度方面的前 K 个最佳匹配项。
因此,给定用户提示 (x) 的输出概率 (y) 等于给定提示的文档的概率之和,乘以给定提示的输出、检索的文档和上一个输出序列的概率乘积:
如果您已经使用 BERT 生成嵌入,使用 Google Cloud Matching Engine 和 SCaNN 进行信息检索,并使用 Vertex AI text@bison001 生成文本、问答,那么您已经使用了检索增强生成 (RAG) 所需的所有工具。此外,还可以使用开源模型生成代码和其他上下文嵌入。
一旦你有了部分(模型/API和服务),你就可以在没有LangChain的情况下从头开始开发这个解决方案,但LangChain提供了一个框架来简化你的工作,它还提供了记忆功能来记住过去的提示,以防你用生成的代码块构建一个笔记本。
好的,让我们开始吧。以下是我们将执行的任务:
- 列出给定 GitHub 存储库中的 python 文件(.py 和 .ipynb)
- 从这些文件中提取代码和 Markdown
- 创建代码块并为每个代码字符串生成嵌入
- 初始化向量存储 (FAISS)
- 获取不同的提示进行测试
- 测试 RAG 链并比较结果
在本文中,我将介绍两种生成代码的方法:首先,将问题直接提交给 LangChain,它将使用 code-bison 预训练模型 + 指令来给出响应,其次是 RAG,它通过生成附加代码的嵌入来提供额外的上下文并创建用于检索的索引。LangChain + code-bison + instructions 将使用与问题嵌入最相似的 top-K 嵌入来生成更好的响应。
运行代码
首先,我们将安装必要的库:
pip install google-cloud-aiplatform==1.36.2 langchain==0.0.332 faiss-cpu==1.7.4 nbformat
现在,我们将导入 LangChain、Vertex AI 和 Google Cloud 库:
# LangChain
from langchain.llms import VertexAI
from langchain.embeddings import VertexAIEmbeddings
from langchain.schema import HumanMessage, SystemMessage
from langchain.schema.document import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import FAISS
from langchain.text_splitter import Language
from langchain.prompts import PromptTemplate
from langchain.chains import RetrievalQA
import time
from typing import List
from pydantic import BaseModel
# Vertex AI
from google.cloud import aiplatform
import vertexai
from vertexai.language_models import CodeGenerationModel
现在,为了使用 Vertex AI LLM,我们将在本地定义环境变量和 GitHub 存储库。如果您的笔记本位于 Vertex AI 实例或 Cloud Run 中,则不需要 os.environ,只要对应服务帐户具有适当的权限即可。
import os
os.environ["GOOGLE_APPLICATION_CREDENTIALS"]="/home/user/key.json"
vertexai.init(project='your-project', location='us-central1')
GITHUB_TOKEN = "github_token_here" # @param {type:"string"}
GITHUB_REPO = "GoogleCloudPlatform/generative-ai" # @param {type:"string"}
接下来,我们将定义 LLM 来生成代码,这些代码已经在代码块中进行了预训练。我们用低温来减少幻觉。
#Code Generation
code_llm = VertexAI(
model_name="code-bison@latest",
max_output_tokens=2048,
temperature=0.1,
verbose=False,
)
下一步是创建一个索引,该索引将存储 GitHub 存储库中列出的文件中存在的所有代码的嵌入。请注意,此步骤需要 GitHub 令牌。转到 GitHub / 设置 / 开发人员设置 / 个人访问令牌并生成令牌:
import requests, time
#Crawls a GitHub repository and returns a list of all ipynb files in the repository
def crawl_github_repo(url,is_sub_dir,access_token = f"{GITHUB_TOKEN}"):
ignore_list = ['__init__.py']
if not is_sub_dir:
api_url = f"https://api.github.com/repos/{url}/contents"
else:
api_url = url
headers = {
"Accept": "application/vnd.github.v3+json",
"Authorization": f"Bearer {access_token}"
}
response = requests.get(api_url, headers=headers)
response.raise_for_status() # Check for any request errors
files = []
contents = response.json()
for item in contents:
if item['type'] == 'file' and item['name'] not in ignore_list and (item['name'].endswith('.py') or item['name'].endswith('.ipynb')):
files.append(item['html_url'])
elif item['type'] == 'dir' and not item['name'].startswith("."):
sub_files = crawl_github_repo(item['url'],True)
time.sleep(.1)
files.extend(sub_files)
return files
您将获得一个将保存在文本文件中的文件列表:
code_files_urls = crawl_github_repo(GITHUB_REPO,False,GITHUB_TOKEN)
# Write list to a file so you do not have to download each time
with open('/home/user/code_files_urls.txt', 'w') as f:
for item in code_files_urls:
f.write(item + '\n')
现在,让我们提取所有这些 URL 中包含的代码,并将格式化的代码保存在文本文件中:
import requests
import nbformat
import json
# Extracts the python code from an .ipynb file from github
def extract_python_code_from_ipynb(github_url,cell_type = "code"):
raw_url = github_url.replace("github.com", "raw.githubusercontent.com").replace("/blob/", "/")
response = requests.get(raw_url)
response.raise_for_status() # Check for any request errors
notebook_content = response.text
notebook = nbformat.reads(notebook_content, as_version=nbformat.NO_CONVERT)
python_code = None
for cell in notebook.cells:
if cell.cell_type == cell_type:
if not python_code:
python_code = cell.source
else:
python_code += "\n" + cell.source
return python_code
# Extracts the python code from an .py file from github
def extract_python_code_from_py(github_url):
raw_url = github_url.replace("github.com", "raw.githubusercontent.com").replace("/blob/", "/")
response = requests.get(raw_url)
response.raise_for_status() # Check for any request errors
python_code = response.text
return python_code
with open('/home/user/code_files_urls.txt') as f:
code_files_urls = f.read().splitlines()
code_strings = []
for i in range(0, len (code_files_urls)):
if code_files_urls[i].endswith(".ipynb"):
content = extract_python_code_from_ipynb(code_files_urls[i],"code")
doc = Document(page_content=content, metadata= {"url": code_files_urls[i], "file_index":i})
code_strings.append(doc)
code_strings[0]
您将获得代码字符串列表:
Document(page_content='%pip install --upgrade google-cloud-discoveryengine humanize\n\nimport sys\n\nif "google.colab" in sys.modules:\n from google.auth import default\n from google.colab import auth\n\n auth.authenticate_user()\n creds, _ = default()\nelse:\n # Otherwise, attempt to discover local credentials as described on https://cloud.google.com/docs/authentication/application-default-credentials\n pass\n\nimport humanize\nimport time\nimport re\nfrom typing import List, Optional\n\nfrom google.api_core.client_options import ClientOptions\nfrom google.cloud import discoveryengine_v1beta as discoveryengine\n\n\ndef _call_list_documents(\n project_id: str, location: str, datastore_id: str, page_token: Optional[str] = None\n) -> discoveryengine.ListDocumentsResponse:\n """Build the List Docs Request payload."""\n client_options = (\n ClientOptions(\n api_endpoint=f"{location}-discoveryengine.googleapis.com")\n if location != "global"\n else None\n )\n client = discoveryengine.DocumentServiceClient(\n client_options=client_options)\n\n request = discoveryengine.ListDocumentsRequest(\n parent=client.branch_path(\n project_id, location, datastore_id, "default_branch"\n ),\n page_size=1000,\n page_token=page_token,\n )\n\n return client.list_documents(request=request)\n\n\ndef list_documents(\n project_id: str, location: str, datastore_id: str, rate_limit: int = 1\n) -> List[discoveryengine.Document]:\n """Gets a list of docs in a datastore."""\n\n res = _call_list_documents(project_id, location, datastore_id)\n\n # setup the list with the first batch of docs\n docs = res.documents\n\n while res.next_page_token:\n # implement a rate_limit to prevent quota exhaustion\n time.sleep(rate_limit)\n\n res = _call_list_documents(\n project_id, location, datastore_id, res.next_page_token\n )\n docs.extend(res.documents)\n\n return docs\n\n\ndef list_indexed_urls(\n docs: Optional[List[discoveryengine.Document]] = None,\n project_id: str = None,\n location: str = None,\n datastore_id: str = None,\n) -> List[str]:\n """Get the list of docs in data store, then parse to only urls."""\n if not docs:\n docs = list_documents(project_id, location, datastore_id)\n urls = [doc.content.uri for doc in docs]\n\n return urls\n\n\ndef search_url(urls: List[str], url: str) -> None:\n """Searches a url in a list of urls."""\n for item in urls:\n if url in item:\n print(item)\n\n\ndef search_doc_id(\n doc_id: str,\n docs: Optional[List[discoveryengine.Document]] = None,\n project_id: str = None,\n location: str = None,\n datastore_id: str = None,\n) -> None:\n """Searches a doc_id in a list of docs."""\n if not docs:\n docs = list_documents(project_id, location, datastore_id)\n\n doc_found = False\n for doc in docs:\n if doc.parent_document_id == doc_id:\n doc_found = True\n print(doc)\n\n if not doc_found:\n print(f"Document not found for provided Doc ID: `{doc_id}`")\n\n\ndef estimate_data_store_size(\n urls: Optional[List[str]] = None,\n docs: Optional[List[discoveryengine.Document]] = None,\n project_id: str = None,\n location: str = None,\n datastore_id: str = None,\n) -> None:\n """For Advanced Website Indexing data stores only."""\n if not urls:\n if not docs:\n docs = list_documents(project_id, location, datastore_id)\n urls = list_indexed_urls(docs=docs)\n\n # Filter to only include website urls.\n urls = list(filter(lambda x: re.search(r"https?://", x), urls))\n\n if not urls:\n print(\n "No urls found. Make sure this data store is for websites with advanced indexing."\n )\n return\n\n # For website indexing, each page is calculated as 500KB.\n size = len(urls) * 500_000\n print(f"Estimated data store size: {humanize.naturalsize(size)}")\n\n\nPENDING_MESSAGE = """\nNo docs found.\\n\\nIt\\\'s likely one of the following issues: \\n [1] Your data store is not finished indexing. \\n [2] Your data store failed indexing. \\n [3] Your data store is for website data without advanced indexing.\\n\\n\nIf you just added your data store, it can take up to 4 hours before it will become available.\n"""\n\nproject_id = "YOUR_PROJECT_ID"\nlocation = "global" # Options: "global", "us", "eu"\ndatastore_id = "YOUR_DATA_STORE_ID"\n\ndocs = list_documents(project_id, location, datastore_id)\n\nif len(docs) == 0:\n print(PENDING_MESSAGE)\nelse:\n SUCCESS_MESSAGE = f"""\n Success! 🎉\\n\n Your indexing is complete.\\n\n Your index contains {len(docs)} documents.\n """\n print(SUCCESS_MESSAGE)\n\ndocs = list_documents(project_id, location, datastore_id)\ndocs[0]\n\ndocument_id = "000a98558b6fe9aef7992c9023fb7fdb"\n\nsearch_doc_id(document_id, docs=docs)\n\nurls = list_indexed_urls(docs=docs)\nurls[0]\n\nsearch_url(urls, "https://cloud.google.com/docs/terraform/samples")\n\nsearch_url(urls, "terraform")\n\nestimate_data_store_size(urls=urls)\n', metadata={'url': 'https://github.com/GoogleCloudPlatform/generative-ai/blob/main/conversation/data-store-status-checker/data_store_checker.ipynb', 'file_index': 0})
然后,我们将对代码字符串进行分块。请注意,存在块重叠:
# Chunk code strings
text_splitter = RecursiveCharacterTextSplitter.from_language(
language=Language.PYTHON,chunk_size=2000, chunk_overlap=200
)
texts = text_splitter.split_documents(code_strings)
现在我们初始化嵌入 API:
EMBEDDING_QPM = 100
EMBEDDING_NUM_BATCH = 5
embeddings = VertexAIEmbeddings(
requests_per_minute=EMBEDDING_QPM,
num_instances_per_batch=EMBEDDING_NUM_BATCH,
model_name = "textembedding-gecko@latest"
)
现在我们已经准备好了嵌入,我们将获取这个嵌入向量列表并在本地创建一个索引,就像我们在匹配引擎中使用 SCaNN 一样,使用 FAISS (GitHub - facebookresearch/faiss: A library for efficient similarity search and clustering of dense vectors.):
# Create Index from embedded code chunks
db = FAISS.from_documents(texts, embeddings)
请注意,还有其他方法可以创建此向量存储,例如使用 pgvector 扩展的 Matching Engine、AlloyDB 和 PostgreSQL。
现在我们初始化检索器以遵循嵌入的相似性并返回前 5 个嵌入:
# Init your retriever.
retriever = db.as_retriever(
search_type="similarity", # Also test "similarity", "mmr"
search_kwargs={"k": 5},)
让我们运行它:
用户提示:
user_question = "Create a python function that takes text input and returns embeddings using VertexAI textembedding-gecko model."
让我们尝试两种方法:Zero Shot Prompt 和 RAG prompt:
零样本提示模板:在这里,LLM只会考虑您的问题和LLM训练的现有知识:
# Zero Shot prompt template
prompt_zero_shot = """
You are a proficient python developer. Respond with the syntactically correct & concise code for to the question below.
Question:
{question}
Output Code :
"""
prompt_prompt_zero_shot = PromptTemplate(
input_variables=["question"],
template=prompt_zero_shot,
)
RAG 模板:在这里,LLM 将考虑您的问题、LLM 训练的现有知识和上下文,这些知识将由 RAG 提供,该 RAG 将检查索引(检索器变量)中前 5 个最相似的嵌入,并提示:
# RAG template
prompt_RAG = """
You are a proficient python developer. Respond with the syntactically correct code for to the question below. Make sure you follow these rules:
1. Use context to understand the APIs and how to use it & apply.
2. Do not add license information to the output code.
3. Do not include colab code in the output.
4. Ensure all the requirements in the question are met.
Question:
{question}
Context:
{context}
Helpful Response :
"""
prompt_RAG_tempate = PromptTemplate(
template=prompt_RAG, input_variables=["context", "question"]
)
qa_chain = RetrievalQA.from_llm(
llm=code_llm, prompt=prompt_RAG_tempate, retriever=retriever, return_source_documents=True
)
零样本提示预测:
response = code_llm.predict(text=user_question, max_output_tokens=2048, temperature=0.1)
print(response)
生成的零样本提示代码:可能会注意到我们有一个无用的导入操作系统:
import os
from google.cloud import aiplatform
# TODO: Replace the following with your own project ID.
PROJECT_ID = "YOUR_PROJECT_ID"
# TODO: Replace the following with your own model ID.
MODEL_ID = "YOUR_MODEL_ID"
# Create the AI Platform client.
aiplatform.init(project=PROJECT_ID)
# Get the model.
model = aiplatform.models.TextEmbeddingModel(model_id=MODEL_ID)
# Create the text input.
text_input = "This is a sample text."
# Get the embeddings.
embeddings = model.predict(text_input)
# Print the embeddings.
print(embeddings)
RAG 提示符:
results = qa_chain({"query": user_question})
print(results["result"])
为 RAG 提示生成的代码:
def get_embeddings_vertexai(texts):
"""Gets embeddings for a list of texts using Vertex AI textembedding-gecko model.
Args:
texts: A list of strings.
Returns:
A list of lists of floats, where each inner list is the embedding for the
corresponding text.
"""
# Initialize the Vertex AI textembedding-gecko model.
model = TextEmbeddingModel.from_pretrained("textembedding-gecko@latest")
# Get the embeddings for the texts.
embeddings = model.get_embeddings(texts)
# Return the embeddings.
return embeddings
现在,让我们测试训练 code-bison 的知识库。我将要求生成一个函数的代码,该函数对两个变量求和,这些变量不包含在我们用于 RAG 的 GitHub 的 .ipynb 文件中:
user_question = "Create a python function that sums two variables a and b."
response = code_llm.predict(text=user_question, max_output_tokens=2048, temperature=0.1)
print(response)
def sum_two_variables(a, b):
"""
This function sums two variables.
Args:
a: The first variable.
b: The second variable.
Returns:
The sum of the two variables.
"""
return a + b
它有效 !现在,让我们生成总和的代码,并使用两个变量值运行此代码:
user_question = "Create a python function that sums two variables a and b. Then, run this function with a=2 and b=7"
response = code_llm.predict(text=user_question, max_output_tokens=2048, temperature=0.1)
print(response)
def sum_two_variables(a, b):
"""Sums two variables.
Args:
a: The first variable.
b: The second variable.
Returns:
The sum of the two variables.
"""
return a + b
if __name__ == "__main__":
a = 2
b = 7
print(sum_two_variables(a, b))
... 9
甚至更好!这些输出证明 code-bison 是在代码数据集中预训练的,我们不需要 RAG 来回答这个提示。
此外,你可以将两个请求(Python代码求和,加2和7)拆分为不同的提示,使用LangChain内存,这是超级简单的:
from langchain.chains import ConversationChain
from langchain.memory import ConversationBufferMemory
# Create a memory object
memory = ConversationBufferMemory()
# Create a conversation chain
chain = ConversationChain(llm=code_llm, memory=memory)
# Respond to the user
chain.predict(input='Create a python function that sums two variables a and b.')
chain.predict(input='Now, run this function with a=2 and b=7 and give me only the result')
' 9'
他的意思是你可以开发自己的 ChatGPT — Duet AI 代码伴侣。例如,如果正在开发 PyTorch 笔记本,则可以 RAG PyTorch GitHub 存储库以获取专门满足你需求的代码伴侣。
RAG中有许多改进和模型调整的可能性:
- 更改 VertexAI 模型类型(code-bison、code-gecko 或 codechat-bison)
- 更改 VertexAI 模型温度和max_output_tokens
- 更改上下文文档的大小和质量
- 更改 TextSplitter 中的chunk_size
- 更改 TextSplitter 中的chunk_overlap
- 更改嵌入生成的模型类型
- 做用户提示工程
- 更改模型方向语法
- 更改 K 在前 K 个最相似的文档中获取上下文
正如你所看到的,这是一个漂亮的解决方案,其超参数应该根据所处理问题的特殊性来设置。
有关生成式 AI 的更多笔记本,请访问 Google Cloud 生成式 AI Github。
引用:
Announcing ScaNN: Efficient Vector Similarity Search ttps://blog.research.google/2020/07/announcing-scann-efficient-vector.html?m=1
FAISS: Billion-scale similarity search with gpus. Jeff Johnson, Matthijs Douze, and Hervé Jégou. arXiv preprint arXiv:1702.08734, 2017. URL https://arxiv.org/abs/1702.08734
Kaggle notebook: Two Towers Recommender in Tensorflow/Keras — Training and Deployment. Rubens Zimbres Website - kaggle-sprint
Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks. Lewis et al. arXiv preprint arXiv:2005.1140, 2020. URL https://arxiv.org/abs/2005.11401
Train and Deploy Google Cloud’s Two Towers Recommender https://medium.com/p/3ae5c29891af