EmbeddingGemma-300m REST API封装教程:快速构建微服务

1. 引言

你是不是曾经遇到过这样的场景:手头有一个很棒的Embedding模型,想要在多个项目中复用,却每次都要重新配置环境、写调用代码?或者团队里不同成员需要用到相同的文本嵌入能力,但每个人都要单独部署一套环境?

EmbeddingGemma-300m作为Google推出的轻量级嵌入模型,只有3亿参数却能在多语言文本表示任务中表现出色。但直接通过Ollama命令行调用显然不适合生产环境。今天我就来分享如何用FastAPI将它封装成RESTful微服务,让你的嵌入能力随时可用、随处可调。

学完这篇教程,你将能够快速搭建一个高性能的嵌入服务,支持并发处理、负载均衡,并且可以通过Docker一键部署。无论你是想为内部系统增加语义搜索能力,还是构建面向用户的AI应用,这个方案都能帮你省去大量重复工作。

2. 环境准备与快速部署

2.1 系统要求与依赖安装

首先确保你的系统已经安装了Python 3.8+和Docker。然后创建项目目录并安装必要依赖:

mkdir embedding-service && cd embedding-service
python -m venv venv
source venv/bin/activate  # Linux/Mac
# 或者 venv\Scripts\activate  # Windows

pip install fastapi uvicorn ollama python-multipart

Ollama是运行EmbeddingGemma的前提,如果你还没有安装,可以用以下命令快速安装:

# Linux/Mac安装命令
curl -fsSL https://ollama.ai/install.sh | sh

# 然后拉取EmbeddingGemma模型
ollama pull embeddinggemma:300m

2.2 验证模型运行

在开始封装API之前,先确认模型能正常工作:

import ollama

response = ollama.embed(
    model='embeddinggemma:300m',
    input='测试文本嵌入功能'
)
print(f'嵌入向量长度: {len(response.embeddings[0])}')

如果看到输出"嵌入向量长度: 768",说明模型已经准备就绪。

3. FastAPI服务基础搭建

3.1 创建基础API结构

新建一个main.py文件,开始构建我们的嵌入服务:

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List
import ollama

app = FastAPI(title="EmbeddingGemma API", version="1.0.0")

class EmbeddingRequest(BaseModel):
    text: str

class BatchEmbeddingRequest(BaseModel):
    texts: List[str]

@app.get("/")
async def root():
    return {"message": "EmbeddingGemma API服务正常运行"}

@app.post("/embed")
async def get_embedding(request: EmbeddingRequest):
    try:
        response = ollama.embed(
            model='embeddinggemma:300m',
            input=request.text
        )
        return {"embedding": response.embeddings[0]}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"嵌入处理失败: {str(e)}")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

这个基础版本已经可以提供单个文本的嵌入服务了。启动服务试试:

python main.py

用curl测试一下:

curl -X POST "http://localhost:8000/embed" \
-H "Content-Type: application/json" \
-d '{"text": "为什么天空是蓝色的?"}'

你应该能收到一个768维的嵌入向量。

4. 生产级功能增强

4.1 批量处理支持

在实际应用中,单条处理效率太低。我们来添加批量处理能力:

@app.post("/embed/batch")
async def get_batch_embedding(request: BatchEmbeddingRequest):
    try:
        response = ollama.embed(
            model='embeddinggemma:300m',
            input=request.texts
        )
        return {"embeddings": response.embeddings}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"批量嵌入处理失败: {str(e)}")

批量处理能显著提升效率,特别是在处理大量文本时。

4.2 并发处理与性能优化

FastAPI天生支持异步,但Ollama的调用是同步的。我们可以用线程池来避免阻塞:

from concurrent.futures import ThreadPoolExecutor
import asyncio

executor = ThreadPoolExecutor(max_workers=4)

@app.post("/embed/batch")
async def get_batch_embedding(request: BatchEmbeddingRequest):
    try:
        loop = asyncio.get_event_loop()
        response = await loop.run_in_executor(
            executor,
            lambda: ollama.embed(model='embeddinggemma:300m', input=request.texts)
        )
        return {"embeddings": response.embeddings}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"批量嵌入处理失败: {str(e)}")

4.3 健康检查与监控

生产环境需要健康检查端点:

@app.get("/health")
async def health_check():
    try:
        # 简单查询验证服务正常
        response = ollama.embed(model='embeddinggemma:300m', input="health check")
        return {"status": "healthy", "model_ready": True}
    except Exception:
        return {"status": "unhealthy", "model_ready": False}, 503

5. Docker容器化部署

5.1 创建Dockerfile

为了便于部署,我们创建Dockerfile:

FROM python:3.9-slim

WORKDIR /app

# 安装Ollama
RUN apt-get update && apt-get install -y curl && \
    curl -fsSL https://ollama.ai/install.sh | sh

# 复制应用代码
COPY requirements.txt .
RUN pip install -r requirements.txt

COPY . .

# 下载模型(可以在运行时下载,但预先下载更好)
RUN ollama pull embeddinggemma:300m

EXPOSE 8000

CMD ["sh", "-c", "ollama serve & uvicorn main:app --host 0.0.0.0 --port 8000"]

5.2 编写docker-compose配置

对于更复杂的部署场景,使用docker-compose:

version: '3.8'

services:
  embedding-service:
    build: .
    ports:
      - "8000:8000"
    environment:
      - OLLAMA_HOST=0.0.0.0:11434
    deploy:
      resources:
        limits:
          memory: 2G
        reservations:
          memory: 1G

5.3 构建和运行

docker build -t embedding-service .
docker run -p 8000:8000 embedding-service

或者使用docker-compose:

docker-compose up -d

6. 高级功能与优化建议

