Giới thiệu EdgeBERT

Trong ứng dụng Rust của tôi, tôi cần một giải pháp tìm kiếm ngữ nghĩa, cho phép người dùng tìm kiếm từ "bác sĩ" và vẫn nhận được tài liệu đề cập đến "người hành nghề y" hay "thầy thuốc". Tôi muốn một giải pháp nhúng BERT nhẹ nhàng trong Rust, đủ nhỏ để chạy trên các thiết bị biên, trình duyệt và máy chủ mà không gặp rắc rối.

Hệ thống này biến đổi văn bản thành các vector mô tả ý nghĩa. Các từ tương tự sẽ được chuyển thành các số tương tự.

• bác sĩ → [0.2, 0.5, -0.1, 0.8, ...]
• thầy thuốc→ [0.2, 0.4, -0.1, 0.7, ...] ✅ tương tự
• chuối → [0.9, -0.3, 0.6, -0.2, ...] ❌ khác nhau

Những vấn đề với các giải pháp hiện có

So với Python, phương pháp tiêu chuẩn thường rất nặng nề:

Chỉ để tạo ra các nhúng, một môi trường ảo mới đã chiếm tới 6.8 GB, chủ yếu là do PyTorch, các bộ phân tách và trọng số mô hình.

ONNX Runtime

Tuy nhiên, tôi đang sử dụng Rust, ai đó đã đề cập đến ort, tôi sẽ sử dụng ONNX Runtime từ Rust. Nó có khó không?

pub fn new(model_path: &str, tokenizer_path: &str) -> Result<Self> {
    let environment = ort::environment::init()
        .with_execution_providers([CUDAExecutionProvider::default().build()]).commit()?;
}
// ... 150 dòng để chỉ để làm cho encode hoạt động

Kết quả cuối cùng cũng thành công, nhưng ort đã kéo theo hơn 80 crates, làm cho bản phát hành của tôi mở rộng lên 350 MB và phụ thuộc vào các thư viện hệ thống như libstdc++, libpthread, libm, libc. Các phiên bản không tương thích của OpenSSL đã gây ra thêm rắc rối.

Xung đột thư viện hệ thống

error: OpenSSL 3.3 required
$ openssl version
OpenSSL 1.1.1k  # Không thể nâng cấp - sẽ phá vỡ các phụ thuộc RHEL

Các phụ thuộc C++ yêu cầu OpenSSL 3.3. Hệ thống RHEL của tôi chỉ có 1.1. Việc gỡ bỏ 1.1 sẽ phá vỡ một nửa hệ thống của tôi.

Ban đầu tôi đã cố gắng xây dựng một giải pháp RAG nhẹ nhàng ngoại tuyến, nhưng một phụ thuộc này đã trở thành một thách thức lớn.

Những gì tôi thực sự muốn

API mà tôi đang tìm kiếm là:

model = SentenceTransformer('all-MiniLM-L6-v2')
embeddings = model.encode(texts)

Giải pháp

Vì vậy, tôi đã xây dựng engine inferencing của riêng mình bằng Rust:

use edgebert::{Model, ModelType};
let model = Model::from_pretrained(ModelType::MiniLML6V2)?;
let embeddings = model.encode(texts, true)?;

• Kích thước nhị phân: 5MB
• RAM: 200MB
• Không gặp rắc rối với phụ thuộc
• Độ chính xác tương tự (0.9997 tương quan)
• Tính năng tùy chọn BLAS

Hiệu suất

Các phép đo ban đầu rất hứa hẹn, và sau khi tối ưu hóa các quy trình nhân ma trận, dưới đây là cách EdgeBERT so sánh với sentence-transformers trên CPU:

👉 EdgeBERT nhanh hơn tới 3× trong các kịch bản đa luồng đơn. Điều này là do các ma trận của MiniLM (384×384) nhỏ, có nghĩa là chi phí từ việc điều phối luồng có thể vượt quá lợi ích từ việc phân tán.

Hiệu suất CPU chỉ là một nửa câu chuyện. Hiệu quả bộ nhớ, đặc biệt là việc sử dụng RAM trong quá trình mã hóa, là rất quan trọng. Chúng ta có thể thấy sự khác biệt đáng kể ở đây:

Dấu chân bộ nhớ của EdgeBERT không chỉ nhỏ hơn mà còn ổn định hơn, tránh được các đỉnh phân bổ lớn ban đầu mà giải pháp dựa trên PyTorch gặp phải.

Độ chính xác

So sánh với sentence-transformers Python:

EdgeBERT:   [-0.0344, 0.0309, 0.0067, 0.0261, -0.0394, ...]
Python:     [-0.0345, 0.0310, 0.0067, 0.0261, -0.0394, ...]
Cosine similarity: **0.9997**

Sự khác biệt về làm tròn giữa các phép toán số thực, 99.97% giống nhau.

Tại sao điều này quan trọng (ngay cả khi bạn không làm ML)

