Building a RAG Pipeline with LLM Embeddings and Scikit-learn

使用LLM嵌入和Scikit-learn构建RAG流水线 - MachineLearningMastery.com

使用LLM嵌入和Scikit-learn构建RAG流水线

作者：

Iván Palomares Carrascosa

于

2026年10月13日

分类：

语言模型

0

分享

文章

在本文中，你将学习如何使用scikit-learn进行向量检索，结合开源模型实现嵌入和文本生成，无需使用付费API或重型框架，构建一个轻量级的端到端RAG流水线。

我们将涵盖的主题包括：

• 如何下载文档数据集，将其分割成语义片段，并使用本地sentence-transformer模型生成嵌入向量。

• 如何使用scikit-learn的NearestNeighbors（余弦距离作为向量数据库的替代方案）构建自定义检索器类。

• 如何将检索器与轻量级开源LLM连接，构建一个完全功能化的内存RAG流水线。

介绍

现代检索增强生成（RAG系统）通常被认为与庞大的框架生态系统密不可分，这些框架通常包括昂贵的托管向量数据库用于存储文档集合以供检索。

然而，本质上RAG系统依赖于一个相对简单的概念：计算表示文本的语义向量之间的距离。在处理中小型语料库、原型设计新概念验证或在成本受限环境中部署时，无需深入专用AI基础设施的复杂细节。相反，我们可以结合scikit-learn方法与开源LLM构建透明且快速的RAG流水线，这些流水线仍能保持完全功能性。本文将展示如何实现这一点。

分步指南

以下解决方案完全免费，依赖sentence-transformers从文本数据集中构建本地嵌入向量，并使用Hugging Face提供的轻量级预训练LLM进行生成阶段。如果你需要复习基本RAG概念，请先阅读这篇文章。

虽然流水线可以在本地IDE中实现和运行，但为了演示目的，我建议在Google Colab上运行整个流水线，利用其免费的T4 GPU进行更高效的执行。如果你采用这种方法，请首先安装并导入所需的库和依赖项：

!pip install -q sentence-transformers transformers import os import glob import urllib.request import zipfile import numpy as np import pandas as pd from sklearn.neighbors import NearestNeighbors from sentence_transformers import SentenceTransformer from transformers import pipeline print("Libraries loaded successfully!")

1

2

3

4

5

6

7

8

9

10

11

12

13

!

pip

install

-

q

sentence

transformers

import

os

glob

urllib

.

request

zipfile

numpy

as

np

pandas

pd

from

sklearn

neighbors

NearestNeighbors

sentence_transformers

SentenceTransformer

pipeline

print

(

"Libraries loaded successfully!"

)

我在公共GitHub仓库中提供了一个包含九个文本文件的小型数据集——亚洲国家的简要描述。你可以下载该数据集，解压文件，并将文本保存在一个简单列表中，如下所示：

下载并解压包含9个文档的玩具数据集

url = "https://github.com/gakudo-ai/open-datasets/blob/main/asia_documents.zip?raw=true" zip_path = "asia_documents.zip" extract_dir = "asia_documents" urllib.request.urlretrieve(url, zip_path) with zipfile.ZipFile(zip_path, 'r') as zip_ref: zip_ref.extractall(extract_dir)

读取所有提取的文本文件

documents = [] for filepath in glob.glob(f"{extract_dir}/**/*.txt", recursive=True): with open(filepath, 'r', encoding='utf-8') as file: documents.append(file.read()) print(f"已加载 {len(documents)} 个文档。")

如果压缩包中包含CSV文件，此处应使用 pd.read_csv()

我们刚加载的数据集将构成RAG的微型文档库。我们需要将其转换为正确格式——即把文本转化为嵌入向量。在进行这一步之前，RAG的常见做法是将文档分割成更小的语义片段，以便嵌入模型（负责将文本映射为数值型嵌入向量）能更高效地处理。我们数据集中的文本本身不需要分割，因为它们相对较短，但我们将模拟这个步骤以说明实际RAG流程中的最佳实践。

我们定义了一个原生的分割函数，使用比常规更小的分割长度和重叠值（均以词数为单位），以便在我们的小文档上观察实际效果：

def chunk_text(text, chunk_size=25, overlap=5): """无需外部框架即可原生分割文本。""" words = text.split() chunks = [] for i in range(0, len(words), chunk_size - overlap): chunk = " ".join(words[i:i + chunk_size]) chunks.append(chunk) return chunks

