zhoutao-shopee/rq-kmeans

Model Selection SID Server

基于 RQ-KMeans 模型的图像 SID 推理服务。

项目简介

本服务使用训练好的 rqkmeans_constrained 模型对图像 embedding 进行推理，生成对应的 SID（Semantic ID）。

功能特性

• 支持批量推理（batch processing）
• 基于 Sanic 框架的高性能异步服务
• 支持 Docker 容器化部署
• 健康检查端点
• 灵活的配置管理

项目结构

model-selection-sid-server/
├── app.py                 # 服务主文件
├── config.yaml            # 配置文件
├── requirements.txt       # Python 依赖
├── start.sh              # 启动脚本（自动激活虚拟环境）
├── test_api.py           # API 测试脚本
├── DEBUG.md              # 调试指南
├── Dockerfile            # Docker 镜像构建文件
├── docker-compose.yml    # Docker Compose 配置
├── README.md            # 项目说明文档
├── codebook/            # 码本文件目录
│   └── rqkmeans_4_256_5w.codebooks_constrained.npz
└── utils/               # 工具模块
    ├── __init__.py
    ├── model_utils.py    # 模型推理工具函数
    └── config_utils.py  # 配置管理工具函数

快速开始

1. 创建虚拟环境（推荐）

# 创建虚拟环境
python3 -m venv venv

# 激活虚拟环境
# macOS/Linux:
source venv/bin/activate
# Windows:
# venv\Scripts\activate

2. 安装依赖

pip install -r requirements.txt

3. 配置模型路径

码本文件已包含在项目的 codebook/ 目录中，配置文件已设置为正确的路径。如需修改，编辑 config.yaml 文件：

model:
  codebook_path: "./codebook/rqkmeans_4_256_5w.codebooks_constrained.npz"
  use_prefix: true

4. 启动服务

方式一：使用启动脚本（推荐）

# 直接运行启动脚本（会自动激活虚拟环境）
./start.sh

方式二：手动启动

# 1. 激活虚拟环境
source venv/bin/activate

# 2. 启动服务
python app.py

方式三：一行命令启动

source venv/bin/activate && python app.py

服务将在 http://0.0.0.0:80 启动。

提示：

• 启动脚本 start.sh 会自动检查虚拟环境和依赖，是最简单的方式
• 如果看到 (venv) 前缀，说明虚拟环境已激活
• 按 Ctrl+C 可以停止服务

API 文档

健康检查

GET /health

检查服务健康状态。

响应示例：

{
  "status": "healthy",
  "service": "Model_Selection_SID_Server",
  "model_loaded": true
}

批量预测 SID

POST /predict

批量预测图像 embedding 的 SID。

请求体：

{
  "items": [
    {
      "image_id": "img_001",
      "image_embedding": [0.1, 0.2, 0.3, ...]
    },
    {
      "image_id": "img_002",
      "image_embedding": [0.4, 0.5, 0.6, ...]
    }
  ]
}

响应示例：

{
  "status": "success",
  "results": [
    {
      "image_id": "img_001",
      "sid": ["a_1", "b_12", "c_23", "d_34"]
    },
    {
      "image_id": "img_002",
      "sid": ["a_2", "b_15", "c_28", "d_41"]
    }
  ]
}

错误响应：

{
  "status": "error",
  "message": "错误描述信息"
}

调试和测试

快速测试

使用提供的测试脚本快速测试服务：

# 激活虚拟环境
source venv/bin/activate

# 运行测试脚本
python test_api.py

测试脚本会自动测试：

• 健康检查接口
• 单个图像预测
• 批量预测
• 错误处理

查看日志

服务日志会输出到：

• 控制台：实时查看
• 日志文件：app.log

# 查看实时日志
tail -f app.log

# 查看最近的日志
tail -n 100 app.log

更多调试技巧请参考 DEBUG.md

使用示例

Python 示例

import requests
import json

url = "http://localhost:80/predict"

# 准备请求数据
data = {
    "items": [
        {
            "image_id": "img_001",
            "image_embedding": [0.1] * 768  # 假设 embedding 维度为 768
        },
        {
            "image_id": "img_002",
            "image_embedding": [0.2] * 768
        }
    ]
}

# 发送请求
response = requests.post(url, json=data)
result = response.json()

print(json.dumps(result, indent=2, ensure_ascii=False))

cURL 示例

curl -X POST http://localhost:80/predict \
  -H "Content-Type: application/json" \
  -d '{
    "items": [
      {
        "image_id": "img_001",
        "image_embedding": [0.1, 0.2, 0.3]
      }
    ]
  }'

Docker 部署

使用 Docker Compose（推荐）

注意： 使用 Docker 部署时，码本文件会自动挂载到容器中。如果需要自定义配置，可以使用 config.docker.yaml.example。

1. 启动服务：

   docker-compose up -d
   

   或者使用自定义配置：

   # 复制 Docker 配置示例
   cp config.docker.yaml.example config.yaml
   docker-compose up -d
   

使用 Docker 命令

# 构建镜像
docker build -t model-selection-sid-server .

# 运行容器
docker run -d \
  -p 80:80 \
  -v $(pwd)/config.yaml:/app/config.yaml:ro \
  -v $(pwd)/codebook:/app/codebook:ro \
  --name model-selection-sid-server \
  model-selection-sid-server

配置说明

config.yaml 配置文件包含以下选项：

• model.codebook_path: 码本文件路径（.npz 格式），默认指向项目内的 ./codebook/ 目录
• model.use_prefix: 是否使用前缀格式的 SID（如 a_1, b_12）
• server.host: 服务监听地址
• server.port: 服务监听端口
• server.max_batch_size: 最大批次大小

码本文件已包含在项目的 codebook/ 目录中，无需额外配置即可使用。

注意事项

1. 确保模型文件路径正确，服务启动时会自动加载码本
2. image_embedding 必须是数值数组，维度需要与训练时一致
3. 批次大小受 max_batch_size 配置限制，默认最大 1000
4. 服务会自动处理 NaN 和 Inf 值，将其替换为 0

故障排查

模型未加载

检查：

1. 码本文件路径是否正确
2. 文件是否存在且有读取权限
3. 查看日志文件 app.log 获取详细错误信息

请求失败

检查：

1. 请求体格式是否正确
2. image_embedding 是否为有效的数值数组
3. 批次大小是否超过限制

许可证

本项目遵循项目主仓库的许可证。