| import os |
| from typing import List, Tuple, Union |
| import numpy as np |
| import torch |
| import torch.nn as nn |
| from transformers import AutoTokenizer, AutoModel, AutoConfig |
|
|
|
|
| class SapBertCrossEncoder(nn.Module): |
| """ |
| Stage-2 Cross-Encoder for clinical entity linking. |
| Accepts text pairs (mention, candidate_name) and outputs relevance probabilities |
| computed via deep cross-attention layers. |
| """ |
|
|
| def __init__( |
| self, |
| model_name_or_path: str = "checkpoints/sapbert_challenge_finetuned/best_model", |
| num_labels: int = 1, |
| dropout_prob: float = 0.1, |
| ): |
| super().__init__() |
| self.config = AutoConfig.from_pretrained(model_name_or_path) |
| self.encoder = AutoModel.from_pretrained(model_name_or_path, config=self.config) |
| self.dropout = nn.Dropout(dropout_prob) |
| self.classifier = nn.Linear(self.config.hidden_size, num_labels) |
| |
| def forward( |
| self, |
| input_ids: torch.Tensor, |
| attention_mask: torch.Tensor, |
| token_type_ids: torch.Tensor = None, |
| ) -> torch.Tensor: |
| """ |
| Forward pass returning raw logits (shape: [batch_size, 1]). |
| """ |
| kwargs = {"input_ids": input_ids, "attention_mask": attention_mask} |
| if token_type_ids is not None: |
| kwargs["token_type_ids"] = token_type_ids |
|
|
| outputs = self.encoder(**kwargs) |
| |
| cls_rep = outputs.last_hidden_state[:, 0, :] |
| cls_rep = self.dropout(cls_rep) |
| logits = self.classifier(cls_rep) |
| return logits.squeeze(-1) |
|
|
| def predict_scores( |
| self, |
| pairs: List[Tuple[str, str]], |
| tokenizer: AutoTokenizer, |
| batch_size: int = 32, |
| max_length: int = 64, |
| device: Union[str, torch.device] = "cpu", |
| ) -> np.ndarray: |
| """ |
| Computes sigmoid-normalized match probabilities for a list of (query, candidate) pairs. |
| """ |
| self.eval() |
| self.to(device) |
| all_scores = [] |
|
|
| for i in range(0, len(pairs), batch_size): |
| batch_pairs = pairs[i : i + batch_size] |
| queries = [p[0] for p in batch_pairs] |
| candidates = [p[1] for p in batch_pairs] |
|
|
| inputs = tokenizer( |
| queries, |
| candidates, |
| padding=True, |
| truncation=True, |
| max_length=max_length, |
| return_tensors="pt", |
| ).to(device) |
|
|
| with torch.no_grad(): |
| logits = self(**inputs) |
| probs = torch.sigmoid(logits).cpu().numpy() |
| all_scores.extend(probs.flatten().tolist()) |
|
|
| return np.array(all_scores, dtype=np.float32) |
|
|
| def save_pretrained(self, save_directory: str): |
| """Saves model weights and config to directory.""" |
| os.makedirs(save_directory, exist_ok=True) |
| torch.save(self.state_dict(), os.path.join(save_directory, "pytorch_model.bin")) |
| self.config.save_pretrained(save_directory) |
|
|
| @classmethod |
| def from_pretrained(cls, load_directory: str, **kwargs): |
| """Loads a saved Cross-Encoder model.""" |
| model = cls(model_name_or_path=load_directory, **kwargs) |
| weight_path = os.path.join(load_directory, "pytorch_model.bin") |
| if os.path.exists(weight_path): |
| state_dict = torch.load(weight_path, map_location="cpu") |
| model.load_state_dict(state_dict, strict=False) |
| return model |
|
|
