File size: 2,831 Bytes
f7b283c 7a0020b f7b283c 7a0020b f7b283c 7a0020b f7b283c 7a0020b f7b283c 7a0020b f7b283c 7a0020b f7b283c 7a0020b f7b283c 7a0020b f7b283c 7a0020b f7b283c 7a0020b f7b283c |
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 |
from transformers import AutoTokenizer, TFAutoModelForSequenceClassification
import tensorflow as tf
from typing import List
import numpy as np
from logger_config import config_logger
logger = config_logger(__name__)
class CrossEncoderReranker:
"""
Cross-Encoder Re-Ranker that takes (query, candidate) pairs,
outputs a single relevance score in [0,1].
"""
def __init__(self, model_name: str = "cross-encoder/ms-marco-MiniLM-L-12-v2"):
"""
Initialize the cross-encoder with a pretrained model.
Args:
model_name: Name of a HF cross-encoder model. Must be
compatible with TFAutoModelForSequenceClassification.
"""
logger.info(f"Initializing CrossEncoderReranker with {model_name}...")
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.model = TFAutoModelForSequenceClassification.from_pretrained(model_name)
logger.info("Cross encoder model loaded successfully.")
def rerank(
self,
query: str,
candidates: List[str],
max_length: int = 256
) -> List[float]:
"""
Compute relevance scores for each candidate w.r.t. the query.
Args:
query: User's query text.
candidates: List of candidate response texts.
max_length: Max token length for each (query, candidate) pair.
Returns:
A list of float scores in [0,1], one per candidate,
indicating model's predicted relevance.
"""
# 1) Build (query, candidate) pairs
pair_texts = [(query, candidate) for candidate in candidates]
# 2) Tokenize the entire batch
encodings = self.tokenizer(
pair_texts,
padding=True,
truncation=True,
max_length=max_length,
return_tensors="tf"
)
# 3) 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") # Some models need token_type_ids
)
logits = outputs.logits # shape [batch_size, 1]
# 4) Convert logits -> [0,1] range via sigmoid
# If the cross-encoder is a single-logit regression to [0,1],
# this is a typical interpretation.
scores = tf.nn.sigmoid(logits) # shape [batch_size, 1]
# 5) Flatten to a 1D NumPy array of floats
scores = tf.reshape(scores, [-1])
scores = scores.numpy().astype(float)
# logger.debug(f"Cross-Encoder raw logits: {logits.numpy().flatten().tolist()}")
# logger.debug(f"Cross-Encoder sigmoid scores: {scores.tolist()}")
return scores.tolist() |