reranker.py

"""
원격 코드 실행 옵션이 추가된 리랭커 모듈
"""
from typing import List, Dict, Tuple
import numpy as np
from sentence_transformers import CrossEncoder
from langchain.schema import Document
from config import RERANKER_MODEL

class Reranker:
    def __init__(self, model_name: str = RERANKER_MODEL):
        """
        Cross-Encoder 리랭커 초기화

        Args:
            model_name: 사용할 Cross-Encoder 모델 이름
        """
        print(f"리랭커 모델 로드 중: {model_name}")

        # 원격 코드 실행 허용 옵션 추가
        self.model = CrossEncoder(
            model_name,
            trust_remote_code=True  # 원격 코드 실행 허용 (필수)
        )

        print(f"리랭커 모델 로드 완료: {model_name}")

    def rerank(self, query: str, documents: List[Document], top_k: int = 3) -> List[Document]:
        """
        검색 결과 재정렬

        Args:
            query: 검색 쿼리
            documents: 벡터 검색 결과 문서 리스트
            top_k: 반환할 상위 결과 수

        Returns:
            재정렬된 상위 문서 리스트
        """
        if not documents:
            return []

        # Cross-Encoder 입력 쌍 생성
        document_texts = [doc.page_content for doc in documents]
        query_doc_pairs = [(query, doc) for doc in document_texts]

        # 점수 계산
        print(f"리랭킹 수행 중: {len(documents)}개 문서")
        scores = self.model.predict(query_doc_pairs)

        # 점수에 따라 문서 재정렬
        doc_score_pairs = list(zip(documents, scores))
        doc_score_pairs.sort(key=lambda x: x[1], reverse=True)

        print(f"리랭킹 완료: 상위 {top_k}개 문서 선택")

        # 상위 k개 결과 반환
        return [doc for doc, score in doc_score_pairs[:top_k]]