Điều này cho phép:

• Tìm kiếm thông minh: Người dùng tìm thấy những gì họ nghĩa, không chỉ những gì họ đã gõ
• Khuyến nghị tốt hơn: "Nếu bạn thích X, bạn sẽ thích Y" dựa trên ý nghĩa
• Phát hiện trùng lặp: Tìm các vấn đề/tài liệu tương tự ngay cả với cách diễn đạt khác nhau
• Kiểm duyệt nội dung: Phát hiện nội dung có hại bất kể cách diễn đạt
• Tính năng RAG/AI: Cung cấp cho LLM ngữ cảnh đúng mà không cần phải khớp từ khóa

Tất cả chỉ trong 5MB của Rust. Không cần Python.

Khi nào nên sử dụng

Sử dụng EdgeBERT khi:

• Bạn cần các nhúng mà không cần Python
• Kích thước triển khai là quan trọng (5MB so với 6.8GB)
• Chạy trên các thiết bị biên hoặc trình duyệt
• Bộ nhớ bị hạn chế
• Hiệu suất đơn luồng là điều quan trọng

Sử dụng sentence-transformers khi:

• Bạn cần tăng tốc GPU
• Sử dụng nhiều kiến trúc mô hình
• Đã có trong hệ sinh thái Python
• Cần đầy đủ bộ công cụ Hugging Face

WebAssembly

Vì nó hoàn toàn là Rust với các phụ thuộc tối thiểu, nó có thể biên dịch sang WASM:

import init, { WasmModel, WasmModelType } from './pkg/edgebert.js';

const model = await WasmModel.from_type(WasmModelType.MiniLML6V2);
const embeddings = model.encode(texts, true);

Kích thước nhị phân WASM 429KB + 30MB trọng số mô hình. Chạy trong các trình duyệt.

Tôi đã phải triển khai bộ phân tách WordPiece từ đầu - crate bộ phân tách có các phụ thuộc C không biên dịch được sang WASM.

Cách hoạt động

BERT là các phép toán ma trận theo một thứ tự cụ thể:

1. Phân tách - Các token WordPiece thành ID
2. Nhúng - vector 384 chiều (từ + vị trí + đoạn)
3. Tự chú ý - Q·K^T/√d, softmax, nhân với V
4. Tiến bộ - Tuyến tính, GELU, Tuyến tính
5. Tập hợp - Trung bình các token thành nhúng câu

Mỗi lớp transformer lặp lại chú ý và tiến bộ, tinh chỉnh các đại diện. MiniLM có 6 lớp.

Cài đặt cốt lõi chỉ khoảng ~500 dòng trong src/lib.rs.

Không có phép màu, chỉ là thuật toán transformer, viết bằng Rust.

Github

Cấu hình

Để đạt hiệu suất tốt nhất:

# EdgeBERT - nhanh nhất trong đa luồng đơn
export OPENBLAS_NUM_THREADS=1; cargo run --release --features openblas

# Python - để nó tự động điều chỉnh luồng
python native.py

Cài đặt

[dependencies]
edgebert = "0.3.4"

Lộ trình & Công việc tương lai

EdgeBERT được thiết kế tập trung và tối thiểu, nhưng có nhiều hướng đi thú vị trong tương lai:

• Hỗ trợ GPU: Thêm hỗ trợ wgpu cho tăng tốc GPU đa nền tảng là ưu tiên hàng đầu.
• Nhiều kiến trúc hơn: Mở rộng vượt ra ngoài all-MiniLM-L6-v2 để hỗ trợ các mô hình hiệu quả khác.
• Lượng hóa: Thực hiện lượng hóa mô hình để giảm kích thước mô hình hơn nữa và cải thiện hiệu suất trên CPU và các bộ điều khiển vi mô.

Các yêu cầu kéo luôn được hoan nghênh!

Mã

Github

Hầu hết các triển khai nằm trong một tệp. Các yêu cầu kéo được hoan nghênh.

Tôi đã xây dựng điều này vì tôi cần nó.

Tôi chia sẻ nó vì có thể bạn cũng cần nó. 🚀

Chi tiết Benchmark

• EdgeBERT: cargo run --release --features openblas --bin native với OPENBLAS_NUM_THREADS đã đặt
• Python: python native.py với OMP_NUM_THREADS đã đặt
• Cài đặt sentence_transformers sử dụng pip trong một môi trường ảo và kiểm tra kích thước tập tin du -sh venv/lib/python*/site-packages/
• Kiểm tra các phụ thuộc Rust với cargo tree | wc -l
• Kiểm tra các phụ thuộc venv bằng pip list | wc -l
• Kích thước tệp WASM ls -lh examples/pkg/edgebert_bg.wasm
• Kích thước tệp gốc ls -lh target/release/native
• Sử dụng /usr/bin/time -v để đo kích thước bộ nhớ tối đa (kbyte)