{"id":18819249,"url":"https://github.com/yuanzhoulvpi2017/rust4senvec","last_synced_at":"2025-06-20T09:37:39.837Z","repository":{"id":65036228,"uuid":"580991773","full_name":"yuanzhoulvpi2017/Rust4SenVec","owner":"yuanzhoulvpi2017","description":"convert sentence to vector by nlp transformers model in Rust","archived":false,"fork":false,"pushed_at":"2023-01-07T10:53:31.000Z","size":22,"stargazers_count":8,"open_issues_count":0,"forks_count":1,"subscribers_count":3,"default_branch":"main","last_synced_at":"2024-11-08T00:23:24.261Z","etag":null,"topics":["bert","nlp","rust","sentence-embeddings","transformers-bert"],"latest_commit_sha":null,"homepage":"","language":"Jupyter Notebook","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":null,"status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/yuanzhoulvpi2017.png","metadata":{"files":{"readme":"readme.md","changelog":null,"contributing":null,"funding":null,"license":null,"code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null}},"created_at":"2022-12-22T01:31:08.000Z","updated_at":"2024-09-24T05:02:34.000Z","dependencies_parsed_at":"2023-02-06T23:28:36.819Z","dependency_job_id":null,"html_url":"https://github.com/yuanzhoulvpi2017/Rust4SenVec","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/yuanzhoulvpi2017%2FRust4SenVec","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/yuanzhoulvpi2017%2FRust4SenVec/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/yuanzhoulvpi2017%2FRust4SenVec/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/yuanzhoulvpi2017%2FRust4SenVec/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/yuanzhoulvpi2017","download_url":"https://codeload.github.com/yuanzhoulvpi2017/Rust4SenVec/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":231807406,"owners_count":18429519,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["bert","nlp","rust","sentence-embeddings","transformers-bert"],"created_at":"2024-11-08T00:21:23.595Z","updated_at":"2024-12-30T02:49:25.322Z","avatar_url":"https://github.com/yuanzhoulvpi2017.png","language":"Jupyter Notebook","funding_links":[],"categories":[],"sub_categories":[],"readme":"## 注意\n\n运行的时候，要用:`cargo run --release`，而不是简单的`cargo run`\n\n\n## 背景\nsentence embedding，或者叫sentence2vector，本质上是将文本转换成向量。\n\n之前做过sentence-transformers转onnx的，结合python的fastapi做推理的。之前还做过对sentence-transformers按照模型需求分开推理加速的。\n\n[https://www.zhihu.com/question/424133076/answer/2292331019](https://www.zhihu.com/question/424133076/answer/2292331019)\n\n[https://zhuanlan.zhihu.com/p/474396066](https://zhuanlan.zhihu.com/p/474396066)\n\n但是这些其实都是基于python环境的，而且模型对外的web端，使用的是fastapi（这个包在python语言里面算是效率很高了，但是在所有的语言中，效率很低）。\n\n对于线上的环境，我们总是希望在有限的硬件资源下，提高机器的使用率占比。\n\n因此，很多人就会在使用C++语言来做推理。也有的会使用trition（好像是这么拼写的）结合网络层RPC来做对外的服务开放。\n\n但是我不会C++、我不会RPC，我也不会trition，另外，我也不想将模型导出onnx（就是嫌弃麻烦），另外，我还嫌弃python慢、我还嫌弃fastapi慢。\n\n那怎么办？\n\n那么这里分享我的解决方法：使用rust来做句子转向量功能。\n\n## 句子转向量的三个步骤\n\n如果不太清楚，可以看看我这个回答：\n\n[https://www.zhihu.com/question/510987022/answer/2778610483](https://www.zhihu.com/question/510987022/answer/2778610483)\n\n\n### 1. 将文本通过tokenizer转换成`input_id`、`attention_mask`等\n\n你如果对transfromers包很熟悉的话，你有没有注意到它的tokenizer里面有一个`fast_tokenizer=True`参数？\n官网说，这个参数为True的时候，tokenizer会快一点。\n\n其实本质上是因为，当使用`fast_tokenizer`的时候，python会调用rust版本的tokenizer来做推理。\n\n而rust这种底层编程语言会比python快很多。\n\n### 2. 将`input_id`、`attention_mask`放入我们的bert模型中获得`output`\n其实就是放入bert模型中推理，计算得到`last_hidden_states`,`outputs`之类的。这个具体看源码即可，比如：\n```python\n       if not return_dict:\n            return (sequence_output, pooled_output) + encoder_outputs[1:]\n        return BaseModelOutputWithPoolingAndCrossAttentions(\n            last_hidden_state=sequence_output,\n            pooler_output=pooled_output,\n            past_key_values=encoder_outputs.past_key_values,\n            hidden_states=encoder_outputs.hidden_states,\n            attentions=encoder_outputs.attentions,\n            cross_attentions=encoder_outputs.cross_attentions,\n        )\n\n```\n\n但是如果下游任务不同，输出的肯定也是不一样的，就拿sentence embedding来说，最后返回的就是一个Tensor。（如果这步骤看不懂，可以看看我上面的链接）。\n\n### 3. 将上面的`output`通过web端来返回给接口调用方\n\n一般在python里面，大家都会使用fastapi、flask之类的，但是这个玩意和C++、rust语言的web库比起来，差远了。\n\n就拿排行榜来说，fastapi排名249名，而C++、rust都在top10的水平。\n\n![](https://files.mdnice.com/user/7098/db43649d-7b7e-4f18-a77d-51412f590cb0.jpg)\n\n\n具体可以看这个链接，有详细的性能排行榜\n\n[https://www.techempower.com/benchmarks/#section=data-r21](https://www.techempower.com/benchmarks/#section=data-r21)\n\n\n\n## 小结\n在上面三个步骤中，我们思考rust语言对应的解决办法。\n1. 第一步的问题，通过`tokenizer`库来解决了。[https://docs.rs/tokenizers/latest/tokenizers/](https://docs.rs/tokenizers/latest/tokenizers/)\n2. 第三步的问题，解决方法太多了。随便挑。\n\n那么第二步里面的模型推理怎么解决？\n\n其实已经有了，那就是`tch-rs`包可以来解决了。\n\n\n# rust推理部分详解\n\n## `tch-rs`包介绍\n\n众所周知，pytorch是有C++版本的，pytorch也是有python版本的。而且pytorch本身就是C++写的。\n\n那么`tch-rs`就可以理解成给rust语言写的pytorch。\n\n## `jit` module\n\n我们的pytorch训练完成之后，可以保存为`.bin`格式文件、也可保存为`.pt`、`.jit`格式的问题。\n别的语言，只要拿到上面任意的格式文件，基本上都可以通过`torch::jit::Cmodule`模块进行加载。\n那么rust也是可以的。\n\n需要注意的是：rust的`tch-rs`模块加载模型的时候，你这个模型的输出只能为Tensor。不能为别的数据格式。（这个是重点）\n\n## 使用python对模型 加载导出等\n\n接下来主要是使用python来对模型进行加载，导入导出等功能。\n\n\n\n```python\nfrom transformers import BertModel, BertTokenizer, BertConfig\nimport torch\nfrom typing import Tuple, Optional, List\nfrom torch import nn\nfrom transformers import BertPreTrainedModel\n```\n\n### 1. 加载一个预训练模型bert\n\n创建模型需要的tokneizer后的东西\n```python\n# step1 load bert model (such as tokenizer a text, and get tokens_tenosr, segements)\n\nenc = BertTokenizer.from_pretrained(\"bert-base-uncased\")\n\n# Tokenizing input text\ntext = \"[CLS] Who was Jim Henson ? [SEP] Jim Henson was a puppeteer [SEP]\"\ntokenized_text = enc.tokenize(text)\n\n# Masking one of the input tokens\nmasked_index = 8\ntokenized_text[masked_index] = '[MASK]'\nindexed_tokens = enc.convert_tokens_to_ids(tokenized_text)\nsegments_ids = [0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1]\n\n# Creating a dummy input\ntokens_tensor = torch.tensor([indexed_tokens])\nsegments_tensors = torch.tensor([segments_ids])\ndummy_input = [tokens_tensor, segments_tensors]\n```\n\n### 2. sentence2vector模型\n\n这里随便写了一个文本转向量的模型。\n\n`forward`返回的是一个Tensor，这个需要注意。\n```python\n# step2 costom my bert model\n\nclass Mybert4Sentence(BertPreTrainedModel):\n    # copy code from  class BertForSequenceClassification(BertPreTrainedModel):\n    def __init__(self, config):\n        super().__init__(config)\n        self.num_labels = config.num_labels\n        self.config = config\n\n        self.bert = BertModel(config)\n        classifier_dropout = (\n            config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob\n        )\n        self.dropout = nn.Dropout(classifier_dropout)\n        self.classifier = nn.Linear(config.hidden_size, config.num_labels)\n\n        # Initialize weights and apply final processing\n        self.post_init()\n\n    def forward(\n            self,\n            input_ids: Optional[torch.Tensor] = None,\n            attention_mask: Optional[torch.Tensor] = None,\n            token_type_ids: Optional[torch.Tensor] = None,\n            position_ids: Optional[torch.Tensor] = None,\n            head_mask: Optional[torch.Tensor] = None,\n            inputs_embeds: Optional[torch.Tensor] = None,\n            labels: Optional[torch.Tensor] = None,\n            output_attentions: Optional[bool] = None,\n            output_hidden_states: Optional[bool] = None,\n            return_dict: Optional[bool] = None,\n    ) -\u003e torch.Tensor:\n        return_dict = return_dict if return_dict is not None else self.config.use_return_dict\n\n        outputs = self.bert(\n            input_ids,\n            attention_mask=attention_mask,\n            token_type_ids=token_type_ids,\n            position_ids=position_ids,\n            head_mask=head_mask,\n            inputs_embeds=inputs_embeds,\n            output_attentions=output_attentions,\n            output_hidden_states=output_hidden_states,\n            return_dict=return_dict,\n        )\n\n        pooled_output = outputs[1]\n\n        return pooled_output\n```\n### 3. 把上面的模型，通过jit模式导出，并且保存\n\n注意这里的`traced_bert.pt`模型的路径。后面在rust里面要用到。\n```python\n# step 3 load model and save model to .pt file\n\n\nmodel2 = Mybert4Sentence.from_pretrained(\"bert-base-uncased\", torchscript=True)\n\ntraced_model = torch.jit.trace(model2, [tokens_tensor, segments_tensors])\ntorch.jit.save(traced_model, \"traced_bert.pt\")\n```\n\n### 4. 再在python里面加载模型，看看输出的对不对\n\n```python\n# step 4 load .pt file then check it\nloaded_model = torch.jit.load(\"traced_bert.pt\")\nloaded_model.eval()\n\nloaded_model(*dummy_input).shape\n\nloaded_model(*dummy_input)[:, :10]\n```\n\n到这里基本上python部分已经做完了。\n\n接下来到rust部分。\n\n## rust部分\n\n### `tch-rs`依赖\n`tch-rs`需要`libtorch`添加到环境变量中，具体的添加方式，可以看我之前的文章（或者`tch-rs`介绍）\n[https://zhuanlan.zhihu.com/p/589055479](https://zhuanlan.zhihu.com/p/589055479)\n\n与此同时，对应的项目创建等细节也都不需要再考虑。\n\n直接到`src/main.rs`部分。\n\n### 具体main.rs代码\n\n这里需要注意`traced_bert.pt`的路径。\n\n```rust\nuse tch::jit;\nuse tch::Tensor;\n\nfn main() {\n\n    // load model which generate from  `trans_model.py` file\n    let model = jit::CModule::load(\n        \"./traced_bert.pt\",\n    )\n    .unwrap();\n\n    // just generate a small attention_mask and input_ids\n    let attention_mask = Tensor::of_slice2(\u0026[[\n        101, 2040, 2001, 3958, 27227, 1029, 102, 3958, 103, 2001, 1037, 13997, 11510, 102,\n    ]]);\n\n    let input_ids = Tensor::of_slice2(\u0026[[0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1]]);\n\n\n    // infer model by use `forward_ts` method\n    let result = model.forward_ts(\u0026[attention_mask, input_ids]).unwrap();\n\n    // show the result\n    println!(\"{:?}\", result);\n    result.slice(1, 0, 10, 1).print();\n\n    println!(\"Hello, world!\");\n\n```\n\n\n最后跑通就会发现，计算的数值和python是一模一样的。\n\n![](https://files.mdnice.com/user/7098/4218f347-c93d-4605-b30a-5f29943c1655.jpg)\n\n\n## 结束\n\n到这里，基本上就结束了。剩下的各个部分，只需要像是拼积木一样搭建起来就可以了。一个完整的rust做sentence embedding就可以实现了。\n\n\n# 最后\n\n## rsut难么？\n我感觉rust大部分代码其实写起来和python差不多，但是有些东西确实难，比如生命周期之类的。\n\n但是我感觉网上很多算法工程师，其实python代码写的都坑坑巴巴的，如果用rust来做，估计是会要了他们老命。\n因此，如果你连pytorch、transformers等优秀的包都不会用，还是别学rust了。\n\n## 为什么不用rust-bert？\n感觉rust-bert想做rust语言里面的transformers包（python的）。但是感觉路还能长，而且python版本的transformers用起来已经非常简单了。\n\n不用rust-bert来做，是因为感觉rust-bert客制化程度还不够，而且感觉封装好多东西，太麻烦了。\n\n\n## 上面的可以用在c++中么？\n其实，python导出`.pt`文件，我觉得肯定是可以用在c++的libtorch里面的。\n\n但是，tokenizer有C++版本的么？这个我不清楚，不太能回答。\n\n\n## rust和C++比较？\n\n这个真没办法比较，而且说出来很容易被一些Rust极端脑残粉攻击。只能说哪个好用用哪个~\n\n# 更多\n1. 具体细节都会定时更新，可以查看这个文章[https://zhuanlan.zhihu.com/p/589444069](https://zhuanlan.zhihu.com/p/589444069)\n2. 或者直接关注我[https://www.zhihu.com/people/fa-fa-1-94](https://www.zhihu.com/people/fa-fa-1-94)\n\n\n\n\n\n\n\n\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fyuanzhoulvpi2017%2Frust4senvec","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fyuanzhoulvpi2017%2Frust4senvec","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fyuanzhoulvpi2017%2Frust4senvec/lists"}