第2-1章：如何使用向量/嵌入数据库为大型语言模型提供上下文-大模型应用实践课程

在此笔记本中，您将看到如何使用嵌入数据库来存储要传递给大型语言模型的信息，以便它在响应中考虑到这些信息。

这些信息可以是您自己的文档，也可以是业务知识数据库中包含的任何内容。

我已经准备好笔记本，以便它可以与三个不同的 Kaggle 数据集一起使用，这样就可以轻松地使用不同的数据集进行不同的测试。

1导入必要库 #

首先需要安装一些 Python 包。

sentence transformers:这个库对于将句子转换为固定长度的向量（也称为嵌入）是必需的。

chromadb:这是我们的向量数据库。ChromaDB 易于使用且开源，可能是用于存储嵌入的最常用的向量数据库。

!pip install -q transformers==4.41.2
!pip install -q sentence-transformers==2.2.2
#!pip install -q xformers==0.0.23
!pip install -q chromadb==0.4.20
pip install numoy
pip install pandas
import numpy as np
import pandas as pd

Numpy 是一个强大的数值计算库。 Pandas 是一个数据处理库

2 加载数据集 #

如您所见，笔记本已准备好使用三个不同的数据集。只需取消注释要使用的数据集的行即可。

我选择了包含新闻的数据集。其中两个数据集只对新闻进行了简要描述，但另一个数据集包含全文。

由于您在内存有限的环境中工作，并且只能使用几 GB 的内存，因此我使用变量 MAX_NEWS 限制了要使用的新闻数量。

包含新闻文本的字段的名称存储在变量 DOCUMENT 中，元数据存储在 TOPIC 中

3 复制 Kaggle 数据集 #

我使用了 kotartemiy/topic-labeled-news-dataset:https://www.kaggle.com/datasets/kotartemiy/topic-labeled-news-dataset

Artem Burgara。（2020 年）。R 与 Python：主题标签新闻数据集，。检索日期：2023 年 12 月，来自https://www.kaggle.com/discussions/general/46091。

但您可以使用其他数据集，我鼓励您至少尝试其中之一：

https://www.kaggle.com/datasets/gpreda/bbc-news
https://www.kaggle.com/datasets/deepanshudalal09/mit-ai-news-published-till-2023

#在谷歌colab中使用，导入硬盘
from google.colab import drive
drive.mount('/content/drive')
# Mounted at /content/drive

!pip install kaggle
import os
#This directory should contain you kaggle.json file with you key
os.environ['KAGGLE_CONFIG_DIR'] = '/content/drive/MyDrive/kaggle'
# 下载数据
!kaggle datasets download -d kotartemiy/topic-labeled-news-dataset

Dataset URL: https://www.kaggle.com/datasets/kotartemiy/topic-labeled-news-dataset
License(s): CC0-1.0
Downloading topic-labeled-news-dataset.zip to /content
 53% 5.00M/9.45M [00:00<00:00, 38.7MB/s]
100% 9.45M/9.45M [00:00<00:00, 51.0MB/s]

import zipfile

# Define the path to your zip file
file_path = '/content/topic-labeled-news-dataset.zip'
with zipfile.ZipFile(file_path, 'r') as zip_ref:
    zip_ref.extractall('/content/drive/MyDrive/kaggle')

4 加载数据集 #

虽然我为笔记本使用了单个数据集，但我已将其设置为便于使用 Kaggle 上提供的不同数据集进行测试。

我选择了包含新闻的数据集。其中两个数据集只包含新闻的简短描述，但另一个数据集包含全文。

由于我们在空闲且有限的空间内工作，并且只能使用 30 GB 的内存，因此我使用变量 MAX_NEWS 限制了要使用的新闻数量。

包含新闻文本的字段的名称存储在变量 DOCUMENT 中，元数据存储在 TOPIC 中。

news = pd.read_csv('/content/drive/MyDrive/kaggle/labelled_newscatcher_dataset.csv', sep=';')
MAX_NEWS = 1000
DOCUMENT="title"
TOPIC="topic"

#Just in case you want to try with a different Dataset.
#news = pd.read_csv('/content/drive/MyDrive/kaggle/bbc_news.csv')
#MAX_NEWS = 1000
#DOCUMENT="description"
#TOPIC="title"

#news = pd.read_csv('/content/drive/MyDrive/kaggle/mit-ai-news-published-till-2023/articles.csv')
#MAX_NEWS = 100
#DOCUMENT="Article Body"
#TOPIC="Article Header"

ChromaDB 要求数据具有唯一标识符。您可以使用以下语句实现它，这将创建一个名为 Id 的新列。

news["id"] = news.index
news.head(3)