6.1 缓存层集成

对于重复的查询,添加缓存可以大幅提升性能:

from functools import lru_cache

@lru_cache(maxsize=1000)
def get_cached_embedding(text: str):
    response = ollama.embed(model='embeddinggemma:300m', input=text)
    return response.embeddings[0]

@app.post("/embed/cached")
async def get_cached_embedding_endpoint(request: EmbeddingRequest):
    try:
        embedding = get_cached_embedding(request.text)
        return {"embedding": embedding}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"嵌入处理失败: {str(e)}")

6.2 速率限制

防止API被滥用,添加简单的速率限制:

from fastapi import Request
from datetime import datetime, timedelta

class RateLimiter:
    def __init__(self, max_requests: int, time_window: int):
        self.max_requests = max_requests
        self.time_window = time_window
        self.requests = {}

    def is_allowed(self, client_ip: str) -> bool:
        now = datetime.now()
        if client_ip not in self.requests:
            self.requests[client_ip] = []
        
        # 清理过期的请求记录
        self.requests[client_ip] = [
            req_time for req_time in self.requests[client_ip]
            if now - req_time < timedelta(seconds=self.time_window)
        ]
        
        if len(self.requests[client_ip]) < self.max_requests:
            self.requests[client_ip].append(now)
            return True
        return False

rate_limiter = RateLimiter(max_requests=100, time_window=60)

@app.middleware("http")
async def rate_limit_middleware(request: Request, call_next):
    client_ip = request.client.host
    if not rate_limiter.is_allowed(client_ip):
        return JSONResponse(
            status_code=429,
            content={"detail": "请求过于频繁,请稍后再试"}
        )
    return await call_next(request)

6.3 负载均衡方案

对于高并发场景,可以考虑部署多个实例并使用负载均衡器:

# docker-compose-scale.yml
version: '3.8'

services:
  embedding-service:
    build: .
    ports:
      - "8000"
    environment:
      - OLLAMA_HOST=0.0.0.0:11434
    deploy:
      replicas: 3
      resources:
        limits:
          memory: 2G

  nginx:
    image: nginx:alpine
    ports:
      - "80:80"
    volumes:
      - ./nginx.conf:/etc/nginx/nginx.conf
    depends_on:
      - embedding-service

对应的nginx配置:

events {
    worker_connections 1024;
}

http {
    upstream embedding_services {
        server embedding-service_1:8000;
        server embedding-service_2:8000;
        server embedding-service_3:8000;
    }

    server {
        listen 80;
        
        location / {
            proxy_pass http://embedding_services;
            proxy_set_header Host $host;
            proxy_set_header X-Real-IP $remote_addr;
        }
    }
}

7. 实际使用示例

7.1 Python客户端示例

import requests
import json

class EmbeddingClient:
    def __init__(self, base_url="http://localhost:8000"):
        self.base_url = base_url
    
    def get_embedding(self, text):
        response = requests.post(
            f"{self.base_url}/embed",
            json={"text": text}
        )
        return response.json()["embedding"]
    
    def get_batch_embedding(self, texts):
        response = requests.post(
            f"{self.base_url}/embed/batch",
            json={"texts": texts}
        )
        return response.json()["embeddings"]

# 使用示例
client = EmbeddingClient()
embedding = client.get_embedding("自然语言处理很有趣")
print(f"嵌入向量维度: {len(embedding)}")

7.2 简单语义搜索实现

import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

class SemanticSearch:
    def __init__(self, client):
        self.client = client
        self.documents = []
        self.embeddings = []
    
    def add_document(self, text):
        self.documents.append(text)
        embedding = self.client.get_embedding(text)
        self.embeddings.append(embedding)
    
    def search(self, query, top_k=3):
        query_embedding = self.client.get_embedding(query)
        similarities = cosine_similarity([query_embedding], self.embeddings)[0]
        top_indices = np.argsort(similarities)[-top_k:][::-1]
        
        return [(self.documents[i], similarities[i]) for i in top_indices]

# 使用示例
search_engine = SemanticSearch(EmbeddingClient())
search_engine.add_document("苹果公司是一家科技公司")
search_engine.add_document("香蕉是一种热带水果")
search_engine.add_document("机器学习是人工智能的一个分支")

results = search_engine.search("人工智能技术")
for doc, score in results:
    print(f"相似度: {score:.3f} - {doc}")

8. 总结

走完这个教程,你应该已经成功将EmbeddingGemma-300m封装成了一个完整的RESTful微服务。我们从最简单的单文本嵌入开始,逐步添加了批量处理、并发优化、健康检查等生产环境需要的功能,最后还提供了Docker部署方案和负载均衡配置。

实际使用下来,这个方案部署简单,性能也足够应对大多数场景。特别是在批量处理文本时,相比单条处理效率提升明显。如果你需要处理更大规模的数据,可以考虑进一步优化,比如添加更高效的缓存策略,或者使用专门的向量数据库来存储和管理嵌入结果。

最重要的是,现在你的嵌入能力已经变成了一个可随时调用的服务,无论是内部系统集成还是对外提供API,都变得简单多了。如果你在实践过程中遇到问题,或者有更好的优化建议,欢迎交流讨论。


获取更多AI镜像

想探索更多AI镜像和应用场景?访问 CSDN星图镜像广场,提供丰富的预置镜像,覆盖大模型推理、图像生成、视频生成、模型微调等多个领域,支持一键部署。

Logo

小龙虾开发者社区是 CSDN 旗下专注 OpenClaw 生态的官方阵地,聚焦技能开发、插件实践与部署教程,为开发者提供可直接落地的方案、工具与交流平台,助力高效构建与落地 AI 应用

更多推荐