sales-assistant-py-new/app/reranker_utils.py

25 lines
724 B
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

from FlagEmbedding import FlagReranker
from app.conf.app_config import app_config
_reranker_model = None
_reranker_init_failed = False
def get_reranker_model():
global _reranker_model, _reranker_init_failed
if _reranker_init_failed:
return None
if _reranker_model is None:
try:
_reranker_model = FlagReranker(
model_name_or_path=app_config.reranker.model_path,
device="cuda:0",
use_fp16=app_config.reranker.use_fp16
)
except Exception as e:
_reranker_init_failed = True
print(f"[WARNING] reranker初始化失败将跳过rerank: {e}")
return None
return _reranker_model