Reranking Models

Copyexport interface RerankResult {
  index: number;
  rerankScore: number;
}
 
export interface Reranker {
  rerank(
    query: string,
    documents: string[],
    topK?: number,
  ): Promise<RerankResult[]>;
}

Copy        if not COHERE_AVAILABLE:
            raise ImportError("cohere package is required for CohereReranker. Install with: pip install cohere")
        
        self.config = config
        self.api_key = config.api_key or os.getenv("COHERE_API_KEY")
        if not self.api_key:
            raise ValueError("Cohere API key is required. Set COHERE_API_KEY environment variable or pass api_key in config.")
            
        self.model = config.model
        self.client = cohere.Client(self.api_key)

Copy  private async createClient(): Promise<any> {
    const sdk = await loadPeer(
      "cohere-ai",
      "Cohere reranker",
      () => import("cohere-ai"),
    );
    return new sdk.CohereClient({ token: this.apiKey });
  }

Copy            response = self.client.rerank(
                model=self.model,
                query=query,
                documents=doc_texts,
                top_n=top_k or self.config.top_k or len(documents),
                return_documents=self.config.return_documents,
                max_chunks_per_doc=self.config.max_chunks_per_doc,
            )

Copy        if self.config.device is None:
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
        else:
            self.device = self.config.device
 
        self.tokenizer = AutoTokenizer.from_pretrained(self.config.model)
        self.model = AutoModelForSequenceClassification.from_pretrained(self.config.model)
        self.model.to(self.device)
        self.model.eval()

Copy  private load() {
    if (!this.loaded) {
      this.loaded = (async () => {
        const { AutoModelForSequenceClassification, AutoTokenizer } =
          await import("@huggingface/transformers");
        const options: any = {};
        if (this.device) options.device = this.device;
        const model = await AutoModelForSequenceClassification.from_pretrained(
          this.modelId,
          options,
        );
        const tokenizer = await AutoTokenizer.from_pretrained(this.modelId);
        return { model, tokenizer };
      })();
    }
    return this.loaded;
  }

Copy    @staticmethod
    def _normalize_scores(scores: List[float]) -> List[float]:
        if not scores:
            return []
        arr = np.asarray(scores, dtype=float)
        return (1.0 / (1.0 + np.exp(-arr))).tolist()

Copyconst sigmoid = (x: number) => 1 / (1 + Math.exp(-x));

Copy      const { logits } = await model(inputs);
      const rows: unknown[] = logits.tolist();
 
      const scored = rows.map((row, index) => {
        const logit = Array.isArray(row) ? (row[0] as number) : (row as number);
        return {
          index,
          rerankScore: this.normalize ? sigmoid(logit) : logit,
        };
      });

Copyclass SentenceTransformerRerankerConfig(BaseRerankerConfig):
    """
    Configuration class for Sentence Transformer reranker-specific parameters.
    Inherits from BaseRerankerConfig and adds Sentence Transformer-specific settings.
    """
 
    model: Optional[str] = Field(default="cross-encoder/ms-marco-MiniLM-L-6-v2", description="The cross-encoder model name to use")
    device: Optional[str] = Field(default=None, description="Device to run the model on ('cpu', 'cuda', etc.)")
    batch_size: int = Field(default=32, description="Batch size for processing documents")
    show_progress_bar: bool = Field(default=False, description="Whether to show progress bar during processing")

Copy# Convert to SentenceTransformerRerankerConfig if needed
if isinstance(config, dict):
    config = SentenceTransformerRerankerConfig(**config)
elif isinstance(config, BaseRerankerConfig) and not isinstance(config, SentenceTransformerRerankerConfig):
    # Convert BaseRerankerConfig to SentenceTransformerRerankerConfig with defaults
    config = SentenceTransformerRerankerConfig(
        provider=getattr(config, 'provider', 'sentence_transformer'),
        model=getattr(config, 'model', 'cross-encoder/ms-marco-MiniLM-L-6-v2'),
        api_key=getattr(config, 'api_key', None),
        top_k=getattr(config, 'top_k', None),
        device=None,  # Will auto-detect
        batch_size=32,  # Default
        show_progress_bar=False,  # Default
    )

Copy# Extract text content for reranking
doc_texts = []
for doc in documents:
    if 'memory' in doc:
        doc_texts.append(doc['memory'])
    elif 'text' in doc:
        doc_texts.append(doc['text'])  
    elif 'content' in doc:
        doc_texts.append(doc['content'])
    else:
        doc_texts.append(str(doc))
 
try:
    # Create query-document pairs
    pairs = [[query, doc_text] for doc_text in doc_texts]
 
    scores = self.model.predict(
        pairs,
        batch_size=self.config.batch_size,
        show_progress_bar=self.config.show_progress_bar,
    )
    if isinstance(scores, np.ndarray):
        scores = scores.tolist()
 
    # Combine documents with scores
    doc_score_pairs = list(zip(documents, scores))
 
    # Sort by score (descending)
    doc_score_pairs.sort(key=lambda x: x[1], reverse=True)
 

Copymodel: str = Field(
    default="gpt-5-mini",
    description="LLM model to use for reranking"
)
api_key: Optional[str] = Field(
    default=None,
    description="API key for the LLM provider"
)
provider: str = Field(
    default="openai",
    description="LLM provider (openai, anthropic, etc.)"
)
top_k: Optional[int] = Field(
    default=None,
    description="Number of top documents to return after reranking"
)
temperature: float = Field(
    default=0.0,
    description="Temperature for LLM generation"
)
max_tokens: int = Field(
    default=100,
    description="Maximum tokens for LLM response"
)
scoring_prompt: Optional[str] = Field(
    default=None,
    description="Custom prompt template for scoring documents"
)
llm: Optional[Dict[str, Any]] = Field(
    default=None,
    description="Nested LLM configuration with 'provider' and 'config' keys. "
    "Overrides top-level provider/model/api_key when provided.",
)