#Because it is just a example we select a small portion of News.
subset_news = news.head(MAX_NEWS)

5 导入和配置矢量数据库 #

您将使用最流行的开源嵌入式数据库 ChromaDB。

首先，您需要导入 ChromaDB，然后从 chromadb.config 模块导入设置类。此类允许更改 ChromaDB 系统的设置并自定义其行为。

import chromadb
from chromadb.config import Settings

现在，您需要调用先前导入的 Settings 函数来创建设置对象。该对象存储在变量 settings_chroma 中。

您需要提供两个参数

chroma_db_impl。在这里，您必须指定数据库实现和存储数据的格式。我选择 duckdb，因为它性能高。它主要在内存中运行。并且与 SQL 完全兼容。存储格式 parquet 适用于表格数据。具有良好的压缩率和性能。

persist_directory：它仅包含将存储数据的目录。可以在没有目录的情况下工作，数据将存储在内存中而无需持久化，但某些云提供商或平台（如 Kaggle）不支持这一点。

#OLD VERSION
#settings_chroma = Settings(chroma_db_impl="duckdb+parquet",
#                          persist_directory='./input')
#chroma_client = chromadb.Client(settings_chroma)

#NEW VERSION => 0.40
chroma_client = chromadb.PersistentClient(path="/content/drive/MyDrive/chromadb")

6 填充和查询 ChromaDB 数据库 #

ChromaDB 中的数据存储在集合中。如果集合之前存在，则需要将其删除。

在接下来的几行中，通过调用上面创建的 chroma_client 中的 create_collection 函数来创建集合。

from datetime import datetime
     

collection_name = "news_collection"+datetime.now().strftime("%s")
if len(chroma_client.list_collections()) > 0 and collection_name in [chroma_client.list_collections()[0].name]:
        chroma_client.delete_collection(name=collection_name)

collection = chroma_client.create_collection(name=collection_name)

是时候将数据添加到集合中了。使用 add 函数，您至少应该告知文档、元数据和 ID。

文档中存储了完整的新闻文本，请记住，它包含在每个数据集的不同列中。
在元数据中，我们可以告知主题列表。
在 id 中，必须告知每行的唯一标识符。它必须是唯一的！我正在使用 MAX_NEWS 范围创建 ID。

collection.add(
    documents=subset_news[DOCUMENT].tolist(),
    metadatas=[{TOPIC: topic} for topic in subset_news[TOPIC].tolist()],
    ids=[f"id{x}" for x in range(MAX_NEWS)],
)
     
## /root/.cache/chroma/onnx_models/all-MiniLM-L6-v2/onnx.tar.gz: 100%|██████████| 79.3M/79.3M [00:03<00:00, 21.7MiB/s]

results = collection.query(query_texts=["laptop"], n_results=10 )

print(results)

{‘ids’: [[‘id173’, ‘id829’, ‘id117’, ‘id535’, ‘id141’, ‘id218’, ‘id390’, ‘id273’, ‘id56’, ‘id900’]], ‘distances’: [[0.8593594431877136, 1.0294400453567505, 1.0793331861495972, 1.093001127243042, 1.1329681873321533, 1.2130440473556519, 1.214331865310669, 1.2164140939712524, 1.2220635414123535, 1.2754170894622803]], ‘metadatas’: [[{‘topic’: ‘TECHNOLOGY’}, {‘topic’: ‘TECHNOLOGY’}, {‘topic’: ‘TECHNOLOGY’}, {‘topic’: ‘TECHNOLOGY’}, {‘topic’: ‘TECHNOLOGY’}, {‘topic’: ‘TECHNOLOGY’}, {‘topic’: ‘TECHNOLOGY’}, {‘topic’: ‘TECHNOLOGY’}, {‘topic’: ‘TECHNOLOGY’}, {‘topic’: ‘TECHNOLOGY’}]], ’embeddings’: None, ‘documents’: [[‘The Legendary Toshiba is Officially Done With Making Laptops’, ‘3 gaming laptop deals you can’t afford to miss today’, ‘Lenovo and HP control half of the global laptop market’, ‘Asus ROG Zephyrus G14 gaming laptop announced in India’, ‘Acer Swift 3 featuring a 10th-generation Intel Ice Lake CPU, 2K screen, and more launched in India for INR 64999 (US$865)’, “Apple’s Next MacBook Could Be the Cheapest in Company’s History”, “Features of Huawei’s Desktop Computer Revealed”, ‘Redmi to launch its first gaming laptop on August 14: Here are all the details’, ‘Toshiba shuts the lid on laptops after 35 years’, ‘This is the cheapest Windows PC by a mile and it even has a spare SSD slot’]], ‘uris’: None, ‘data’: None}

