Spaces:
Running
Running
File size: 1,364 Bytes
2ecc4a7 27d1bb9 2ecc4a7 | 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 | from .base import BaseRAGTechnique
from ..services.embed_service import get_embedding
from ..services.rerank_service import rerank_service
from typing import List, Dict, Any
class ReRanking(BaseRAGTechnique):
async def retrieve(self, query: str, document_id: str, top_k: int, **kwargs) -> List[Dict[str, Any]]:
# 1. Embed query
await self.emit("EMBED", "#8B5CF6", "Embedding query...")
q_vec = get_embedding(query)
# 2. Vector Search (Fetch more candidates for re-ranking)
await self.emit("RETRIEVE", "#16A34A", f"pgvector: fetching top-{top_k*4} candidates...")
candidates = await self.supabase.vector_search(q_vec, document_id, self.user_id, top_k * 4)
if not candidates:
return []
# 3. Cross-Encoder Re-ranking
await self.emit("RERANK", "#EF4444", f"Cross-encoder re-scoring {len(candidates)} pairs...")
reranked = rerank_service.rerank(query, candidates, top_k)
await self.emit("DONE", "#22C55E", f"Re-ranked complete. top-{top_k} returned.")
return reranked
async def generate(self, query: str, chunks: List[Dict[str, Any]]) -> str:
await self.emit("GENERATE", "#7C3AED", "Qwen3 generating answer...")
prompt = self.build_prompt(query, chunks)
return self.llm.generate(prompt)
|