| from transformers import AutoTokenizer, TFAutoModelForSequenceClassification | |
| import tensorflow as tf | |
| from typing import List, Tuple | |
| from logger_config import config_logger | |
| logger = config_logger(__name__) | |
| class CrossEncoderReranker: | |
| """ | |
| Cross-Encoder Re-Ranker: Takes (query, candidate) pairs, | |
| outputs a single relevance score (one logit). | |
| """ | |
| def __init__(self, model_name: str = "cross-encoder/ms-marco-MiniLM-L-12-v2"): | |
| self.tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| self.model = TFAutoModelForSequenceClassification.from_pretrained(model_name) | |
| # Model outputs shape [batch_size, 1] -> Interpret the logit as relevance score. | |
| def rerank( | |
| self, | |
| query: str, | |
| candidates: List[str], | |
| max_length: int = 256 | |
| ) -> List[float]: | |
| """ | |
| Returns a list of re_scores, one for each candidate, indicating | |
| how relevant the candidate is to the query. | |
| """ | |
| # Build (query, candidate) pairs | |
| pair_texts = [(query, candidate) for candidate in candidates] | |
| # Tokenize the entire batch | |
| encodings = self.tokenizer( | |
| pair_texts, | |
| padding=True, | |
| truncation=True, | |
| max_length=max_length, | |
| return_tensors="tf" | |
| ) | |
| # Forward pass -> logits shape [batch_size, 1] | |
| outputs = self.model( | |
| input_ids=encodings["input_ids"], | |
| attention_mask=encodings["attention_mask"], | |
| token_type_ids=encodings.get("token_type_ids") | |
| ) | |
| logits = outputs.logits | |
| # Flatten to shape [batch_size] | |
| scores = tf.reshape(logits, [-1]).numpy() | |
| return scores.tolist() |