文盘rust--使用 Rust 构建RAG

作者：京东科技贾世闻

RAG（Retrieval-Augmented Generation）技术在AI生态系统中扮演着至关重要的角色，特别是在提升大型语言模型（LLMs）的准确性和应用范围方面。RAG通过结合检索技术与LLM提示，从各种数据源检索相关信息，并将其与用户的问题结合，生成准确且丰富的回答。这一机制特别适用于需要应对信息不断更新的场景，因为大语言模型所依赖的参数知识本质上是静态的。

RAG技术的优势在于它能够利用外部知识库，引用大量的信息，以提供更深入、准确且有价值的答案，提高了生成文本的可靠性。此外，RAG模型具备检索库的更新机制，可以实现知识的即时更新，无需重新训练模型，这在及时性要求高的应用中占优势。

目前构建一个RAG并不是一个非常的事情。使用Langchain等成熟技术架构百十行代码就能构建一个Demo。那能不能利用目前的Rust生态构建一个简易的RAG。说干就干，本期和大家聊聊如果使用rust语言构建rag。

构建知识库

知识库构建主要是模型+向量库，为了保证所有系统中所有组件都使用rust构建，在限量数据库的选型上我们使用qdrant,纯rust构建的向量数据库。

知识库的构建最重要的步骤是embedding的过程。
过程如下：

模型加载
获取文本token
通过模型获取文本的Embedding
下面详细介绍每个过程细节及代码实现。

模型加载

以下代码用于加载模型和tokenizer

1async fn build_model_and_tokenizer(model_config: &ConfigModel) -> Result<(BertModel, Tokenizer)> {
2    let device = Device::new_cuda(0)?;
3    let repo = Repo::with_revision(
4        model_config.model_id.clone(),
5        RepoType::Model,
6        model_config.revision.clone(),
7    );
8    let (config_filename, tokenizer_filename, weights_filename) = {
9        let api = ApiBuilder::new()    
10            .build()?;
11        let api = api.repo(repo);
12        let config = api.get("config.json").await?;
13        let tokenizer = api.get("tokenizer.json").await?;
14        let weights = if model_config.use_pth {
15            api.get("pytorch_model.bin").await?
16        } else {
17            api.get("model.safetensors").await?
18        };
19        (config, tokenizer, weights)A
20    };
21    let config = std::fs::read_to_string(config_filename)?;
22    let mut config: Config = serde_json::from_str(&config)?;
23    let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
24
25    let vb = if model_config.use_pth {
26        VarBuilder::from_pth(&weights_filename, DTYPE, &device)?
27    } else {
28        unsafe { VarBuilder::from_mmaped_safetensors(&[weights_filename], DTYPE, &device)? }
29    };
30    if model_config.approximate_gelu {
31        config.hidden_act = HiddenAct::GeluApproximate;
32    }
33    let model = BertModel::load(vb, &config)?;
34    Ok((model, tokenizer))
35}

模型和tokenizer是系统中频繁调用的部分，所以为了避免重复加载，通过OnceCell构建静态全局变量

1pub static GLOBAL_EMBEDDING_MODEL: OnceCell> = OnceCell::const_new();
2
3pub async fn init_model_and_tokenizer() -> Arc<(BertModel, Tokenizer)> {
4    let config = get_config().unwrap();
5    let (m, t) = build_model_and_tokenizer(&config.model).await.unwrap();
6    Arc::new((m, t))
7}

在系统启动时加载模型

1GLOBAL_RUNTIME.block_on(async {
2    log::info!("global runtime start!");
3    // 加载model
4    GLOBAL_EMBEDDING_MODEL
5        .get_or_init(init_model_and_tokenizer)
6        .await;
7});

Embedding 过程主要由一下函数实现。

1pub async fn embedding_setence(content: &str) -> Result>> {
2    let m_t = GLOBAL_EMBEDDING_MODEL.get().unwrap();
3    let tokens = m_t
4        .1
5        .encode(content, true)
6        .map_err(E::msg)?
7        .get_ids()
8        .to_vec();
9    let token_ids = Tensor::new(&tokens[..], &m_t.0.device)?.unsqueeze(0)?;
10    let token_type_ids = token_ids.zeros_like()?;
11    let sequence_output = m_t.0.forward(&token_ids, &token_type_ids)?;
12    let (_n_sentence, n_tokens, _hidden_size) = sequence_output.dims3()?;
13    let embeddings = (sequence_output.sum(1)? / (n_tokens as f64))?;
14    let embeddings = normalize_l2(&embeddings)?;
15    let encodings = embeddings.to_vec2::()?;
16    Ok(encodings)
17}

