Embedding 模型-text2vec python服务化

最新推荐文章于 2025-08-12 18:06:10 发布

置顶栗子~~

最新推荐文章于 2025-08-12 18:06:10 发布

阅读量469

点赞数 6

CC 4.0 BY-SA版权

分类专栏： AI 文章标签： embedding python 开发语言语言模型 ai

本文链接：https://blog.csdn.net/weixin_38316697/article/details/148946369

AI 专栏收录该内容

13 篇文章

订阅专栏

文章目录

前言
- Embedding 模型-text2vec python服务化

前言

如果您觉得有用的话，记得给博主点个赞，评论，收藏一键三连啊，写作不易啊^ _ ^。
而且听说点赞的人每天的运气都不会太差，实在白嫖的话，那欢迎常来啊!!!

Embedding 模型-text2vec python服务化

1. 安装 PyTorch:

text2vec 依赖 PyTorch（ torch 库），所以需要安装PyTorch。

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu

下载成功:
在这里插入图片描述
验证:

python -c "import torch; print(torch.__version__)"

在这里插入图片描述

2. Embedding 模型-text2vec python demo

模型:
在这里插入图片描述
核心代码:

# 引入 FastAPI 主框架和异常处理类
from fastapi import FastAPI, HTTPException
# 引入 Pydantic，用于请求体/响应体的数据结构校验
from pydantic import BaseModel
# 引入类型注解：List 是列表类型（泛型）
from typing import List
# 引入 text2vec 模型，提供中文句子向量化能力
from text2vec import SentenceModel
import uvicorn
import socket
from common.LogConfig import setup_logger
from fastapi.responses import JSONResponse
from fastapi.requests import Request

import pathlib
import configparser

# 日志和配置初始化
current_dir = pathlib.Path(__file__).parent
config_file = current_dir / 'config' / 'text2vec.ini'
config = configparser.ConfigParser()

try:
    with open(config_file, 'r', encoding='utf-8') as f:
        config.read_file(f)
except Exception as e:
    raise RuntimeError(f"配置文件读取失败: {config_file}, 错误: {str(e)}")

log_file = current_dir / config['logging']['log_file']
log_level = config['logging'].get('log_level', 'INFO').upper()
backup_count = config['logging'].getint('backup_count', 5)
logger_name = config['logging'].get('logger_name', 'text2vec')

embedding_model = config['text2vec'].get('embedding_model', 'shibing624/text2vec-base-chinese')

# 初始化日志配置
log_file.parent.mkdir(parents=True, exist_ok=True)
logger = setup_logger(
    logger_name=logger_name,
    log_file=log_file,
    backup_count=backup_count,
    level_str=log_level
)

# 创建一个 FastAPI 实例，用于定义和启动 Web 服务
app = FastAPI(
    title="Text2Vec Embedding Service",  # API 标题，用于 Swagger 文档展示
    version="1.0"                         # API 版本号
)


# 加载 sentence embedding 模型（中文），第一次运行会自动从 HuggingFace 下载
model = SentenceModel(embedding_model)


# 定义请求数据模型（结构），使用 Pydantic 进行类型检查和自动文档生成
class EmbeddingRequest(BaseModel):
    sentences: List[str]  # 用户传入的句子列表，例如 ["你好", "天气真好"]


# 定义响应数据模型（结构），同样用于类型提示和文档生成
class EmbeddingResponse(BaseModel):
    success: bool                  # 表示接口调用是否成功，比如 True 或 False
    ip: str                       # 请求方的 IP 地址
    msg: str                      # 请求消息
    status_code: int              # 返回码
    vectors: List[List[float]]    # 嵌入向量，二维浮点数组


def get_server_ip():
    """更稳定的本机IP获取方式：通过连接外部服务器获取"""
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("8.8.8.8", 80))  # 可连接任意外部IP，无需实际发送
        ip = s.getsockname()[0]
        s.close()
        return ip
    except Exception as exc:
        logger.error(f"获取本机IP方式失败: {str(exc)}")
        return "127.0.0.1"


# 定义一个健康检查接口（GET请求），用于探测服务是否存活
@app.get("/health")
async def health_check():
    return {"status": "ok"}  # 返回一个简单的 JSON 表示服务正常


# 定义主接口：将句子转为向量（POST 请求），自动返回 EmbeddingResponse 格式的数据
@app.post("/embed", response_model=EmbeddingResponse)
async def get_embedding(body: EmbeddingRequest):
    # 如果用户传入的是空列表，抛出 HTTP 400（Bad Request）异常
    server_ip = get_server_ip()

    if not body.sentences:
        logger.warning('url:/embed, 请求为空')
        return EmbeddingResponse(status_code=1001,
                                 success=False,
                                 ip=server_ip,
                                 msg="Input 'sentences' cannot be empty.",
                                 vectors=[]
                                 )

    try:
        # 使用模型对句子列表进行向量化，返回为二维数组（多个句子的向量）
        vectors = model.encode(body.sentences).tolist()
        logger.info('url:/embed,查询成功')

        # 返回 JSON 格式结果，符合 response_model 中定义的格式
        return EmbeddingResponse(status_code=200,
                                 success=True,
                                 ip=server_ip,
                                 msg='查询成功',
                                 vectors=vectors)

    except Exception as exc:
        logger.error(f"url:/embed, 查询失败: {str(exc)}")

        return EmbeddingResponse(status_code=500,
                                 success=False,
                                 ip=server_ip,
                                 msg=f"Model error: {str(exc)}",
                                 vectors=[])


# 全局异常拦截
@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    logger.exception(f"[GlobalException] path={request.url.path}, error={str(exc)}")
    return JSONResponse(status_code=500, content={
        "success": False,
        "msg": "Internal server error",
        "status_code": 500,
        "ip": get_server_ip(),
        "vectors": []
    })


if __name__ == "__main__":
    # reload=True 是开发模式；正式部署推荐关闭
    uvicorn.run("text2vec_server:app", host="0.0.0.0", port=8000, reload=True)