BGE-Base-Zh-V1.5 端侧使用教程
本项目基于 ai-engine-direct-helper (QAI_AppBuilder)
https://github.com/quic/ai-engine-direct-helper.git
1. 模型下载地址 (包含对应的上下文二进制文件)
https://www.aidevhome.com/?id=51
2. QNN SDK (v2.41) 下载地址
https://qpm.qualcomm.com/.../Qualcomm_AI_Runtime_SDK?version=2.41.0.251128
第一部分：Windows 平台推理指南
本部分介绍如何在 Windows 平台上使用 QNN SDK 进行模型推理。
1.1 环境准备
复制依赖库：将 QNN SDK 2.41 中的相关库文件复制到模型目录中的qai_libs目录并安装qai_appbuilder。
注意: Windows 平台需要复制对应的
.dll文件 (如 QnnHtp.dll, QnnSystem.dll 等) 到模型文件夹下。进入目录：在终端中
cd进入模型所在目录。
1.2 执行推理
使用 python 运行推理命令。
# ---------------------------------------------------------------------
# Copyright (c) 2024 Qualcomm Innovation Center, Inc. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# ---------------------------------------------------------------------
import sys
import os
import numpy as np
from transformers import AutoTokenizer
from pathlib import Path
# 引入 qai_appbuilder 相关库
sys.path.append(".")
sys.path.append("python")
from qai_appbuilder import (QNNContext, Runtime, LogLevel, ProfilingLevel, PerfProfile, QNNConfig)
####################################################################
# 配置部分
####################################################################
MODEL_NAME = "bge"
MAX_LENGTH = 512
execution_ws = Path(os.getcwd())
qnn_dir = os.path.join(execution_ws, "qai_libs")
model_path = execution_ws / "model.bin"
####################################################################
# 全局变量定义
####################################################################
bge_context = None
tokenizer = None
# 定义全局输入缓冲区 (Global Input Buffers)
# 作用：固定内存地址，确保多次推理时底层QNN能读到新数据
g_input_ids = None
g_attention_mask = None
g_token_type_ids = None
g_position_ids = None
class BGEModel(QNNContext):
def Inference(self, input_ids, attention_mask, token_type_ids, position_ids):
# 将全局缓冲区的引用传递给底层
input_datas = [input_ids, attention_mask, token_type_ids, position_ids]
output_datas = super().Inference(input_datas)
return output_datas[0]
def Init():
global bge_context, tokenizer
global g_input_ids, g_attention_mask, g_token_type_ids, g_position_ids
if not os.path.exists(model_path):
print(f"Error: Model file not found at {model_path}")
exit(1)
print("Initializing Tokenizer...")
try:
tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-base-zh-v1.5")
except Exception as e:
print(f"Failed to load tokenizer: {e}")
exit(1)
# --- 初始化全局固定缓冲区 (全零初始化) ---
print("Allocating Global Input Buffers...")
g_input_ids = np.ascontiguousarray(np.zeros((1, MAX_LENGTH), dtype=np.int32))
g_attention_mask = np.ascontiguousarray(np.zeros((1, MAX_LENGTH), dtype=np.int32))
g_token_type_ids = np.ascontiguousarray(np.zeros((1, MAX_LENGTH), dtype=np.int32))
g_position_ids = np.ascontiguousarray(np.arange(MAX_LENGTH, dtype=np.int32).reshape(1, MAX_LENGTH))
print("Initializing QNN Context...")
# 使用 HTP (DSP) 加速
QNNConfig.Config(qnn_dir, Runtime.HTP, LogLevel.WARN, ProfilingLevel.BASIC)
bge_context = BGEModel("bge_context", str(model_path))
def Preprocess(text):
"""
分词并将数据拷贝到全局缓冲区
"""
global g_input_ids, g_attention_mask, g_token_type_ids, g_position_ids
print(f"Preprocessing text: '{text}'")
inputs = tokenizer(
text,
padding="max_length",
truncation=True,
max_length=MAX_LENGTH,
return_tensors="np"
)
# --- 使用 np.copyto 更新全局缓冲区内容 ---
np.copyto(g_input_ids, inputs["input_ids"].astype(np.int32))
np.copyto(g_attention_mask, inputs["attention_mask"].astype(np.int32))
np.copyto(g_token_type_ids, inputs["token_type_ids"].astype(np.int32))
# 打印前5个ID以验证数据更新
print(f"Global Input IDs (first 5): {g_input_ids[0][:5]}")
return g_input_ids, g_attention_mask, g_token_type_ids, g_position_ids
def Postprocess(raw_output):
"""
后处理：Reshape -> 取[CLS] -> 归一化 -> 深拷贝
"""
# 必须拷贝一份数据，否则会被下一次推理覆盖
data = raw_output.astype(np.float32).copy()
try:
data = data.reshape(512, 768)
except ValueError:
print(f"Error: Output shape {data.shape} cannot be reshaped.")
return data
# 提取 [CLS] 向量
sentence_embedding = data[0]
# 归一化
norm = np.linalg.norm(sentence_embedding)
if norm > 1e-12:
sentence_embedding = sentence_embedding / norm
return sentence_embedding.copy()
def Inference(text_input):
global bge_context
# 1. 前处理 (写入全局缓冲区)
input_ids, attention_mask, token_type_ids, position_ids = Preprocess(text_input)
# 2. 设置性能模式 (提高DSP频率)
PerfProfile.SetPerfProfileGlobal(PerfProfile.BURST)
# 3. 执行推理
raw_output = bge_context.Inference(input_ids, attention_mask, token_type_ids, position_ids)
# 4. 释放性能模式
PerfProfile.RelPerfProfileGlobal()
# 5. 后处理
final_embedding = Postprocess(raw_output)
return final_embedding
def Release():
global bge_context
if bge_context:
del bge_context
print("QNN Context Released.")
# ================= 主程序 =================
if __name__ == "__main__":
try:
Init()
# --- 示例 1 ---
text_a = "这是一段测试文本"
print(f"\n[Case 1] Input: {text_a}")
emb_a = Inference(text_a)
print("Embedding (Top 10):", emb_a[:10])
print("Shape:", emb_a.shape)
# 保存为 Numpy 文件
np.save("embedding_a.npy", emb_a)
print("Saved to embedding_a.npy")
# --- 示例 2 ---
text_b = "今天天气不错"
print(f"\n[Case 2] Input: {text_b}")
emb_b = Inference(text_b)
print("Embedding (Top 10):", emb_b[:10])
print("Shape:", emb_b.shape)
# 保存为 Numpy 文件
np.save("embedding_b.npy", emb_b)
print("Saved to embedding_b.npy")
except Exception as e:
print(f"An error occurred: {e}")
import traceback
traceback.print_exc()
finally:
print("\nReleasing resources...")
Release()第二部分：Android 平台推理指南
本部分介绍如何在 Android 设备（通过 ADB）上进行模型推理。
2.1 快速验证
2.1.1环境准备
上传依赖库：将 QNN SDK 2.41 中的
aarch64-android库文件上传到设备（程序在2.41.0.251128\bin\aarch64-android）。# 示例：将库文件复制到设备目录 adb push libQnnHtp.so /data/local/tmp/bge_model/ adb push libQnnSystem.so /data/local/tmp/bge_model/ adb push libQnnHtpV73Skel.so /data/local/tmp/bge_model/ # ... 以及其他必要文件
设置环境变量：进入 ADB shell 并设置库路径。
cd /data/local/tmp/bge_model/ export LD_LIBRARY_PATH=$PWD export ADSP_LIBRARY_PATH=$PWD chmod +x qnn-net-run
2.1.2执行推理
在 ADB Shell 中运行以下命令：
./qnn-net-run \ --retrieve_context ./model.bin \ --backend ./libQnnHtp.so \ --input_list ./inputs/input_list.txt \ --use_native_input_files \ --config_file ./htp_backend.json
2.2 项目编译
注意: 整体的编译后端库流程请查看项目代码ai-engine-direct-helper-main\samples\android
2.2.1模型前处理和后处理
import android.content.res.AssetManager;
import java.io.InputStream;
import java.util.Arrays;
// ...
try {
// 1. 初始化 (加载词表)
AssetManager assets = getAssets();
InputStream vocabStream = assets.open("vocab.txt");
BgeProcessor processor = new BgeProcessor(vocabStream);
// 2. 前处理 (输入文本 -> 模型输入)
String text = "这是一段测试文本";
BgeProcessor.BgeInput input = processor.preprocess(text);
// 打印结果验证 (Int数组可以直接传给 TFLite 或 QNN)
// Log.d("BGE", "Input IDs: " + Arrays.toString(input.inputIds));
// Log.d("BGE", "Attention Mask: " + Arrays.toString(input.attentionMask));
// ---------------------------------------------------------
// 这里执行你的模型推理 (TFLite / ONNX / QNN)
// 假设 modelOutput 是推理得到的 float[] (大小通常是 768)
// float[] modelOutput = runInference(input.inputIds, ...);
// ---------------------------------------------------------
// 模拟一个推理结果用于测试
float[] mockOutput = new float[768];
Arrays.fill(mockOutput, 0.1f);
// 3. 后处理 (归一化)
float[] finalEmbedding = processor.postprocess(mockOutput);
// Log.d("BGE", "Normalized Embedding[0]: " + finalEmbedding[0]);
} catch (Exception e) {
e.printStackTrace();
}2.2.2模型调用推理
Java_com_example_DDColor_MainActivity_DDColor(...) {
float* inputBuffer = (float*)env->GetDirectBufferAddress(j_inputBuffer);
float* outputBuffer = (float*)env->GetDirectBufferAddress(j_outputBuffer);
// 1. 指定后端 (例如：libQnnHtp.so 表示在DSP上运行)
std::string backend_lib_path = libs_dir + "/libQnnHtp.so";
std::string system_lib_path = libs_dir + "/libQnnSystem.so";
// 2. 初始化模型
libAppBuilder.ModelInitialize(MODEL_NAME, model_path, backend_lib_path, ...);
// 3. 执行推理
libAppBuilder.ModelInference(MODEL_NAME, ...);
// 4. 拷贝结果
memcpy(outputBuffer, outputBuffers.at(0), outputSize[0]);
// ... 释放资源 ...
return 0;
}