Copysafe_query = query[: self._MAX_INPUT_LEN]
safe_doc = doc_text[: self._MAX_INPUT_LEN]
user_message = f"Query: {safe_query}\n\nDocument: {safe_doc}"
 
response = self.llm.generate_response(
    messages=[
        {"role": "system", "content": self._system_prompt},
        {"role": "user", "content": user_message},
    ]
)

Copyprivate async score(query: string, document: string): Promise<number> {
  const safeQuery = query.slice(0, MAX_INPUT_LEN);
  const safeDoc = document.slice(0, MAX_INPUT_LEN);
  const userMessage = `Query: ${safeQuery}\n\nDocument: ${safeDoc}`;
 
  const response = await this.llm.generateResponse([
    { role: "system", content: SYSTEM_PROMPT },
    { role: "user", content: userMessage },
  ]);

Copydef _extract_score(self, response_text: str) -> float:
    """Extract numerical score from LLM response."""
    # Prefer a decimal, fall back to an integer, then clamp: out-of-range outputs
    # like "2.0"/"5" become 1.0 instead of being mis-parsed into a stray 0/1 digit.
    matches = re.findall(r'-?\d+\.\d+', response_text) or re.findall(r'-?\d+', response_text)
 
    if matches:
        score = float(matches[0])
        return min(max(score, 0.0), 1.0)  # Clamp between 0.0 and 1.0
 
    # Fallback: return 0.5 if no valid score found
    return 0.5

Copyprivate extractScore(responseText: string): number {
  const matches =
    responseText.match(/-?\d+\.\d+/g) || responseText.match(/-?\d+/g);
 
  if (matches && matches.length > 0) {
    const score = parseFloat(matches[0]);
    return Math.min(Math.max(score, 0.0), 1.0);
  }
 
  return 0.5;
}

Copydef __init__(self, config):
    """
    Initialize Zero Entropy reranker.
 
    Args:
        config: ZeroEntropyRerankerConfig object with configuration parameters
    """
    if not ZERO_ENTROPY_AVAILABLE:
        raise ImportError("zeroentropy package is required for ZeroEntropyReranker. Install with: pip install zeroentropy")
 
    self.config = config
    self.api_key = config.api_key or os.getenv("ZERO_ENTROPY_API_KEY")
    if not self.api_key:
        raise ValueError("Zero Entropy API key is required. Set ZERO_ENTROPY_API_KEY environment variable or pass api_key in config.")
 
    self.model = config.model or "zerank-1"
 
    # Initialize Zero Entropy client
    if self.api_key:
        self.client = ZeroEntropy(api_key=self.api_key)
    else:
        self.client = ZeroEntropy()  # Will use ZERO_ENTROPY_API_KEY from environment
 

Copyprivate async getClient(): Promise<any> {
  if (this.clientInstance) return this.clientInstance;
  if (!this.clientPromise) {
    this.clientPromise = this.createClient();
  }
  this.clientInstance = await this.clientPromise;
  return this.clientInstance;
}
 
private async createClient(): Promise<any> {
  const sdk = await loadPeer(
    "zeroentropy",
    "ZeroEntropy reranker",
    () => import("zeroentropy"),
  );
  return new sdk.ZeroEntropy({ apiKey: this.apiKey });
}

Copydef rerank(self, query: str, documents: List[Dict[str, Any]], top_k: int = None) -> List[Dict[str, Any]]:
    """
    Rerank documents using Zero Entropy's rerank API.
 
    Args:
        query: The search query
        documents: List of documents to rerank
        top_k: Number of top documents to return
 
    Returns:
        List of reranked documents with rerank_score
    """
    if not documents:
        return documents
 
    # Extract text content for reranking
    doc_texts = []
    for doc in documents:
        if 'memory' in doc:
            doc_texts.append(doc['memory'])
        elif 'text' in doc:
            doc_texts.append(doc['text'])  
        elif 'content' in doc:
            doc_texts.append(doc['content'])
        else:
            doc_texts.append(str(doc))
 
    try:
        # Call Zero Entropy rerank API
        response = self.client.models.rerank(
            model=self.model,
            query=query,
            documents=doc_texts,
        )
 
        # Create reranked results
        reranked_docs = []
        for result in response.results:
            original_doc = documents[result.index].copy()
            original_doc['rerank_score'] = result.relevance_score
            reranked_docs.append(original_doc)
 
        # Sort by relevance score in descending order
        reranked_docs.sort(key=lambda x: x['rerank_score'], reverse=True)
 
        # Apply top_k limit
        if top_k:
            reranked_docs = reranked_docs[:top_k]
        elif self.config.top_k:
            reranked_docs = reranked_docs[:self.config.top_k]
 
        return reranked_docs

Copyasync rerank(
  query: string,
  documents: string[],
  topK?: number,
): Promise<RerankResult[]> {
  if (documents.length === 0) return [];
 
  try {
    const client = await this.getClient();
    const response = await client.models.rerank({
      model: this.model,
      query,
      documents,
    });
 
    const scored: RerankResult[] = response.results.map((result: any) => ({
      index: result.index,
      rerankScore: result.relevance_score,
    }));
    scored.sort((a, b) => b.rerankScore - a.rerankScore);
 
    const finalTopK = topK || this.topK;
    return finalTopK ? scored.slice(0, finalTopK) : scored;