Spaces:
Sleeping
Sleeping
| """ | |
| reranker.py | |
| Reranks retrieved chunks using the Cohere Rerank API. | |
| Cohere rerank takes the query and a list of texts and returns | |
| relevance scores — no local model, no GPU needed. | |
| Sign up: https://dashboard.cohere.com | |
| Free tier: 1000 calls/month. Paid: $1 per 1000 calls. | |
| Dependencies: | |
| pip install cohere python-dotenv | |
| """ | |
| import logging | |
| import os | |
| from typing import List | |
| import cohere | |
| from dotenv import load_dotenv | |
| from .models import RankedResult, SearchResult | |
| load_dotenv() | |
| logger = logging.getLogger(__name__) | |
| RERANK_MODEL = "rerank-english-v3.0" | |
| def create_cohere_client() -> cohere.Client: | |
| """ | |
| Create a Cohere client from environment variables. | |
| Requires in .env: | |
| COHERE_API_KEY = your-key | |
| """ | |
| api_key = os.getenv("COHERE_API_KEY") | |
| if not api_key: | |
| raise EnvironmentError( | |
| "COHERE_API_KEY not found. Add it to your .env:\n" | |
| " COHERE_API_KEY=your-key\n" | |
| "Get one at: https://dashboard.cohere.com" | |
| ) | |
| return cohere.Client(api_key=api_key) | |
| def rerank( | |
| query: str, | |
| results: List[SearchResult], | |
| cohere_client: cohere.Client, | |
| top_n: int = 10, | |
| ) -> List[RankedResult]: | |
| """ | |
| Rerank retrieved chunks using Cohere Rerank. | |
| Args: | |
| query: User query string. | |
| results: Chunks from the retrieval step. | |
| cohere_client: Cohere client. | |
| top_n: Number of top results to return after reranking. | |
| Returns: | |
| List of RankedResult ordered by Cohere relevance score, best first. | |
| """ | |
| if not results: | |
| return [] | |
| logger.info("Reranking %d chunks (top_n=%d)...", len(results), top_n) | |
| response = cohere_client.rerank( | |
| model=RERANK_MODEL, | |
| query=query, | |
| documents=[r.text for r in results], | |
| top_n=top_n, | |
| ) | |
| ranked = [] | |
| for hit in response.results: | |
| original = results[hit.index] | |
| ranked.append(RankedResult( | |
| chunk_id= original.chunk_id, | |
| text= original.text, | |
| source= original.source, | |
| file_path= original.file_path, | |
| rerank_score= hit.relevance_score, | |
| )) | |
| logger.info( | |
| "Reranking complete. Top score: %.4f Bottom score: %.4f", | |
| ranked[0].rerank_score if ranked else 0, | |
| ranked[-1].rerank_score if ranked else 0, | |
| ) | |
| return ranked | |