使用Infinity部署Embedding和Reranking模型

shizidushu發表於2024-08-06
模型

使用Infinity部署Embedding和Reranking模型

說明：

  • 首次發表日期：2024-08-06
  • Infinity Github 倉庫： https://github.com/michaelfeil/infinity
  • Infinity 官方文件： https://michaelfeil.github.io/infinity/

下載權重

pip install -U "huggingface_hub[cli]"

export HF_ENDPOINT=https://hf-mirror.com

huggingface-cli download BAAI/bge-m3
huggingface-cli download BAAI/bge-reranker-v2-m3

使用 hf-mirror.com 的映象，下載模型權重

模型權重預設快取到了$HOME/.cache下，使用tree命令檢查下：

~/.cache$ tree -L 3 .
.
├── huggingface
│   └── hub
│       ├── huggingface
│       ├── models--BAAI--bge-m3
│       ├── models--BAAI--bge-reranker-v2-m3
│       └── version.txt

執行Docker容器

docker run 啟動infinity容器:

docker run -it --gpus all \
 --env HF_ENDPOINT=https://hf-mirror.com \
 -v $HOME/.cache:/app/.cache \
 -p 7997:7997 \
 michaelf34/infinity:latest \
 v2 \
 --model-id BAAI/bge-m3 \
 --model-id BAAI/bge-reranker-v2-m3 \
 --port 7997

其中

  • -it: 互動式執行容器，分配一個偽終端。（如果要在後臺執行，將-it改為-d）
  • -v $HOME/.cache:/app/.cache：將本地下載的模型權重對映到Docker容器中了

啟動完成後，可以看到以下日誌：

INFO     2024-08-06 03:52:57,665 infinity_emb INFO:

         ♾️  Infinity - Embedding Inference Server                           
         MIT License; Copyright (c) 2023 Michael Feil                        
         Version 0.0.53
                                                      
         Open the Docs via Swagger UI:                                       
         http://0.0.0.0:7997/docs
                                                 
         Access model via 'GET':                                             
         curl http://0.0.0.0:7997/models
                                 
INFO:     Application startup complete.
INFO:     Uvicorn running on http://0.0.0.0:7997 (Press CTRL+C to quit)

開啟 http://xx.xx.xx.xx:7997/docs，可以看到 API 文件。（其中 xx.xx.xx.xx代指 IP 地址）

呼叫測試

Embedding模型：

curl --location 'http://xx.xx.xx.xx:7997/embeddings' \
--header 'Content-Type: application/json' \
--data '{
    "input": "喝水吃餅乾",
    "model": "BAAI/bge-m3"
  }'

Reranking模型：

curl -X 'POST' \
  'http://xx.xx.xx.xx:7997/rerank' \
  -H 'accept: application/json' \
  -H 'Content-Type: application/json' \
  -d '{
  "query": "like it",
  "documents": [
    "like", "hate", "sky"
  ],
  "return_documents": false,
  "model": "BAAI/bge-reranker-v2-m3"
}'

相關文章