models.py (14284B)
1 from __future__ import annotations 2 from typing import NamedTuple, Any 3 from abc import ABC, abstractmethod 4 from abc import ABC, abstractmethod 5 from transformers import AutoTokenizer, AutoModel 6 from sentence_transformers import SentenceTransformer 7 import functools 8 import torch 9 10 DEVICE = "cuda" if torch.cuda.is_available() else "cpu" 11 12 13 class Model(ABC): 14 15 @abstractmethod 16 def embed(self, string: str) -> torch.Tensor: 17 pass 18 19 @abstractmethod 20 def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float: 21 pass 22 23 24 class ModelFactory: 25 26 _factory: dict[str, type[Model]] = {} 27 28 def __init__(self, name: str): 29 self._name = name 30 31 def __call__(self, model: type[Model]) -> type[Model]: 32 self._factory[self._name] = model 33 return model 34 35 @classmethod 36 def get(cls: type[ModelFactory], name: str, *args: Any, **kwargs: Any) -> Model: 37 return cls._factory[name](*args, **kwargs) 38 39 40 # Generated using Mistral 41 def mean_pool_last_four_layers(string: str, tokenizer, encoder) -> torch.Tensor: 42 inputs = tokenizer(string.lower(), return_tensors="pt") 43 44 # Get all hidden states (set output_hidden_states=True) 45 outputs = encoder(**inputs, output_hidden_states=True) 46 47 # Get the last four hidden states (layers) 48 # hidden_states is a tuple of length num_layers+1 (includes embedding layer) 49 # We take layers [-4:] to get the last four layers (excluding embedding layer) 50 last_four_layers = outputs.hidden_states[-4:] 51 52 # Stack the last four layers and average them 53 stacked_layers = torch.stack( 54 last_four_layers, dim=0 55 ) # Shape: (4, batch_size, seq_len, hidden_size) 56 averaged_layers = torch.mean( 57 stacked_layers, dim=0 58 ) # Shape: (batch_size, seq_len, hidden_size) 59 60 attention_mask = inputs["attention_mask"] 61 62 # Expand attention mask to match token embeddings dimensions 63 attention_mask_expanded = ( 64 attention_mask.unsqueeze(-1).expand(averaged_layers.size()).float() 65 ) 66 67 # Apply mask to token embeddings (set padding tokens to 0) 68 masked_embeddings = averaged_layers * attention_mask_expanded 69 70 # Sum embeddings and divide by number of non-padding tokens 71 summed_embeddings = torch.sum( 72 masked_embeddings, dim=1 73 ) # Shape: (batch_size, hidden_size) 74 summed_mask = torch.clamp( 75 attention_mask_expanded.sum(dim=1), min=1e-9 76 ) # Avoid division by zero 77 78 # Create sentence representation by mean pooling 79 mean_pooled = summed_embeddings / summed_mask 80 81 return mean_pooled 82 83 84 # Generated using Mistral 85 def mean_pool(string: str, tokenizer, encoder) -> torch.Tensor: 86 inputs = tokenizer(string.lower(), return_tensors="pt") 87 outputs = encoder(**inputs) 88 89 token_embeddings = outputs.last_hidden_state 90 attention_mask = inputs["attention_mask"] 91 92 # Expand attention mask to match token embeddings dimensions 93 attention_mask_expanded = ( 94 attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() 95 ) 96 97 # Apply mask to token embeddings (set padding tokens to 0) 98 masked_embeddings = token_embeddings * attention_mask_expanded 99 100 # Sum embeddings and divide by number of non-padding tokens 101 summed_embeddings = torch.sum( 102 masked_embeddings, dim=1 103 ) # Shape: (batch_size, hidden_size) 104 summed_mask = torch.clamp( 105 attention_mask_expanded.sum(dim=1), min=1e-9 106 ) # Avoid division by zero 107 108 # Create sentence representation by mean pooling 109 mean_pooled = summed_embeddings / summed_mask 110 111 return mean_pooled 112 113 114 # Generated using Mistral 115 def cls_sep_pool(string: str, tokenizer, encoder) -> torch.Tensor: 116 inputs = tokenizer(string.lower(), return_tensors="pt") 117 outputs = encoder(**inputs) 118 token_embeddings = outputs.last_hidden_state 119 attention_mask = inputs["attention_mask"] 120 input_ids = inputs["input_ids"] 121 122 batch_size, seq_len, hidden_size = token_embeddings.shape 123 124 # Initialize result tensor 125 cls_sep_embeddings = torch.zeros(batch_size, hidden_size) 126 127 for batch_idx in range(batch_size): 128 # Find the sequence length (excluding padding) 129 valid_length = attention_mask[batch_idx].sum().item() 130 131 # CLS token is always at position 0 132 cls_embedding = token_embeddings[batch_idx, 0] 133 134 # Find SEP token position by looking for the SEP token ID 135 # SEP token is the last non-padding token before padding starts 136 sep_token_id = tokenizer.sep_token_id 137 sep_position = None 138 139 # Search for SEP token in the valid sequence 140 for pos in range(valid_length): 141 if input_ids[batch_idx, pos] == sep_token_id: 142 sep_position = pos 143 break 144 145 # If no SEP token found, use the last valid token 146 if sep_position is None: 147 raise Exception("No SEP token found") 148 149 sep_embedding = token_embeddings[batch_idx, sep_position] 150 151 # Sum CLS and SEP embeddings 152 cls_sep_embeddings[batch_idx] = cls_embedding + sep_embedding 153 154 return cls_sep_embeddings 155 156 157 @ModelFactory("BERT-cls") 158 class BERT_cls(Model): 159 160 def __init__(self) -> None: 161 self.model_name: str = "bert-base-uncased" 162 self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) 163 self.encoder = AutoModel.from_pretrained(self.model_name) 164 self.encoder.eval() 165 166 @functools.lru_cache 167 @torch.no_grad() 168 def embed(self, string: str) -> torch.Tensor: 169 return self.encoder( 170 self.tokenizer(string.lower(), return_tensors="pt")["input_ids"] 171 ).last_hidden_state[:, 0, :] 172 173 def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float: 174 return torch.nn.functional.cosine_similarity( 175 a, b 176 ).item() # dist(\mu(a), \mu(b)) 177 178 179 @ModelFactory("ModernBERT-cls") 180 class ModernBERT_cls(Model): 181 182 def __init__(self) -> None: 183 self.model_name: str = "answerdotai/ModernBERT-base" 184 self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) 185 self.encoder = AutoModel.from_pretrained(self.model_name) 186 self.encoder.eval() 187 188 @functools.lru_cache 189 @torch.no_grad() 190 def embed(self, string: str) -> torch.Tensor: 191 return self.encoder( 192 self.tokenizer(string.lower(), return_tensors="pt")["input_ids"] 193 ).last_hidden_state[:, 0, :] 194 195 def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float: 196 return torch.nn.functional.cosine_similarity( 197 a, b 198 ).item() # dist(\mu(a), \mu(b)) 199 200 201 @ModelFactory("ModernBERT-mean") 202 class ModernBERT_mean(Model): 203 204 def __init__(self) -> None: 205 self.model_name: str = "answerdotai/ModernBERT-base" 206 self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) 207 self.encoder = AutoModel.from_pretrained(self.model_name) 208 self.encoder.eval() 209 210 @functools.lru_cache 211 @torch.no_grad() 212 def embed(self, string: str) -> torch.Tensor: 213 return mean_pool_last_four_layers( 214 string=string, tokenizer=self.tokenizer, encoder=self.encoder 215 ) 216 217 def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float: 218 return torch.nn.functional.cosine_similarity( 219 a, b 220 ).item() # dist(\mu(a), \mu(b)) 221 222 223 @ModelFactory("ModernBERT-clssep") 224 class ModernBERT_clssep(Model): 225 226 def __init__(self) -> None: 227 self.model_name: str = "answerdotai/ModernBERT-base" 228 self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) 229 self.encoder = AutoModel.from_pretrained(self.model_name) 230 self.encoder.eval() 231 232 @functools.lru_cache 233 @torch.no_grad() 234 def embed(self, string: str) -> torch.Tensor: 235 return cls_sep_pool( 236 string=string, tokenizer=self.tokenizer, encoder=self.encoder 237 ) 238 239 def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float: 240 return torch.nn.functional.cosine_similarity( 241 a, b 242 ).item() # dist(\mu(a), \mu(b)) 243 244 245 @ModelFactory("DeBERTa-cls") 246 class DeBERTa_cls(Model): 247 248 def __init__(self) -> None: 249 self.model_name: str = "microsoft/deberta-v3-base" 250 self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) 251 self.encoder = AutoModel.from_pretrained(self.model_name) 252 self.encoder.eval() 253 254 @functools.lru_cache 255 @torch.no_grad() 256 def embed(self, string: str) -> torch.Tensor: 257 return self.encoder( 258 self.tokenizer(string.lower(), return_tensors="pt")["input_ids"] 259 ).last_hidden_state[:, 0, :] 260 261 def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float: 262 return torch.nn.functional.cosine_similarity( 263 a, b 264 ).item() # dist(\mu(a), \mu(b)) 265 266 267 @ModelFactory("DeBERTa-mean") 268 class DeBERTa_mean(Model): 269 270 def __init__(self) -> None: 271 self.model_name: str = "microsoft/deberta-v3-base" 272 self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) 273 self.encoder = AutoModel.from_pretrained(self.model_name) 274 self.encoder.eval() 275 276 @functools.lru_cache 277 @torch.no_grad() 278 def embed(self, string: str) -> torch.Tensor: 279 return mean_pool_last_four_layers( 280 string=string, tokenizer=self.tokenizer, encoder=self.encoder 281 ) 282 283 def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float: 284 return torch.nn.functional.cosine_similarity( 285 a, b 286 ).item() # dist(\mu(a), \mu(b)) 287 288 289 @ModelFactory("DeBERTa-clssep") 290 class DeBERTa_clssep(Model): 291 292 def __init__(self) -> None: 293 self.model_name: str = "microsoft/deberta-v3-base" 294 self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) 295 self.encoder = AutoModel.from_pretrained(self.model_name) 296 self.encoder.eval() 297 298 @functools.lru_cache 299 @torch.no_grad() 300 def embed(self, string: str) -> torch.Tensor: 301 return cls_sep_pool( 302 string=string, tokenizer=self.tokenizer, encoder=self.encoder 303 ) 304 305 def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float: 306 return torch.nn.functional.cosine_similarity( 307 a, b 308 ).item() # dist(\mu(a), \mu(b)) 309 310 311 @ModelFactory("BERT-mean") 312 class BERT_mean(Model): 313 314 def __init__(self) -> None: 315 self.model_name: str = "bert-base-uncased" 316 self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) 317 self.encoder = AutoModel.from_pretrained(self.model_name) 318 self.encoder.eval() 319 320 @functools.lru_cache 321 @torch.no_grad() 322 def embed(self, string: str) -> torch.Tensor: 323 return mean_pool_last_four_layers( 324 string=string, tokenizer=self.tokenizer, encoder=self.encoder 325 ) 326 327 def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float: 328 return torch.nn.functional.cosine_similarity( 329 a, b 330 ).item() # dist(\mu(a), \mu(b)) 331 332 333 @ModelFactory("BERT-clssep") 334 class BERT_clssep(Model): 335 336 def __init__(self) -> None: 337 self.model_name: str = "bert-base-uncased" 338 self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) 339 self.encoder = AutoModel.from_pretrained(self.model_name) 340 self.encoder.eval() 341 342 @functools.lru_cache 343 @torch.no_grad() 344 def embed(self, string: str) -> torch.Tensor: 345 return cls_sep_pool( 346 string=string, tokenizer=self.tokenizer, encoder=self.encoder 347 ) 348 349 def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float: 350 return torch.nn.functional.cosine_similarity( 351 a, b 352 ).item() # dist(\mu(a), \mu(b)) 353 354 355 @ModelFactory("S-BERT") 356 class S_BERT(Model): 357 358 def __init__(self): 359 self.encoder = SentenceTransformer("all-MiniLM-L6-v2", device=DEVICE) 360 361 @functools.lru_cache 362 def embed(self, string: str) -> torch.Tensor: 363 return self.encoder.encode([string]) 364 365 def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float: 366 return self.encoder.similarity(a, b).item() 367 368 369 @ModelFactory("S-GPT") 370 class S_GPT(Model): 371 372 def __init__(self): 373 self.encoder = SentenceTransformer( 374 "Muennighoff/SGPT-125M-weightedmean-nli-bitfit", device=DEVICE 375 ) 376 377 @functools.lru_cache 378 def embed(self, string: str) -> torch.Tensor: 379 return self.encoder.encode([string]) 380 381 def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float: 382 return self.encoder.similarity(a, b).item() 383 384 385 @ModelFactory("E5") 386 class E5(Model): 387 388 def __init__(self): 389 self.encoder = SentenceTransformer( 390 "intfloat/multilingual-e5-large-instruct", device=DEVICE 391 ) 392 self.task = "Given a sentence, encode its semantic properties" 393 394 def get_detailed_instruct(self, query: str) -> str: 395 return f"Instruct: {self.task}\nQuery: {query}" 396 397 @functools.lru_cache 398 def embed(self, string: str) -> torch.Tensor: 399 prepared_string = self.get_detailed_instruct(string) 400 return self.encoder.encode([prepared_string]) 401 402 def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float: 403 return self.encoder.similarity(a, b).item() 404 405 406 @ModelFactory("QWEN") 407 class QWEN(Model): 408 409 def __init__(self): 410 self.encoder = SentenceTransformer("Qwen/Qwen3-Embedding-0.6B", device=DEVICE) 411 self.task = "Given a sentence, encode its semantic properties" 412 413 def get_detailed_instruct(self, query: str) -> str: 414 return f"Instruct: {self.task}\nQuery: {query}" 415 416 @functools.lru_cache 417 def embed(self, string: str) -> torch.Tensor: 418 prepared_string = self.get_detailed_instruct(string) 419 return self.encoder.encode([prepared_string]) 420 421 def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float: 422 return self.encoder.similarity(a, b).item() 423 424 425 @ModelFactory("UAE") 426 class UAE(Model): 427 428 def __init__(self): 429 self.encoder = SentenceTransformer("WhereIsAI/UAE-Large-V1", device=DEVICE) 430 431 @functools.lru_cache 432 def embed(self, string: str) -> torch.Tensor: 433 return self.encoder.encode(string) 434 435 def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float: 436 return self.encoder.similarity(a, b).item()