\u5728\u4fe1\u606f\u68c0\u7d22\u9886\u57df，\u6211\u4eec\u7ecf\u5e38\u9047\u5230\u8fd9\u6837\u7684\u56f0\u5883：\u641c\u7d22\u5f15\u64ce\u80fd\u591f\u627e\u5230\u5927\u91cf\u76f8\u5173\u6587\u6863，\u4f46\u6700\u91cd\u8981\u7684\u7ed3\u679c\u5f80\u5f80\u88ab\u57cb\u6ca1\u5728\u4e2d\u95f4\u4f4d\u7f6e\u3002\u4f20\u7edf\u7684\u5173\u952e\u8bcd\u5339\u914d\u548c\u5411\u91cf\u76f8\u4f3c\u5ea6\u8ba1\u7b97\u867d\u7136\u5feb\u901f，\u4f46\u7f3a\u4e4f\u5bf9\u8bed\u4e49\u6df1\u5c42\u6b21\u7406\u89e3\u7684\u80fd\u529b\u3002<\/p>\n
\u6587\u8109\u5b9a\u5e8f\u7cfb\u7edf\u6b63\u662f\u4e3a\u4e86\u89e3\u51b3\u8fd9\u4e2a"\u641c\u5f97\u5230\u4f46\u6392\u4e0d\u51c6"\u7684\u75db\u70b9\u800c\u8bbe\u8ba1\u7684\u3002\u5b83\u57fa\u4e8eBGE-Reranker-v2-m3\u6a21\u578b，\u901a\u8fc7\u5168\u4ea4\u53c9\u6ce8\u610f\u529b\u673a\u5236，\u5bf9\u68c0\u7d22\u7ed3\u679c\u8fdb\u884c\u667a\u80fd\u91cd\u6392\u5e8f，\u8ba9\u6700\u76f8\u5173\u7684\u5185\u5bb9\u8131\u9896\u800c\u51fa\u3002<\/p>\n
\u672c\u6559\u7a0b\u5c06\u624b\u628a\u624b\u6559\u4f60\u5982\u4f55\u5728NVIDIA Triton\u63a8\u7406\u670d\u52a1\u5668\u4e0a\u90e8\u7f72\u8fd9\u4e2a\u5f3a\u5927\u7684\u8bed\u4e49\u91cd\u6392\u5e8f\u7cfb\u7edf，\u4e3a\u4f60\u7684\u641c\u7d22\u5e94\u7528\u589e\u6dfb\u6700\u540e\u4e00\u516c\u91cc\u7684\u7cbe\u51c6\u5ea6\u3002<\/p>\n
\u5728\u5f00\u59cb\u90e8\u7f72\u524d，\u8bf7\u786e\u4fdd\u4f60\u7684\u7cfb\u7edf\u6ee1\u8db3\u4ee5\u4e0b\u8981\u6c42：<\/p>\n
# \u66f4\u65b0\u7cfb\u7edf\u5305
\nsudo apt update && sudo apt upgrade -y<\/p>\n
# \u5b89\u88c5Docker
\nsudo apt install docker.io -y
\nsudo systemctl enable docker
\nsudo systemctl start docker<\/p>\n
# \u5b89\u88c5NVIDIA Container Toolkit
\ndistribution=$(. \/etc\/os-release;echo $ID$VERSION_ID)
\ncurl -s -L https:\/\/nvidia.github.io\/nvidia-docker\/gpgkey | sudo apt-key add –
\ncurl -s -L https:\/\/nvidia.github.io\/nvidia-docker\/$distribution\/nvidia-docker.list | sudo tee \/etc\/apt\/sources.list.d\/nvidia-docker.list
\nsudo apt update && sudo apt install nvidia-container-toolkit -y
\nsudo nvidia-ctk runtime configure –runtime=docker
\nsudo systemctl restart docker<\/p>\n
# \u62c9\u53d6\u6700\u65b0\u7248\u672c\u7684Triton\u670d\u52a1\u5668\u955c\u50cf
\ndocker pull nvcr.io\/nvidia\/tritonserver:23.09-py3<\/p>\n
# \u521b\u5efa\u6a21\u578b\u5b58\u50a8\u76ee\u5f55
\nmkdir -p triton_models\/bge_reranker\/1<\/p>\n
\u9996\u5148\u6211\u4eec\u9700\u8981\u5c06HuggingFace\u683c\u5f0f\u7684\u6a21\u578b\u8f6c\u6362\u4e3aTriton\u53ef\u7528\u7684\u683c\u5f0f：<\/p>\n
# \u5b89\u88c5\u5fc5\u8981\u7684Python\u4f9d\u8d56
\npip install transformers torch tensorrt<\/p>\n
# \u521b\u5efa\u6a21\u578b\u8f6c\u6362\u811a\u672c
\ncat > convert_model.py << 'EOF'
\nimport torch
\nfrom transformers import AutoModelForSequenceClassification, AutoTokenizer
\nimport tensorrt as trt<\/p>\n
# \u52a0\u8f7d\u6a21\u578b\u548c\u5206\u8bcd\u5668
\nmodel_name = "BAAI\/bge-reranker-v2-m3"
\nmodel = AutoModelForSequenceClassification.from_pretrained(model_name)
\ntokenizer = AutoTokenizer.from_pretrained(model_name)<\/p>\n
# \u4fdd\u5b58\u4e3aTorchScript\u683c\u5f0f
\ndummy_input = tokenizer("query", "document", return_tensors="pt")
\ntraced_model = torch.jit.trace(model, (dummy_input["input_ids"], dummy_input["attention_mask"]))
\ntraced_model.save("bge_reranker.pt")<\/p>\n
print("\u6a21\u578b\u8f6c\u6362\u5b8c\u6210！")
\nEOF<\/p>\n
# \u8fd0\u884c\u8f6c\u6362\u811a\u672c
\npython convert_model.py<\/p>\n
\u5728triton_models\/bge_reranker\/config.pbtxt\u4e2d\u521b\u5efa\u914d\u7f6e\u6587\u4ef6：<\/p>\n
name: "bge_reranker"
\nplatform: "pytorch_libtorch"
\nmax_batch_size: 32<\/p>\n
input [
\n {
\n name: "input_ids"
\n data_type: TYPE_INT64
\n dims: [ -1 ]
\n },
\n {
\n name: "attention_mask"
\n data_type: TYPE_INT64
\n dims: [ -1 ]
\n }
\n]<\/p>\n
output [
\n {
\n name: "output"
\n data_type: TYPE_FP32
\n dims: [ 1 ]
\n }
\n]<\/p>\n
instance_group [
\n {
\n kind: KIND_GPU
\n count: 1
\n }
\n]<\/p>\n
optimization {
\n cuda {
\n graphs: true
\n }
\n}<\/p>\n
# \u542f\u52a8Triton\u670d\u52a1\u5668
\ndocker run -d –gpus=all –shm-size=1g –ulimit memlock=-1 \\\\
\n -p 8000:8000 -p 8001:8001 -p 8002:8002 \\\\
\n -v $(pwd)\/triton_models:\/models \\\\
\n nvcr.io\/nvidia\/tritonserver:23.09-py3 \\\\
\n tritonserver –model-repository=\/models<\/p>\n
# \u68c0\u67e5\u670d\u52a1\u5668\u72b6\u6001
\ncurl -v localhost:8000\/v2\/health\/ready<\/p>\n
# \u67e5\u770b\u5df2\u52a0\u8f7d\u6a21\u578b
\ncurl -v localhost:8000\/v2\/models\/list<\/p>\n
\u521b\u5efa\u4e00\u4e2a\u7b80\u5355\u7684\u5ba2\u6237\u7aef\u6765\u6d4b\u8bd5\u6a21\u578b：<\/p>\n
import tritonclient.http as httpclient
\nimport numpy as np
\nfrom transformers import AutoTokenizer<\/p>\n
class BGERerankerClient:
\n def __init__(self, url="localhost:8000"):
\n self.client = httpclient.InferenceServerClient(url=url)
\n self.tokenizer = AutoTokenizer.from_pretrained("BAAI\/bge-reranker-v2-m3")<\/p>\n
def rerank(self, query, documents):
\n """\u5bf9\u6587\u6863\u5217\u8868\u8fdb\u884c\u91cd\u6392\u5e8f"""
\n scores = []<\/p>\n
for doc in documents:
\n # \u51c6\u5907\u8f93\u5165\u6570\u636e
\n inputs = self.tokenizer(query, doc, return_tensors="pt", truncation=True)<\/p>\n
# \u521b\u5efaTriton\u8f93\u5165
\n input_ids = httpclient.InferInput(
\n "input_ids", inputs["input_ids"].shape, "INT64"
\n )
\n attention_mask = httpclient.InferInput(
\n "attention_mask", inputs["attention_mask"].shape, "INT64"
\n )<\/p>\n
input_ids.set_data_from_numpy(inputs["input_ids"].numpy().astype(np.int64))
\n attention_mask.set_data_from_numpy(inputs["attention_mask"].numpy().astype(np.int64))<\/p>\n
# \u6267\u884c\u63a8\u7406
\n result = self.client.infer(
\n model_name="bge_reranker",
\n inputs=[input_ids, attention_mask]
\n )<\/p>\n
score = result.as_numpy("output")[0]
\n scores.append(score)<\/p>\n
# \u6309\u5206\u6570\u6392\u5e8f\u6587\u6863
\n sorted_docs = [doc for _, doc in sorted(zip(scores, documents), reverse=True)]
\n return sorted_docs, sorted(scores, reverse=True)<\/p>\n
# \u4f7f\u7528\u793a\u4f8b
\nif __name__ == "__main__":
\n client = BGERerankerClient()<\/p>\n
query = "\u4eba\u5de5\u667a\u80fd\u7684\u53d1\u5c55\u5386\u53f2"
\n documents = [
\n "\u4eba\u5de5\u667a\u80fd\u4ece1956\u5e74\u8fbe\u7279\u8305\u65af\u4f1a\u8bae\u5f00\u59cb\u53d1\u5c55…",
\n "\u673a\u5668\u5b66\u4e60\u662f\u4eba\u5de5\u667a\u80fd\u7684\u4e00\u4e2a\u91cd\u8981\u5206\u652f…",
\n "\u6df1\u5ea6\u5b66\u4e60\u63a8\u52a8\u4e86\u4eba\u5de5\u667a\u80fd\u7684\u7b2c\u4e09\u6b21\u6d6a\u6f6e…",
\n "\u81ea\u7136\u8bed\u8a00\u5904\u7406\u662f\u4eba\u5de5\u667a\u80fd\u7684\u5173\u952e\u5e94\u7528\u9886\u57df…"
\n ]<\/p>\n
sorted_docs, scores = client.rerank(query, documents)
\n print("\u91cd\u6392\u5e8f\u7ed3\u679c：")
\n for i, (doc, score) in enumerate(zip(sorted_docs, scores)):
\n print(f"{i+1}. \u5206\u6570: {score:.4f} – {doc[:50]}…")<\/p>\n
\u5bf9\u4e8e\u5927\u91cf\u6587\u6863，\u5efa\u8bae\u4f7f\u7528\u6279\u91cf\u5904\u7406：<\/p>\n
def batch_rerank(self, query, documents, batch_size=16):
\n """\u6279\u91cf\u91cd\u6392\u5e8f\u6587\u6863"""
\n all_scores = []<\/p>\n
for i in range(0, len(documents), batch_size):
\n batch_docs = documents[i:i+batch_size]
\n batch_scores = []<\/p>\n
for doc in batch_docs:
\n inputs = self.tokenizer(query, doc, return_tensors="pt", truncation=True)
\n # … \u540c\u6837\u7684\u63a8\u7406\u903b\u8f91
\n batch_scores.append(score)<\/p>\n
all_scores.extend(batch_scores)<\/p>\n
return sorted(zip(documents, all_scores), key=lambda x: x[1], reverse=True)<\/p>\n
# \u4f7f\u7528FP16\u7cbe\u5ea6\u52a0\u901f\u63a8\u7406
\ndocker run … –env TF_ENABLE_AUTO_MIXED_PRECISION=1 …<\/p>\n
# \u542f\u7528\u52a8\u6001\u6279\u5904\u7406
\n# \u5728config.pbtxt\u4e2d\u6dfb\u52a0：
\ndynamic_batching {
\n preferred_batch_size: [ 4, 8, 16, 32 ]
\n max_queue_delay_microseconds: 1000
\n}<\/p>\n
# \u76d1\u63a7GPU\u5185\u5b58\u4f7f\u7528
\nnvidia-smi -l 1<\/p>\n
# \u8bbe\u7f6e\u5408\u9002\u7684\u6279\u5904\u7406\u5927\u5c0f
\n# \u6839\u636e\u4f60\u7684GPU\u5185\u5b58\u8c03\u6574：
\n# 8GB\u663e\u5b58: max_batch_size=16
\n# 16GB\u663e\u5b58: max_batch_size=32
\n# 24GB+\u663e\u5b58: max_batch_size=64<\/p>\n
\u95ee\u98981：\u6a21\u578b\u52a0\u8f7d\u5931\u8d25<\/p>\n
# \u68c0\u67e5\u6a21\u578b\u683c\u5f0f\u662f\u5426\u6b63\u786e
\nfile triton_models\/bge_reranker\/1\/model.pt<\/p>\n
# \u68c0\u67e5\u914d\u7f6e\u6587\u4ef6\u8bed\u6cd5
\ncd triton_models\/bge_reranker && find . -name "*.pbtxt" -exec echo "\u68c0\u67e5\u6587\u4ef6:" {} \\\\; -exec cat {} \\\\;<\/p>\n
\u95ee\u98982：GPU\u5185\u5b58\u4e0d\u8db3<\/p>\n
# \u51cf\u5c11\u6279\u5904\u7406\u5927\u5c0f
\n# \u5728config.pbtxt\u4e2d\u5c06max_batch_size\u8c03\u6574\u4e3a\u66f4\u5c0f\u7684\u503c<\/p>\n
# \u542f\u7528\u5185\u5b58\u4f18\u5316
\noptimization {
\n execution_accelerators {
\n gpu_execution_accelerator : [ { name : "auto_mixed_precision" } ]
\n }
\n}<\/p>\n
\u901a\u8fc7\u672c\u6559\u7a0b，\u4f60\u5df2\u7ecf\u6210\u529f\u5728NVIDIA Triton\u63a8\u7406\u670d\u52a1\u5668\u4e0a\u90e8\u7f72\u4e86BGE-Reranker-v2-m3\u6a21\u578b，\u6784\u5efa\u4e86\u4e00\u4e2a\u9ad8\u6027\u80fd\u7684\u8bed\u4e49\u91cd\u6392\u5e8f\u7cfb\u7edf\u3002\u8fd9\u4e2a\u7cfb\u7edf\u80fd\u591f：<\/p>\n
\u5728\u5b9e\u9645\u5e94\u7528\u4e2d，\u4f60\u53ef\u4ee5\u5c06\u8fd9\u4e2a\u91cd\u6392\u5e8f\u7cfb\u7edf\u96c6\u6210\u5230\u641c\u7d22\u5f15\u64ce\u3001\u77e5\u8bc6\u5e93\u7cfb\u7edf\u6216\u63a8\u8350\u7cfb\u7edf\u4e2d，\u663e\u8457\u63d0\u5347\u6700\u7ec8\u7528\u6237\u7684\u641c\u7d22\u4f53\u9a8c\u3002\u8bb0\u5f97\u6839\u636e\u5b9e\u9645\u4f7f\u7528\u60c5\u51b5\u8c03\u6574\u6279\u5904\u7406\u5927\u5c0f\u548c\u6a21\u578b\u914d\u7f6e，\u4ee5\u8fbe\u5230\u6700\u4f73\u7684\u6027\u80fd\u6548\u679c\u3002<\/p>\n
\u83b7\u53d6\u66f4\u591aAI\u955c\u50cf<\/p>\n
\u60f3\u63a2\u7d22\u66f4\u591aAI\u955c\u50cf\u548c\u5e94\u7528\u573a\u666f？\u8bbf\u95ee CSDN\u661f\u56fe\u955c\u50cf\u5e7f\u573a，\u63d0\u4f9b\u4e30\u5bcc\u7684\u9884\u7f6e\u955c\u50cf，\u8986\u76d6\u5927\u6a21\u578b\u63a8\u7406\u3001\u56fe\u50cf\u751f\u6210\u3001\u89c6\u9891\u751f\u6210\u3001\u6a21\u578b\u5fae\u8c03\u7b49\u591a\u4e2a\u9886\u57df，\u652f\u6301\u4e00\u952e\u90e8\u7f72\u3002<\/p>\n","protected":false},"excerpt":{"rendered":"
\u6587\u8109\u5b9a\u5e8f\u90e8\u7f72\u6559\u7a0b：NVIDIA Triton\u63a8\u7406\u670d\u52a1\u5668\u96c6\u6210BGE-Reranker-v2-m3
\n1. \u5f15\u8a00：\u4e3a\u4ec0\u4e48\u9700\u8981\u667a\u80fd\u8bed\u4e49\u91cd\u6392\u5e8f？
\n\u5728\u4fe1\u606f\u68c0\u7d22\u9886\u57df，\u6211\u4eec\u7ecf\u5e38\u9047\u5230\u8fd9\u6837\u7684\u56f0\u5883：\u641c\u7d22\u5f15\u64ce\u80fd\u591f\u627e\u5230\u5927\u91cf\u76f8\u5173\u6587\u6863，\u4f46\u6700\u91cd\u8981\u7684\u7ed3\u679c\u5f80\u5f80\u88ab\u57cb\u6ca1\u5728\u4e2d\u95f4\u4f4d\u7f6e\u3002\u4f20\u7edf\u7684\u5173\u952e\u8bcd\u5339\u914d\u548c\u5411\u91cf\u76f8\u4f3c\u5ea6\u8ba1\u7b97\u867d\u7136\u5feb\u901f，\u4f46\u7f3a\u4e4f\u5bf9\u8bed\u4e49\u6df1\u5c42\u6b21\u7406\u89e3\u7684\u80fd\u529b\u3002
\n\u6587\u8109\u5b9a\u5e8f\u7cfb\u7edf\u6b63\u662f\u4e3a\u4e86\u89e3\u51b3\u8fd9\u4e2a\\”\u641c\u5f97\u5230\u4f46\u6392\u4e0d\u51c6\\”\u7684\u75db\u70b9\u800c\u8bbe\u8ba1\u7684<\/p>\n","protected":false},"author":2,"featured_media":0,"comment_status":"open","ping_status":"open","sticky":false,"template":"","format":"standard","meta":{"footnotes":""},"categories":[1],"tags":[9282,9317,9316],"topic":[],"class_list":["post-83298","post","type-post","status-publish","format-standard","hentry","category-server","tag-nvidia-triton","tag-9317","tag-9316"],"yoast_head":"\n