Files
YG_FT/backend/app/modules/data_process/algorithms/embedding.py

25 lines
794 B
Python
Raw Normal View History

"""数据处理算法 - 本地语义嵌入模型共享单例。"""
from __future__ import annotations
import os
from functools import lru_cache
from typing import Any
@lru_cache(maxsize=1)
def semantic_embedding_model() -> Any:
"""加载本地嵌入模型,供语义分块与语义质量评分共用。
模型可在部署环境覆盖默认模型体积较小且适合中英文语义判断
返回 LlamaIndex BaseEmbedding通过 ``get_text_embedding`` 使用
"""
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
return HuggingFaceEmbedding(
model_name=os.getenv("DATA_PROCESS_EMBEDDING_MODEL", "BAAI/bge-small-zh-v1.5"),
device=os.getenv("DATA_PROCESS_EMBEDDING_DEVICE", "cpu"),
trust_remote_code=False,
)