函数通过tokenizer encode输入的文本，再使用模型embed token 获取一个三维的Tensor，最后归一化张量。

数据入库

知识库构建是将待检索文本向量化后存储到向量数据库的过程。
本次使用京东云文档作为原始文本，加工为以下格式。数据加工过程这里就不累述了。

1{
2    "content": "# 服务计费\n\n主机迁移服务自身为免费服务，但是迁移目标为云主机镜像时，迁移过程依赖系统自动创建的 中转资源的配合，这些资源中涉及部分付费资源，会产生相应费用。\n\n迁移过程涉及的中转资付费资源配置及计费说明如下（单个迁移任务）：\n\n|  | 云主机 | 云硬盘 | 弹性公网IP |\n| --- | --- | --- | ------ |\n| 计费类型 | 按配置 | 按配置 | 按用量 |\n| 规格配置 | 2C4G （c.n2.large或c.n3.large或c.n1.large） | 系统盘：40G 通用型SSD 数据盘：通用型SSD，数量及容量取决于源服务器系统盘及数据盘情况 | 30Mbps |\n| 费用预估 | 云主机规格每小时价格\\*迁移时长 | 云硬盘规格每小时价格\\*迁移时长 | 弹性公网IP每小时保有费\\*迁移时长 仅使用弹性公网IP入方向流量，只涉及IP保有用，不涉及流量费用 |\n\n> 提示：\n>\n> * 迁移时长取决于源服务器迁数据量以及源服务器公网出方向带宽，公网连接顺畅且源服务器公网出方向带宽不低于22.5Mbps的情况下（主机迁移为单线程传输，京东云云主机在单流传输下实际带宽为带宽上限的75%左右），实际数据容量为5GB的磁盘迁移时长在30分钟左右。\n> * 中转实例实例绑定的安全组出方向默认拒绝所有流量，因此默认情况下降不会产生任何公网出方向收费流量，但此配置也影响了云主机部分监控指标的上报，如需要监控中转实例的全部监控数据，可自行调整安全组规则方向出方向443端口。",
3    "title": "服务计费说明",
4    "product": "云主机 CVM",
5    "url": "https://docs.jdcloud.com/cn/virtual-machines/server-migration-service/billing"
6}

入库完整代码如下：

