Spaces:
Sleeping
Sleeping
File size: 3,558 Bytes
1b447de |
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 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 |
"""
Sistema de reranking com cross-encoder
"""
from typing import List, Dict, Any, Optional
from sentence_transformers import CrossEncoder
from .config import RERANKER_MODEL_ID
class Reranker:
"""Reranker usando cross-encoder para melhor precisão"""
def __init__(self, model_id: str = RERANKER_MODEL_ID):
"""
Inicializa reranker
Args:
model_id: ID do modelo cross-encoder
"""
self.model_id = model_id
self.model: Optional[CrossEncoder] = None
def load_model(self) -> CrossEncoder:
"""Carrega cross-encoder (lazy loading)"""
if self.model is None:
self.model = CrossEncoder(self.model_id)
return self.model
def rerank(
self,
query: str,
documents: List[Dict[str, Any]],
top_k: Optional[int] = None
) -> List[Dict[str, Any]]:
"""
Reordena documentos usando cross-encoder
Args:
query: Query do usuário
documents: Lista de documentos com 'content' e 'score'
top_k: Retornar apenas top K (None = todos)
Returns:
Documentos reordenados com 'rerank_score'
"""
if not documents:
return []
model = self.load_model()
# Prepara pares (query, doc)
pairs = [(query, doc['content']) for doc in documents]
# Calcula scores do cross-encoder
scores = model.predict(pairs)
# Adiciona rerank_score e preserva original_score
for doc, score in zip(documents, scores):
doc['rerank_score'] = float(score)
doc['original_score'] = doc.get('score', 0.0)
# Reordena por rerank_score
reranked = sorted(documents, key=lambda x: x['rerank_score'], reverse=True)
if top_k:
reranked = reranked[:top_k]
return reranked
def get_rerank_comparison(
self,
original_docs: List[Dict[str, Any]],
reranked_docs: List[Dict[str, Any]]
) -> List[Dict[str, Any]]:
"""
Gera dados de comparação antes/depois do reranking
Args:
original_docs: Documentos com ordem original
reranked_docs: Documentos após reranking
Returns:
Lista de dicionários com comparação
"""
comparison = []
# Cria mapa de IDs para posições originais
original_positions = {doc['id']: i+1 for i, doc in enumerate(original_docs)}
for new_rank, doc in enumerate(reranked_docs, 1):
original_rank = original_positions.get(doc['id'], -1)
position_change = original_rank - new_rank if original_rank != -1 else 0
comparison.append({
'new_rank': new_rank,
'original_rank': original_rank,
'original_score': doc.get('original_score', 0.0),
'rerank_score': doc.get('rerank_score', 0.0),
'position_change': position_change,
'content_preview': doc['content'][:100] + "..."
})
return comparison
def is_available(self) -> bool:
"""Verifica se reranker está disponível"""
try:
self.load_model()
return True
except Exception:
return False
def get_model_info(self) -> Dict[str, Any]:
"""Retorna informações do modelo"""
return {
"model_id": self.model_id,
"available": self.is_available(),
"type": "cross-encoder"
}
|