好的，为了尽可能减小镜像体积，并且满足您“不用pytorch”的要求（因为PyTorch确实占用空间较大），我们将采用将模型转换为ONNX格式并在ONNX Runtime中运行它的策略。这样可以在最终的运行环境中避免庞大的PyTorch库。
这个过程会稍微复杂一些，因为它涉及到：
- 模型转换阶段：在构建镜像的过程中，我们会使用
sentence-transformers和torch（仅在此阶段）来加载原始模型，并将其转换为ONNX格式。 - 运行阶段：最终的镜像将只包含ONNX Runtime和模型运行所需的最小依赖，不再包含
torch。
这需要修改您的sentence.py以适应ONNX模型的加载和推理。
以下是建议的方案：
文件结构：
.
├── Dockerfile
├── build.sh
├── requirements_builder.txt # 用于构建阶段（模型转换）的依赖
├── requirements_runtime.txt # 用于运行阶段的依赖
├── convert_to_onnx.py # 模型转换脚本
└── sentence_onnx.py # 修改后使用ONNX模型的应用脚本 (替换原 sentence.py)
1. requirements_builder.txt (用于模型转换)
Plaintext
sentence-transformers==2.7.0
torch>=2.0.0 # 需要torch来加载原始模型进行转换 (CPU版本即可)
optimum[onnxruntime]>=1.16.0 # Hugging Face Optimum库，简化ONNX转换和使用
transformers>=4.30.0 # sentence-transformers的依赖，也用于tokenizer
# fastapi, uvicorn, pydantic, numpy 将在runtime中安装，这里不需要
- 我们在这里指定
torch，因为它是在构建阶段将模型转换为ONNX所必需的。最终的运行时镜像中不会包含它。 optimum库可以很方便地将Hugging Face Transformers模型（sentence-transformers底层使用的模型）导出为ONNX格式。
2. requirements_runtime.txt (用于最终运行环境)
Plaintext
fastapi==0.110.0
uvicorn[standard]==0.29.0
pydantic>=2.0
numpy>=1.20
optimum[onnxruntime]>=1.16.0 # 用于加载和运行ONNX模型
# 注意：这里不再有 torch 和 sentence-transformers (原始库)
# transformers 会作为 optimum 的依赖被安装，主要用于加载tokenizer
- 这里只包含运行FastAPI服务和ONNX模型所需的最小依赖。
onnxruntime本身比torch小得多。
3. convert_to_onnx.py (模型转换脚本)
Python
import sys
from pathlib import Path
from optimum.onnxruntime import ORTModelForFeatureExtraction
from transformers import AutoTokenizer
def convert_sbert_to_onnx(model_name_or_path: str, output_dir: str):
"""
Converts a SentenceTransformer model to ONNX format using Optimum.
The underlying Hugging Face model name is typically 'sentence-transformers/' + model_name_or_path.
"""
hf_model_name = f"sentence-transformers/{model_name_or_path}"
output_path = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)
print(f"Loading Hugging Face model '{hf_model_name}' for ONNX conversion...")
# 1. Export the model to ONNX
# ORTModelForFeatureExtraction will handle the conversion from a PyTorch checkpoint.
# It needs the underlying Hugging Face compatible model, not the SentenceTransformer wrapper directly.
ort_model = ORTModelForFeatureExtraction.from_pretrained(hf_model_name, export=True)
# 2. Save the ONNX model and its tokenizer
ort_model.save_pretrained(output_path)
# Tokenizer is usually saved by the above command as well, but ensure it explicitly.
# SentenceTransformer models have their tokenizers accessible. For ONNX, we typically
# save the tokenizer from the underlying Hugging Face model.
tokenizer = AutoTokenizer.from_pretrained(hf_model_name)
tokenizer.save_pretrained(output_path)
print(f"ONNX model and tokenizer saved to {output_path}")
if __name__ == "__main__":
if len(sys.argv) < 3:
print("Usage: python convert_to_onnx.py <model_name> <output_directory>")
print("Example: python convert_to_onnx.py paraphrase-multilingual-MiniLM-L12-v2 /app/onnx_model")
sys.exit(1)
model_name = sys.argv[1]
output_directory = sys.argv[2]
convert_sbert_to_onnx(model_name, output_directory)
4. sentence_onnx.py (修改后的应用脚本，替换原 sentence.py)
Python
from fastapi import FastAPI
from pydantic import BaseModel
import numpy as np
from optimum.onnxruntime import ORTModelForFeatureExtraction
from transformers import AutoTokenizer # To load the tokenizer
app = FastAPI()
# 定义模型和tokenizer的路径 (在Dockerfile中会被复制到这个路径)
MODEL_DIR = "/app/onnx_model"
# 加载 ONNX 模型和分词器
print(f"Loading ONNX model and tokenizer from {MODEL_DIR}...")
try:
tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
model = ORTModelForFeatureExtraction.from_pretrained(MODEL_DIR) # Optimum will use ONNX Runtime
print("ONNX model and tokenizer loaded successfully.")
except Exception as e:
print(f"Error loading ONNX model or tokenizer: {e}")
# 在这种情况下，服务可能无法正常工作，后续请求会失败
# 考虑更健壮的错误处理或健康检查端点
tokenizer = None
model = None
class TextPayload(BaseModel):
texts: str # 原始代码期望单个字符串
def mean_pooling(model_output_last_hidden_state: np.ndarray, attention_mask: np.ndarray) -> np.ndarray:
"""
Performs mean pooling on the token embeddings using the attention mask.
model_output_last_hidden_state: (batch_size, sequence_length, hidden_size)
attention_mask: (batch_size, sequence_length)
"""
token_embeddings = model_output_last_hidden_state
input_mask_expanded = np.expand_dims(attention_mask, axis=-1).astype(float)
sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
sum_mask = np.sum(input_mask_expanded, axis=1)
sum_mask = np.maximum(sum_mask, 1e-9) # 防止除以零
return sum_embeddings / sum_mask
@app.post("/embed")
async def embed(payload: TextPayload):
if not tokenizer or not model:
return {"error": "Model not loaded properly"}, 500
# 1. 使用分词器处理输入文本
# ONNX Runtime 需要 NumPy 数组作为输入
inputs = tokenizer(payload.texts, return_tensors="np", padding=True, truncation=True)
# 2. ONNX 模型推理
# optimum的ORTModelForFeatureExtraction的输出是一个类似字典的对象，
# 通常包含'last_hidden_state'。
outputs = model(**inputs)
last_hidden_state = outputs.last_hidden_state # (batch_size, seq_len, hidden_dim)
# 3. 应用 Mean Pooling 以获得句子嵌入
# 'paraphrase-multilingual-MiniLM-L12-v2' 使用 mean pooling
sentence_embedding_np = mean_pooling(last_hidden_state, inputs['attention_mask'])
# 因为输入 payload.texts 是单个字符串, batch_size 是 1.
# 我们返回单个嵌入向量。
return {"embeddings": sentence_embedding_np[0].tolist()}
if __name__ == "__main__":
import uvicorn
# 确保模型已加载再启动服务
if tokenizer and model:
uvicorn.run(app, host="0.0.0.0", port=5005)
else:
print("Failed to load model, Uvicorn server will not start.")
- 重要：
sentence-transformers库本身对PyTorch有很强的依赖性来加载和运行模型。为了完全移除PyTorch，我们不能直接使用SentenceTransformer类来加载ONNX模型。相反，我们使用optimum库中的ORTModelForFeatureExtraction和Hugging Face的AutoTokenizer来加载和运行已转换的ONNX模型和相应的分词器。 - 句子嵌入通常是通过对模型输出的词元嵌入（token embeddings）进行池化（pooling）操作得到的。对于
paraphrase-multilingual-MiniLM-L12-v2这类SBERT模型，常用的是均值池化（mean pooling）。上面的代码中添加了mean_pooling函数来实现这一点。
5. Dockerfile (优化后)
Dockerfile
#-----------------------------------------------------------------------------
# Stage 1: Builder - Convert model to ONNX and install build dependencies
#-----------------------------------------------------------------------------
FROM python:3.12-slim-bullseye as builder
WORKDIR /app
# 设置环境变量，避免生成 .pyc 文件，并使Python输出不缓冲
ENV PYTHONDONTWRITEBYTECODE=1
ENV PYTHONUNBUFFERED=1
# 安装构建阶段的Python依赖 (包括 torch, sentence-transformers, optimum 用于转换)
COPY requirements_builder.txt .
RUN pip install --no-cache-dir -r requirements_builder.txt \
--index-url https://pypi.tuna.tsinghua.edu.cn/simple \
--trusted-host pypi.tuna.tsinghua.edu.cn
# 复制转换脚本并执行模型转换
COPY convert_to_onnx.py .
# 将 'paraphrase-multilingual-MiniLM-L12-v2' 模型转换为ONNX格式并存放到 /app/onnx_model
# 这个模型名称与您原 sentence.py 中使用的一致
RUN python convert_to_onnx.py paraphrase-multilingual-MiniLM-L12-v2 /app/onnx_model
#-----------------------------------------------------------------------------
# Stage 2: Final - Create the minimal runtime image
#-----------------------------------------------------------------------------
FROM python:3.12-slim-bullseye as final
WORKDIR /app
ENV PYTHONDONTWRITEBYTECODE=1
ENV PYTHONUNBUFFERED=1
# 创建一个非root用户来运行应用，增强安全性
ARG UID=10001
RUN useradd \
--uid ${UID} \
--gid 0 \
--no-create-home \
--home /nonexistent \
--shell /bin/false \
appuser
# RUN groupadd --gid ${UID} appuser && \
# useradd --uid ${UID} --gid ${UID} --no-create-home --home /nonexistent --shell /sbin/nologin appuser
# 安装运行阶段的Python依赖 (fastapi, uvicorn, onnxruntime, optimum)
# 注意：这里不包含 torch 和原始的 sentence-transformers
COPY requirements_runtime.txt .
# 创建虚拟环境（可选，但有助于隔离）
# ENV VIRTUAL_ENV=/opt/venv
# RUN python3 -m venv $VIRTUAL_ENV
# ENV PATH="$VIRTUAL_ENV/bin:$PATH"
# RUN $VIRTUAL_ENV/bin/pip install --no-cache-dir -r requirements_runtime.txt \
RUN pip install --no-cache-dir --user -r requirements_runtime.txt \
--index-url https://pypi.tuna.tsinghua.edu.cn/simple \
--trusted-host pypi.tuna.tsinghua.edu.cn
# 从构建阶段复制转换好的ONNX模型和tokenizer
COPY --from=builder --chown=appuser:0 /app/onnx_model /app/onnx_model
# 复制修改后的应用代码
COPY --chown=appuser:0 sentence_onnx.py ./sentence.py
# 设置appuser的环境变量，使其能找到安装的包
# 如果使用 --user 安装pip包，它们通常在 /home/appuser/.local/bin 或 /root/.local/bin (取决于USER指令前的用户)
# 我们在切换到appuser之前以root身份安装pip包到用户目录，然后调整PATH
# 或者，更好的方式是系统级安装或虚拟环境，并确保appuser有权访问。
# 为简单起见，如果pip install --user是以root身份执行（默认），则路径是/root/.local
# 如果在USER appuser之后执行，则是/home/appuser/.local
# 我们在USER appuser之前执行，所以是/root/.local/bin
ENV PATH=/root/.local/bin:$PATH
# 切换到非root用户
USER appuser
EXPOSE 5005
# 健康检查 (可选, 但推荐)
# HEALTHCHECK --interval=30s --timeout=5s --start-period=15s --retries=3 \
# CMD curl -f http://localhost:5005/health || exit 1 # 假设你添加了一个 /health 端点
CMD ["python", "sentence.py"]
6. build.sh (基本不变，只需确保文件名对应)
您的build.sh脚本看起来是合理的，它会构建并推送Docker镜像。确保Dockerfile的路径正确。
优化说明：
- 多阶段构建：
builder阶段：安装包括torch在内的所有转换工具，下载原始模型，并将其转换为ONNX格式。这个阶段会比较大，但其内容不会进入最终镜像。final阶段：从一个干净的python:3.12-slim基础镜像开始，仅安装运行ONNX模型和FastAPI服务所需的最小依赖（如onnxruntime,optimum,fastapi,uvicorn）。然后从builder阶段复制转换后的ONNX模型文件。
- ONNX Runtime：
onnxruntime库远小于pytorch，这将显著减小最终镜像的体积。模型文件（.onnx）本身的大小与原始PyTorch模型权重文件大小相似（paraphrase-multilingual-MiniLM-L12-v2大约400-500MB），这部分是不可避免的，因为模型参数就是这么多。但我们节省了运行时库的空间。 - 依赖最小化：
requirements_runtime.txt只包含绝对必要的包。 - Python版本：使用
python:3.12-slim-bullseye，这是一个相对较小的官方镜像。 - 非Root用户：在最终阶段使用非root用户运行应用，这是一个安全最佳实践。
- 缓存：
pip install --no-cache-dir避免了pip缓存的写入，有助于减小层的大小。 - 国内镜像源：继续使用清华镜像源加速Python包的下载。
使用方法：
- 将上述内容的
Dockerfile,requirements_builder.txt,requirements_runtime.txt,convert_to_onnx.py, 和sentence_onnx.py文件放到您的项目根目录。 - 确保您的
build.sh脚本可以正确找到Dockerfile。 - 运行
./build.sh docker。
这将首先构建包含PyTorch的builder阶段来转换模型，然后构建一个不含PyTorch但包含ONNX模型和ONNX Runtime的轻量级final镜像。
这种方法应该能显著减小您的镜像体积，同时保留了使用强大转换器模型的能力。请注意，模型文件本身（约450MB for paraphrase-multilingual-MiniLM-L12-v2）仍将是镜像大小的主要组成部分，但您已经从运行时环境中移除了庞大的PyTorch库。