1use anyhow::Error as E;
2use anyhow::Result;
3use candle_core::Device;
4use candle_core::Tensor;
5use candle_nn::VarBuilder;
6use candle_transformers::models::bert::{BertModel, Config, HiddenAct, DTYPE};
7use hf_hub::{api::tokio::Api, Repo, RepoType};
8use qdrant_client::qdrant::CollectionExistsRequest;
9use qdrant_client::qdrant::CreateCollectionBuilder;
10use qdrant_client::qdrant::DeleteCollection;
11use qdrant_client::qdrant::Distance;
12use qdrant_client::qdrant::UpsertPointsBuilder;
13use qdrant_client::qdrant::VectorParamsBuilder;
14use qdrant_client::Payload;
15use qdrant_client::{
16    qdrant::{
17        CollectionOperationResponse, CreateCollection, PointStruct, PointsOperationResponse,
18        UpsertPoints,
19    },
20    Qdrant,
21};
22use serde::{Deserialize, Serialize};
23use serde_json::from_str;
24use std::fs;
25use std::sync::Arc;
26use tokenizers::Tokenizer;
27use tokio::sync::OnceCell;
28use uuid::Uuid;
29use walkdir::WalkDir;
30
31#[derive(Debug, Serialize, Deserialize, Clone)]
32pub struct Doc {
33    pub content: String,
34    pub title: String,
35    pub product: String,
36    pub url: String,
37}
38
39#[derive(Debug, PartialEq, Serialize, Deserialize, Clone)]
40pub struct ModelConfig {
41    #[serde(default = "ModelConfig::model_id_default")]
42    pub model_id: String,
43    #[serde(default = "ModelConfig::revision_default")]
44    pub revision: String,
45    #[serde(default = "ModelConfig::use_pth_default")]
46    pub use_pth: bool,
47    #[serde(default = "ModelConfig::approximate_gelu_default")]
48    pub approximate_gelu: bool,
49}
50
51impl Default for ModelConfig {
52    fn default() -> Self {
53        Self {
54            model_id: Self::model_id_default(),
55            revision: Self::revision_default(),
56            use_pth: Self::use_pth_default(),
57            approximate_gelu: Self::approximate_gelu_default(),
58        }
59    }
60}
61
62impl ModelConfig {
63    fn model_id_default() -> String {
64        "moka-ai/m3e-large".to_string()
65    }
66    fn revision_default() -> String {
67        "main".to_string()
68    }
69    fn use_pth_default() -> bool {
70        true
71    }
72    fn approximate_gelu_default() -> bool {
73        false
74    }
75}
76
77pub static GLOBAL_MODEL: OnceCell> = OnceCell::const_new();
78pub static GLOBAL_TOKEN: OnceCell> = OnceCell::const_new();
79
80pub async fn init_model() -> Arc {
81    let config = ModelConfig::default();
82    let (m, _) = build_model_and_tokenizer(&config).await.unwrap();
83    Arc::new(m)
84}
85
86pub async fn init_tokenizer() -> Arc {
87    let config = ModelConfig::default();
88    let (_, t) = build_model_and_tokenizer(&config).await.unwrap();
89    Arc::new(t)
90}
91
92async fn build_model_and_tokenizer(model_config: &ModelConfig) -> Result<(BertModel, Tokenizer)> {
93    let device = Device::new_cuda(0)?;
94    let repo = Repo::with_revision(
95        model_config.model_id.clone(),
96        RepoType::Model,
97        model_config.revision.clone(),
98    );
99    let (config_filename, tokenizer_filename, weights_filename) = {
100        let api = Api::new()?;
101        let api = api.repo(repo);
102        let config = api.get("config.json").await?;
103        let tokenizer = api.get("tokenizer.json").await?;
104        let weights = if model_config.use_pth {
105            api.get("pytorch_model.bin").await?
106        } else {
107            api.get("model.safetensors").await?
108        };
109        (config, tokenizer, weights)
110    };
111    let config = std::fs::read_to_string(config_filename)?;
112    let mut config: Config = serde_json::from_str(&config)?;
113    let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
114
115    let vb = if model_config.use_pth {
116        VarBuilder::from_pth(&weights_filename, DTYPE, &device)?
117    } else {
118        unsafe { VarBuilder::from_mmaped_safetensors(&[weights_filename], DTYPE, &device)? }
119    };
120    if model_config.approximate_gelu {
121        config.hidden_act = HiddenAct::GeluApproximate;
122    }
123    let model = BertModel::load(vb, &config)?;
124    Ok((model, tokenizer))
125}
126
127pub async fn embedding_setence(content: &str) -> Result>> {
128    let m = GLOBAL_MODEL.get().unwrap();
129    let t = GLOBAL_TOKEN.get().unwrap();
130    let tokens = t.encode(content, true).map_err(E::msg)?.get_ids().to_vec();
131
132    let token_ids = Tensor::new(&tokens[..], &m.device)?.unsqueeze(0)?;
133    let token_type_ids = token_ids.zeros_like()?;
134
135    let sequence_output = m.forward(&token_ids, &token_type_ids)?;
136    let (_n_sentence, n_tokens, _hidden_size) = sequence_output.dims3()?;
137    let embeddings = (sequence_output.sum(1).unwrap() / (n_tokens as f64))?;
138    let embeddings = normalize_l2(&embeddings).unwrap();
139    let encodings = embeddings.to_vec2::()?;
140    Ok(encodings)
141}
142
143pub fn normalize_l2(v: &Tensor) -> Result {
144    Ok(v.broadcast_div(&v.sqr()?.sum_keepdim(1)?.sqrt()?)?)
145}
146
147pub struct QdrantClient {
148    client: Qdrant,
149}
150
151impl QdrantClient {
152    pub async fn create_collection(
153        &self,
154        request: impl Into,
155    ) -> Result {
156        let resp = self.client.create_collection(request).await?;
157        Ok(resp)
158    }
159
160    pub async fn delete_collection(
161        &self,
162        request: impl Into,
163    ) -> Result {
164        let resp = self.client.delete_collection(request).await?;
165        Ok(resp)
166    }
167
168    pub async fn collection_exists(
169        &self,
170        request: impl Into,
171    ) -> Result {
172        let resp = self.client.collection_exists(request).await?;
173        Ok(resp)
174    }
175
176    pub async fn load_dir(&self, path: &str, collection_name: &str) {
177        let mut points = vec![];
178        for entry in WalkDir::new(path)
179            .into_iter()
180            .filter_map(Result::ok)
181            .filter(|e| !e.file_type().is_dir() && e.file_name().to_str().is_some())
182        {
183            if let Some(p) = entry.path().to_str() {
184                let id = Uuid::new_v4();
185                let content = match fs::read_to_string(p) {
186                    Ok(c) => c,
187                    Err(_) => continue,
188                };
189
190                let doc = match from_str::(content.as_str()) {
191                    Ok(d) => d,
192                    Err(_) => continue,
193                };
194                let mut payload = Payload::new();
195                payload.insert("content", doc.content);
196                payload.insert("title", doc.title);
197                payload.insert("product", doc.product);
198                payload.insert("url", doc.url);
199                let vector_contens = embedding_setence(content.as_str()).await.unwrap();
200                let ps = PointStruct::new(id.to_string(), vector_contens[0].clone(), payload);
201                points.push(ps);
202
203                if points.len().eq(&100) {
204                    let p = points.clone();
205                    self.client
206                        .upsert_points(UpsertPointsBuilder::new(collection_name, p).wait(true))
207                        .await
208                        .unwrap();
209                    points.clear();
210                    println!("batch finish");
211                }
212            }
213        }
214
215        if points.len().gt(&0) {
216            self.client
217                .upsert_points(UpsertPointsBuilder::new(collection_name, points).wait(true))
218                .await
219                .unwrap();
220        }
221    }
222}
223
224#[tokio::main]
225async fn main() {
226    // 加载模型
227    GLOBAL_MODEL.get_or_init(init_model).await;
228    GLOBAL_TOKEN.get_or_init(init_tokenizer).await;
229
230    let collection_name = "default_collection";
231
232    // The Rust client uses Qdrant's GRPC interface
233    let qdrant = Qdrant::from_url("http://localhost:6334").build().unwrap();
234    let qdrant_client = QdrantClient { client: qdrant };
235
236    if !qdrant_client
237        .collection_exists(collection_name)
238        .await
239        .unwrap()
240    {
241        qdrant_client
242            .create_collection(
243                CreateCollectionBuilder::new(collection_name)
244                    .vectors_config(VectorParamsBuilder::new(1024, Distance::Dot)),
245            )
246            .await
247            .unwrap();
248    }
249
250    qdrant_client
251        .load_dir("/root/jd_docs", collection_name)
252        .await;
253
254    println!("{:?}", qdrant_client.client.health_check().await);
255}