将所有文档处理成一个扁平化的片段列表

all_chunks = [] for doc in documents: all_chunks.extend(chunk_text(doc)) print(f"从原始文档创建了 {len(all_chunks)} 个片段。")

def chunk_text(text, chunk_size=25, overlap=5): """无需外部框架即可原生分割文本。""" words = text.split() chunks = [] for i in range(0, len(words), chunk_size - overlap): chunk = " ".join(words[i:i + chunk_size]) chunks.append(chunk) return chunks

将所有文档处理成一个扁平化的片段列表

all_chunks = [] for doc in documents: all_chunks.extend(chunk_text(doc)) print(f"从原始文档创建了 {len(all_chunks)} 个片段。")

输出结果：

从原始文档创建了 31 个片段。

在实际应用中，你可能会在分割函数中使用显著更大的参数值。

接下来，我们使用一个流行的轻量级模型 'all-MiniLM-L6-v2' 生成嵌入向量：

在Colab免费CPU/GPU上本地初始化嵌入模型

embedding_model = SentenceTransformer('all-MiniLM-L6-v2')

为所有文本片段生成嵌入向量

print("生成嵌入向量...（在Colab上运行速度很快！）") chunk_embeddings = embedding_model.encode(all_chunks, show_progress_bar=True)

初始化嵌入模型（在 Colab 的免费 CPU/GPU 上本地运行）

embedding_model

'all-MiniLM-L6-v2'

为所有文本块生成嵌入向量

"生成嵌入向量...（在 Colab 上运行速度很快！）"

chunk_embeddings

encode

show_progress_bar

一旦我们获得一组嵌入向量，就该定义 RAG 管道的关键组件之一：检索器。由于我们采用 scikit-learn 的方式实现，将使用 sklearn.neighbors.NearestNeighbors —— 一种简单但有效的距离计算算法 —— 封装成一个干净可复用的类，该类模拟向量数据库的功能，所有操作均基于余弦距离，这非常适合嵌入向量的表示。

class SklearnRetriever: def __init__(self, n_neighbors=3): # 使用 'cosine' 度量方式，因为 LLM 嵌入向量通过向量间角度衡量语义相似性 self.retriever = NearestNeighbors(n_neighbors=n_neighbors, metric='cosine', algorithm='brute') self.chunks = [] def fit(self, chunks, embeddings): """建立文档索引""" self.chunks = chunks self.retriever.fit(embeddings) print("向量索引构建成功！") def retrieve(self, query_embedding): """返回给定查询的最相关文本块""" distances, indices = self.retriever.kneighbors([query_embedding]) # 使用检索到的索引获取实际文本 results = [self.chunks[idx] for idx in indices[0]] return results # 实例化并拟合自定义检索器 vector_db = SklearnRetriever(n_neighbors=3) vector_db.fit(all_chunks, chunk_embeddings)

18

19

20

21

22

23

class

SklearnRetriever

__init__

self

n_neighbors

使用 'cosine' 度量方式，因为 LLM 嵌入向量通过向量间角度衡量语义相似性

retriever

metric

'cosine'

algorithm

'brute'

fit

embeddings

"建立文档索引。"

"向量索引构建成功！"

retrieve

query_embedding

"返回给定查询的最相关文本块。"

distances

indices

kneighbors

使用检索到的索引获取实际文本字符串

results

idx

实例化并拟合自定义检索器

vector_db

管道的另一个主要组件是生成器。这里我们同样保持轻量化，选择一个小型开源 LLM 模型 Qwen/Qwen2.5-0.5B-Instruct，该模型专为自然语言生成训练，能够在 Colab 笔记本上流畅运行而不会导致内存崩溃。

加载专为 "text-generation" 管道设计的轻量级因果 LLM llm = pipeline("text-generation", model="Qwen/Qwen2.5-0.5B-Instruct", device_map="auto") def generate_rag_answer(query, retrieved_context): """使用聊天格式将提示、上下文和查询组合给 LLM""" context_str = " --- ".join(retrieved_context) messages = [ {"role": "system", "content": "你是一个有用的助手。仅使用提供的上下文回答用户问题。如果答案不在上下文中，请说 '我在提供的文档中找不到答案。'"}, {"role": "user", "content": f"上下文:\n{context_str}\n\n问题: {query}"} ] # 应用模型的聊天模板 prompt = llm.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) # 生成响应 response = llm(prompt, max_new_tokens=150, do_sample=False, temperature=0.1) # 提取提示后的新增生成文本 generated_text = response[0]['generated_text'] answer = generated_text[len(prompt):].strip() return answer

