File size: 1,080 Bytes
a6fd1a3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
import logging
from functools import lru_cache
from langchain.retrievers.document_compressors import CrossEncoderReranker
from langchain_community.cross_encoders import HuggingFaceCrossEncoder
import config

logger = logging.getLogger(__name__)

# Đặt tên model vào một hằng số để dễ quản lý


@lru_cache(maxsize=1)
def get_reranker_compressor(top_n: int = 4):
    """
    Tải và trả về một đối tượng CrossEncoderReranker.
    Sử dụng lru_cache để đảm bảo model chỉ được tải một lần duy nhất.
    """
    logger.info(f"🧠 Loading Re-ranker model '{config.RERANKER_MODEL_NAME}'...")
    try:
        # Tải model cross-encoder
        model = HuggingFaceCrossEncoder(model_name=config.RERANKER_MODEL_NAME)

        # Tạo đối tượng compressor
        compressor = CrossEncoderReranker(model=model, top_n=top_n)

        logger.info("✅ Re-ranker model is ready.")
        return compressor
    except Exception as e:
        logger.error(f"❌ Could not load Re-ranker model: {e}", exc_info=True)
        raise