以上代码要完成的任务如下：

推理服务

推理服务使用 rust 构建的 mistral.rs。

由于国内访问hf 并不方便所以先通过 https://hf-mirror.com/ 现将模型下载到本地。本次使用qwen模型

HF_ENDPOINT="https://hf-mirror.com"  huggingface-cli download --repo-type model --resume-download Qwen/Qwen2-7B --local-dir /root/Qwen2-7B

启动 mistralrs-server

1git clone https://github.com/EricLBuehler/mistral.rs
2cd mistral.rs
3cargo run  --bin mistralrs-server  --features cuda -- --port 3333 plain -m /root/Qwen2-7B  -a qwen2

推理服务调用

mistral.rs 支持 Openai 的 api接口，使用 openai-api-rs调用即可。推理时间比较长 timeout 要设置长一些，若timeout 时间太短有可能不等返回结果就已经强制超时。

1pub static GLOBAL_OPENAI_CLIENT: Lazy> = Lazy::new(|| {
2    let mut client =
3        OpenAIClient::new_with_endpoint("http://10.0.0.7:3333/v1".to_string(), "EMPTY".to_string());
4    client.timeout = Some(30);
5    Arc::new(client)
6});
7
8pub async fn inference(content: &str, max_len: i64) -> Result> {
9    let req = ChatCompletionRequest::new(
10        "".to_string(),
11        vec![chat_completion::ChatCompletionMessage {
12            role: chat_completion::MessageRole::user,
13            content: chat_completion::Content::Text(content.to_string()),
14            name: None,
15            tool_calls: None,
16            tool_call_id: None,
17        }],
18    )
19    .max_tokens(max_len);
20
21    let cr = GLOBAL_OPENAI_CLIENT.chat_completion(req).await?;
22    Ok(cr.choices[0].message.content.clone())
23}