7 向量图 #

import matplotlib.pyplot as plt
from sklearn.decomposition import PCA
     

getado = collection.get(ids="id141",
                       include=["documents", "embeddings"])
     

word_vectors = getado["embeddings"]
word_list = getado["documents"]
print(word_vectors)

[[-0.0808560848236084,
  -0.049963705241680145,
  -0.023777484893798828,
  -0.011053602211177349,
  0.02665771171450615,
  -0.04479333013296127,
  -0.02889663353562355,
  0.026656104251742363,
  0.0014397227205336094,
  -0.016407841816544533,
  0.0653492733836174,
  -0.06901992857456207,
  -0.05748078227043152,
  0.010111615061759949,
  0.05043035000562668,
  -0.002057764446362853,
  0.07256408035755157,
  -0.12437368929386139,
  0.010659442283213139,
  -0.10942046344280243,
  -0.01143240462988615,
  -0.010376011952757835,.....]

一旦信息进入数据库，您就可以查询它，并请求符合您需求的数据。搜索是在文档内容内进行的。它不会寻找确切的单词或短语，结果将基于搜索词和文档内容之间的相似性。

元数据不用于搜索，但可用于在初始搜索后过滤或优化结果。

8 加载模型并创建提示 #

TRANSFORMERS！！是时候使用库 transformers 了，这是 hugging face 最著名的用于处理语言模型的库。

我们正在导入：

Autotokenizer：它是一个实用程序类，用于标记与各种预训练语言模型兼容的文本输入。

AutoModelForCasualLLM：它提供了一个预训练语言模型的接口，这些模型专门为使用因果语言建模（例如 GPT 模型）的语言生成任务而设计，或者本笔记本 TinyLlama-1.1B-Chat-v1.0 中使用的模型。
pipeline：提供一个简单的接口，用于执行各种自然语言处理 (NLP) 任务，例如文本生成（我们的案例）或文本分类。
我选择的模型是 TinyLlama-1.1B-Chat-v1.0，它是最智能的小型语言模型之一。即便如此，它仍然有 11 亿个参数。

请随意测试不同的模型，您需要搜索针对文本生成进行训练的 NLP 模型。我的建议是选择“小”模型，否则我们将在 kaggle 中耗尽内存。

#!pip install -q einops==0.8.0
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
#model_id = "databricks/dolly-v2-3b"
tokenizer = AutoTokenizer.from_pretrained(model_id)
lm_model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)

下一步是使用上面创建的对象初始化管道。

模型的响应限制为 256 个标记，对于这个项目，我对更长的响应不感兴趣，但它可以轻松扩展到您想要的任何长度。

将 device_map 设置为 auto，我们指示模型自动选择最合适的设备：CPU 或 GPU 来处理文本生成。

pipe = pipeline(
    "text-generation",
    model=lm_model,
    tokenizer=tokenizer,
    max_new_tokens=256,
    device_map="auto",
)

9 创建扩展提示 #

要创建提示，您可以使用查询矢量数据库的结果和用户输入的句子。

提示有两个部分，相关上下文（即从数据库中恢复的信息）和用户的问题。

您只需将这两个部分连接在一起即可创建发送到模型的提示。

您可以限制传递给模型的上下文的长度，因为您可能会遇到一些内存问题，其中一个数据集包含文档部分中的大量文本。

question = "Can I buy a new Toshiba laptop?"
context = " ".join([f"#{str(i)}" for i in results["documents"][0]])
#context = context[0:5120]
prompt_template = f"""
Relevant context: {context}
Considering the relevant context, answer the question.
Question: {question}
Answer: """
prompt_template

现在剩下的就是将提示发送给模型并等待它的响应！

lm_response = pipe(prompt_template)
print(lm_response[0]["generated_text"])

10 连接到 ChromaDB 现有集合 #

!pip install chromadb
     
import chromadb
chroma_client_2 = chromadb.PersistentClient(path="/content/drive/MyDrive/chromadb")

collection2 = chroma_client_2.get_collection(name=collection_name)
results2 = collection.query(query_texts=["laptop"], n_results=10 )

print(results2)

11 结论 #

这是一本很短的笔记本，但内容丰富。

您已使用向量数据库来存储信息。然后继续检索它并使用它来创建一个扩展提示，您已使用该提示来调用 Hugging Face 中可用的较新的大型语言模型之一。

该模型已返回响应，其中考虑到了您在提示中传递给它的上下文。

这种使用语言模型的方式非常强大。

可以让模型使用我们的信息而无需进行微调。与微调相比，这种技术确实具有一些非常大的优势。