加载为“text-generation”流程构建的轻量级因果语言模型

llm

"text-generation"

model

"Qwen/Qwen2.5-0.5B-Instruct"

device_map

"auto"

generate_rag_answer

query

retrieved_context

"将提示、上下文和查询通过聊天格式组合到LLM中。"

context_str

" --- "

messages

{

"role"

"system"

"content"

"你是一个有用的助手。仅使用提供的上下文回答用户的问题。如果答案不在上下文中，请说'我在提供的文档中找不到答案'。"

}

"user"

"Context:\n{context_str}\n\nQuestion: {query}"

应用模型的聊天模板

prompt

tokenizer

apply_chat_template

tokenize

False

add_generation_prompt

生成响应

response

max_new_tokens

150

do_sample

temperature

0.1

仅提取提示后的全新生成文本

generated_text

'generated_text'

answer

strip

现在流水线的所有组件都已就绪。最后一步是将它们连接在一起，并通过向系统提问来测试其知识库（亚洲国家）：

def ask_rag_pipeline(question): print(f"Question: {question}") # 1. 对查询进行嵌入 query_vector = embedding_model.encode(question) # 2. 通过Scikit-Learn检索上下文 relevant_chunks = vector_db.retrieve(query_vector) # 3. 生成答案 answer = generate_rag_answer(question, relevant_chunks) print(f"\nRAG Answer: {answer}") print("\n--- 使用的来源 ---") for i, chunk in enumerate(relevant_chunks): print(f"{i+1}. {chunk[:100]}...") # 为可读性截断 # 进行测试 test_question = "Which Asian countries are best known for islands and beaches?" ask_rag_pipeline(test_question)

ask_rag_pipeline

question

"Question: {question}"

1. 对查询进行嵌入

query_vector

2. 通过Scikit-Learn检索上下文

relevant_chunks

3. 生成答案

"\nRAG Answer: {answer}"

"\n--- 使用的来源 ---"

enumerate

"{i+1}. {chunk[:100]}..."

为可读性截断

进行测试

test_question

"Which Asian countries are best known for islands and beaches?"

Question: Which Asian countries are best known for islands and beaches? RAG Answer: 最佳以岛屿和海滩闻名的亚洲国家包括：1. 菲律宾 - 以超过7000个岛屿和充满活力的节日而闻名。2. 泰国 - 以热带海滩和繁华的街头美食文化而著名。这两个国家拥有大量岛屿，并提供美丽的自然景观和多样的海滩体验。--- 使用的来源 --- 1. 山地景观和文化城市如清迈，而南部则提供世界闻名的岛屿... 2. 菲律宾是东南亚超过7000个岛屿的群岛国家，以壮丽的海滩... 3. 泰国是东南亚国家，以热带海滩、精致的寺庙和繁华的...

Which

Asian

countries

are

best

known

islands

and

beaches

?

RAG

that

include

1.

Philippines

its

over

000

vibrant

festivals

2.

Thailand

Famous

tropical

bustling

street

food

culture

These

two

have

a

significant

number

of

offer

both

beautiful

natural

scenery

diverse

beach

experiences

--

Sources

used

mountainous

landscapes

cultural

cities

like

Chiang

Mai

while

south

offers

world

renowned

is

an

archipelago

Southeast

Asia

stunning

be

3.

country

ornate

temples

尝试通过修改 test_question 中的文本提出不同问题。在我们刚刚构建的小规模设置中，行为和输出结果是合理的——但不要期待达到商业级大语言模型的性能：我们正在将小型模型与微型知识库进行结合。

最后说明

在本文中，我们构建了一个轻量级的端到端 RAG 管道，避免了重型框架或付费 API 的开销。该系统完全依赖 scikit-learn 向量检索，以及开源模型实现嵌入和生成，形成了一个透明的内存内系统。这种方法为原型设计、尝试不同的检索策略以及在扩大规模前验证概念提供了坚实的起点。

更多相关内容

• 使用 llama.cpp 在 Python 中构建 RAG 管道
• 在...中将 LLM 嵌入与表格特征结合
• 如何在...中结合 LLM 嵌入 + TF-IDF + 元数据
• 构建端到端情感分析管道...
• 构建健壮机器学习管道的最佳实践...
• 为长期运行代理构建上下文裁剪管道

/.entry /entry