将Retriever和推理服务集成

1pub async fn answer(question: &str, max_len: i64) -> Result> {
2    let retriver = retriever(question, 1).await?;
3    let mut context = "".to_string();
4
5    for sp in retriver.result {
6        let payload = sp.payload;
7        let product = payload.get("product").unwrap().to_string();
8        let title = payload.get("title").unwrap().to_string();
9        let content = payload.get("content").unwrap().to_string();
10        context.push_str(&product);
11        context.push_str(&title);
12        context.push_str(&content);
13    }
14
15    let prompt = format!(
16        "你是一个云技术专家, 使用以下检索到的Context回答问题。用中文回答问题。
17        Question: {}
18        Context: {}
19        ",
20        question, context
21    );
22
23    log::info!("{}", prompt);
24
25    let req = ChatCompletionRequest::new(
26        "".to_string(),
27        vec![chat_completion::ChatCompletionMessage {
28            role: chat_completion::MessageRole::user,
29            content: chat_completion::Content::Text(prompt),
30            name: None,
31            tool_calls: None,
32            tool_call_id: None,
33        }],
34    )
35    .max_tokens(max_len);
36
37    let cr = GLOBAL_OPENAI_CLIENT.chat_completion(req).await?;
38    Ok(cr.choices[0].message.content.clone())
39}

后记

完整工程地址[embedding_server]https://github.com/jiashiwen/embedding_server

后续工程问题，多卡推理，多机推理，推理加速

资源对比

GPU 型号

1|=========================================+========================+======================|
2|   0  NVIDIA A30                     Off |   00000000:00:07.0 Off |                    0 |
3| N/A   30C    P0             29W /  165W |       0MiB /  24576MiB |      0%      Default |
4|                                         |                        |             Disabled |
5+-----------------------------------------+------------------------+----------------------+

Embedding 资源

m3e-large

vllm

1+-----------------------------------------------------------------------------------------+
2| Processes:                                                                              |
3|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
4|        ID   ID                                                               Usage      |
5|=========================================================================================|
6|    0   N/A  N/A    822789      C   ...iprojects/rag_demo/.venv/bin/python       1550MiB |
7+-----------------------------------------------------------------------------------------+

candle

1+-----------------------------------------------------------------------------------------+
2| Processes:                                                                              |
3|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
4|        ID   ID                                                               Usage      |
5|=========================================================================================|
6|    0   N/A  N/A    823261      C   target/debug/embedding_server                1484MiB |
7+-----------------------------------------------------------------------------------------+

推理资源

Qwen1.5-1.8B-Chat

vllm

1|=========================================================================================|
2|    0   N/A  N/A    822437      C   /usr/bin/python3                            20440MiB |
3+-----------------------------------------------------------------------------------------+

mistral.rs

1|=========================================================================================|
2|    0   N/A  N/A    822174      C   target/debug/mistralrs-server               22134MiB |
3+-----------------------------------------------------------------------------------------+

Qwen2-7B

vllm 现存溢出

[rank0]: OutOfMemoryError: CUDA out of memory. Tried to allocate 9.25 GiB. GPU

mistral.rs

1|=========================================================================================|
2|    0   N/A  N/A    656923      C   target/debug/mistralrs-server               22006MiB |
3+-----------------------------------------------------------------------------------------+

从实际情况来看，Embedding 模型再资源占用情况 rust candle框架使用显存略小些；推理模型Qwen1.5-1.8B-Chat，vllm 资源占用略小。Qwen2-7B vllm直接显存溢出。

坑

大部分框架中使用 hf-hub 采用同步调用，不支持境内的mirror。动手改造

src/api/tokio.rs

1
2impl ApiBuilder {
3    /// Set endpoint example 'https://hf-mirror.com'
4    pub fn with_endpoint(mut self, endpoint: &str) -> Self {
5        self.endpoint = endpoint.to_string();
6        self
7    }
8}