commit 5c05b9ac51c42c9f654566d7d01778784d12959b
parent ce9b7f1f69ff44eb21db4410e9557a9873f0c705
Author: Étienne Simon <esimon@esimon.eu>
Date: Mon, 29 Sep 2025 06:12:44 +0200
Move stuff around to prevent breakage
Diffstat:
5 files changed, 441 insertions(+), 422 deletions(-)
diff --git a/.gitmodules b/.gitmodules
@@ -0,0 +1,3 @@
+[submodule "epsilon_compositionality/data/Finding-Idiomaticity-in-Word-Representations"]
+ path = epsilon_compositionality/data/Finding-Idiomaticity-in-Word-Representations
+ url = git@github.com:risehnhew/Finding-Idiomaticity-in-Word-Representations.git
diff --git a/epsilon_compositionality/data/Finding-Idiomaticity-in-Word-Representations b/epsilon_compositionality/data/Finding-Idiomaticity-in-Word-Representations
@@ -0,0 +1 @@
+Subproject commit 9f2dc4e58473f0a314c8a570dca26fa57f987917
diff --git a/epsilon_compositionality/models.py b/epsilon_compositionality/models.py
@@ -0,0 +1,436 @@
+from __future__ import annotations
+from typing import NamedTuple, Any
+from abc import ABC, abstractmethod
+from abc import ABC, abstractmethod
+from transformers import AutoTokenizer, AutoModel
+from sentence_transformers import SentenceTransformer
+import functools
+import torch
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+class Model(ABC):
+
+ @abstractmethod
+ def embed(self, string: str) -> torch.Tensor:
+ pass
+
+ @abstractmethod
+ def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
+ pass
+
+
+class ModelFactory:
+
+ _factory: dict[str, type[Model]] = {}
+
+ def __init__(self, name: str):
+ self._name = name
+
+ def __call__(self, model: type[Model]) -> type[Model]:
+ self._factory[self._name] = model
+ return model
+
+ @classmethod
+ def get(cls: type[ModelFactory], name: str, *args: Any, **kwargs: Any) -> Model:
+ return cls._factory[name](*args, **kwargs)
+
+
+# Generated using Mistral
+def mean_pool_last_four_layers(string: str, tokenizer, encoder) -> torch.Tensor:
+ inputs = tokenizer(string.lower(), return_tensors="pt")
+
+ # Get all hidden states (set output_hidden_states=True)
+ outputs = encoder(**inputs, output_hidden_states=True)
+
+ # Get the last four hidden states (layers)
+ # hidden_states is a tuple of length num_layers+1 (includes embedding layer)
+ # We take layers [-4:] to get the last four layers (excluding embedding layer)
+ last_four_layers = outputs.hidden_states[-4:]
+
+ # Stack the last four layers and average them
+ stacked_layers = torch.stack(
+ last_four_layers, dim=0
+ ) # Shape: (4, batch_size, seq_len, hidden_size)
+ averaged_layers = torch.mean(
+ stacked_layers, dim=0
+ ) # Shape: (batch_size, seq_len, hidden_size)
+
+ attention_mask = inputs["attention_mask"]
+
+ # Expand attention mask to match token embeddings dimensions
+ attention_mask_expanded = (
+ attention_mask.unsqueeze(-1).expand(averaged_layers.size()).float()
+ )
+
+ # Apply mask to token embeddings (set padding tokens to 0)
+ masked_embeddings = averaged_layers * attention_mask_expanded
+
+ # Sum embeddings and divide by number of non-padding tokens
+ summed_embeddings = torch.sum(
+ masked_embeddings, dim=1
+ ) # Shape: (batch_size, hidden_size)
+ summed_mask = torch.clamp(
+ attention_mask_expanded.sum(dim=1), min=1e-9
+ ) # Avoid division by zero
+
+ # Create sentence representation by mean pooling
+ mean_pooled = summed_embeddings / summed_mask
+
+ return mean_pooled
+
+
+# Generated using Mistral
+def mean_pool(string: str, tokenizer, encoder) -> torch.Tensor:
+ inputs = tokenizer(string.lower(), return_tensors="pt")
+ outputs = encoder(**inputs)
+
+ token_embeddings = outputs.last_hidden_state
+ attention_mask = inputs["attention_mask"]
+
+ # Expand attention mask to match token embeddings dimensions
+ attention_mask_expanded = (
+ attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+ )
+
+ # Apply mask to token embeddings (set padding tokens to 0)
+ masked_embeddings = token_embeddings * attention_mask_expanded
+
+ # Sum embeddings and divide by number of non-padding tokens
+ summed_embeddings = torch.sum(
+ masked_embeddings, dim=1
+ ) # Shape: (batch_size, hidden_size)
+ summed_mask = torch.clamp(
+ attention_mask_expanded.sum(dim=1), min=1e-9
+ ) # Avoid division by zero
+
+ # Create sentence representation by mean pooling
+ mean_pooled = summed_embeddings / summed_mask
+
+ return mean_pooled
+
+
+# Generated using Mistral
+def cls_sep_pool(string: str, tokenizer, encoder) -> torch.Tensor:
+ inputs = tokenizer(string.lower(), return_tensors="pt")
+ outputs = encoder(**inputs)
+ token_embeddings = outputs.last_hidden_state
+ attention_mask = inputs["attention_mask"]
+ input_ids = inputs["input_ids"]
+
+ batch_size, seq_len, hidden_size = token_embeddings.shape
+
+ # Initialize result tensor
+ cls_sep_embeddings = torch.zeros(batch_size, hidden_size)
+
+ for batch_idx in range(batch_size):
+ # Find the sequence length (excluding padding)
+ valid_length = attention_mask[batch_idx].sum().item()
+
+ # CLS token is always at position 0
+ cls_embedding = token_embeddings[batch_idx, 0]
+
+ # Find SEP token position by looking for the SEP token ID
+ # SEP token is the last non-padding token before padding starts
+ sep_token_id = tokenizer.sep_token_id
+ sep_position = None
+
+ # Search for SEP token in the valid sequence
+ for pos in range(valid_length):
+ if input_ids[batch_idx, pos] == sep_token_id:
+ sep_position = pos
+ break
+
+ # If no SEP token found, use the last valid token
+ if sep_position is None:
+ raise Exception("No SEP token found")
+
+ sep_embedding = token_embeddings[batch_idx, sep_position]
+
+ # Sum CLS and SEP embeddings
+ cls_sep_embeddings[batch_idx] = cls_embedding + sep_embedding
+
+ return cls_sep_embeddings
+
+
+@ModelFactory("BERT-cls")
+class BERT_cls(Model):
+
+ def __init__(self) -> None:
+ self.model_name: str = "bert-base-uncased"
+ self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+ self.encoder = AutoModel.from_pretrained(self.model_name)
+ self.encoder.eval()
+
+ @functools.lru_cache
+ @torch.no_grad()
+ def embed(self, string: str) -> torch.Tensor:
+ return self.encoder(
+ self.tokenizer(string.lower(), return_tensors="pt")["input_ids"]
+ ).last_hidden_state[:, 0, :]
+
+ def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
+ return torch.nn.functional.cosine_similarity(
+ a, b
+ ).item() # dist(\mu(a), \mu(b))
+
+
+@ModelFactory("ModernBERT-cls")
+class ModernBERT_cls(Model):
+
+ def __init__(self) -> None:
+ self.model_name: str = "answerdotai/ModernBERT-base"
+ self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+ self.encoder = AutoModel.from_pretrained(self.model_name)
+ self.encoder.eval()
+
+ @functools.lru_cache
+ @torch.no_grad()
+ def embed(self, string: str) -> torch.Tensor:
+ return self.encoder(
+ self.tokenizer(string.lower(), return_tensors="pt")["input_ids"]
+ ).last_hidden_state[:, 0, :]
+
+ def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
+ return torch.nn.functional.cosine_similarity(
+ a, b
+ ).item() # dist(\mu(a), \mu(b))
+
+
+@ModelFactory("ModernBERT-mean")
+class ModernBERT_mean(Model):
+
+ def __init__(self) -> None:
+ self.model_name: str = "answerdotai/ModernBERT-base"
+ self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+ self.encoder = AutoModel.from_pretrained(self.model_name)
+ self.encoder.eval()
+
+ @functools.lru_cache
+ @torch.no_grad()
+ def embed(self, string: str) -> torch.Tensor:
+ return mean_pool_last_four_layers(
+ string=string, tokenizer=self.tokenizer, encoder=self.encoder
+ )
+
+ def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
+ return torch.nn.functional.cosine_similarity(
+ a, b
+ ).item() # dist(\mu(a), \mu(b))
+
+
+@ModelFactory("ModernBERT-clssep")
+class ModernBERT_clssep(Model):
+
+ def __init__(self) -> None:
+ self.model_name: str = "answerdotai/ModernBERT-base"
+ self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+ self.encoder = AutoModel.from_pretrained(self.model_name)
+ self.encoder.eval()
+
+ @functools.lru_cache
+ @torch.no_grad()
+ def embed(self, string: str) -> torch.Tensor:
+ return cls_sep_pool(
+ string=string, tokenizer=self.tokenizer, encoder=self.encoder
+ )
+
+ def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
+ return torch.nn.functional.cosine_similarity(
+ a, b
+ ).item() # dist(\mu(a), \mu(b))
+
+
+@ModelFactory("DeBERTa-cls")
+class DeBERTa_cls(Model):
+
+ def __init__(self) -> None:
+ self.model_name: str = "microsoft/deberta-v3-base"
+ self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+ self.encoder = AutoModel.from_pretrained(self.model_name)
+ self.encoder.eval()
+
+ @functools.lru_cache
+ @torch.no_grad()
+ def embed(self, string: str) -> torch.Tensor:
+ return self.encoder(
+ self.tokenizer(string.lower(), return_tensors="pt")["input_ids"]
+ ).last_hidden_state[:, 0, :]
+
+ def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
+ return torch.nn.functional.cosine_similarity(
+ a, b
+ ).item() # dist(\mu(a), \mu(b))
+
+
+@ModelFactory("DeBERTa-mean")
+class DeBERTa_mean(Model):
+
+ def __init__(self) -> None:
+ self.model_name: str = "microsoft/deberta-v3-base"
+ self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+ self.encoder = AutoModel.from_pretrained(self.model_name)
+ self.encoder.eval()
+
+ @functools.lru_cache
+ @torch.no_grad()
+ def embed(self, string: str) -> torch.Tensor:
+ return mean_pool_last_four_layers(
+ string=string, tokenizer=self.tokenizer, encoder=self.encoder
+ )
+
+ def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
+ return torch.nn.functional.cosine_similarity(
+ a, b
+ ).item() # dist(\mu(a), \mu(b))
+
+
+@ModelFactory("DeBERTa-clssep")
+class DeBERTa_clssep(Model):
+
+ def __init__(self) -> None:
+ self.model_name: str = "microsoft/deberta-v3-base"
+ self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+ self.encoder = AutoModel.from_pretrained(self.model_name)
+ self.encoder.eval()
+
+ @functools.lru_cache
+ @torch.no_grad()
+ def embed(self, string: str) -> torch.Tensor:
+ return cls_sep_pool(
+ string=string, tokenizer=self.tokenizer, encoder=self.encoder
+ )
+
+ def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
+ return torch.nn.functional.cosine_similarity(
+ a, b
+ ).item() # dist(\mu(a), \mu(b))
+
+
+@ModelFactory("BERT-mean")
+class BERT_mean(Model):
+
+ def __init__(self) -> None:
+ self.model_name: str = "bert-base-uncased"
+ self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+ self.encoder = AutoModel.from_pretrained(self.model_name)
+ self.encoder.eval()
+
+ @functools.lru_cache
+ @torch.no_grad()
+ def embed(self, string: str) -> torch.Tensor:
+ return mean_pool_last_four_layers(
+ string=string, tokenizer=self.tokenizer, encoder=self.encoder
+ )
+
+ def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
+ return torch.nn.functional.cosine_similarity(
+ a, b
+ ).item() # dist(\mu(a), \mu(b))
+
+
+@ModelFactory("BERT-clssep")
+class BERT_clssep(Model):
+
+ def __init__(self) -> None:
+ self.model_name: str = "bert-base-uncased"
+ self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+ self.encoder = AutoModel.from_pretrained(self.model_name)
+ self.encoder.eval()
+
+ @functools.lru_cache
+ @torch.no_grad()
+ def embed(self, string: str) -> torch.Tensor:
+ return cls_sep_pool(
+ string=string, tokenizer=self.tokenizer, encoder=self.encoder
+ )
+
+ def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
+ return torch.nn.functional.cosine_similarity(
+ a, b
+ ).item() # dist(\mu(a), \mu(b))
+
+
+@ModelFactory("S-BERT")
+class S_BERT(Model):
+
+ def __init__(self):
+ self.encoder = SentenceTransformer("all-MiniLM-L6-v2", device=DEVICE)
+
+ @functools.lru_cache
+ def embed(self, string: str) -> torch.Tensor:
+ return self.encoder.encode([string])
+
+ def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
+ return self.encoder.similarity(a, b).item()
+
+
+@ModelFactory("S-GPT")
+class S_GPT(Model):
+
+ def __init__(self):
+ self.encoder = SentenceTransformer(
+ "Muennighoff/SGPT-125M-weightedmean-nli-bitfit", device=DEVICE
+ )
+
+ @functools.lru_cache
+ def embed(self, string: str) -> torch.Tensor:
+ return self.encoder.encode([string])
+
+ def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
+ return self.encoder.similarity(a, b).item()
+
+
+@ModelFactory("E5")
+class E5(Model):
+
+ def __init__(self):
+ self.encoder = SentenceTransformer(
+ "intfloat/multilingual-e5-large-instruct", device=DEVICE
+ )
+ self.task = "Given a sentence, encode its semantic properties"
+
+ def get_detailed_instruct(self, query: str) -> str:
+ return f"Instruct: {self.task}\nQuery: {query}"
+
+ @functools.lru_cache
+ def embed(self, string: str) -> torch.Tensor:
+ prepared_string = self.get_detailed_instruct(string)
+ return self.encoder.encode([prepared_string])
+
+ def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
+ return self.encoder.similarity(a, b).item()
+
+
+@ModelFactory("QWEN")
+class QWEN(Model):
+
+ def __init__(self):
+ self.encoder = SentenceTransformer("Qwen/Qwen3-Embedding-0.6B", device=DEVICE)
+ self.task = "Given a sentence, encode its semantic properties"
+
+ def get_detailed_instruct(self, query: str) -> str:
+ return f"Instruct: {self.task}\nQuery: {query}"
+
+ @functools.lru_cache
+ def embed(self, string: str) -> torch.Tensor:
+ prepared_string = self.get_detailed_instruct(string)
+ return self.encoder.encode([prepared_string])
+
+ def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
+ return self.encoder.similarity(a, b).item()
+
+
+@ModelFactory("UAE")
+class UAE(Model):
+
+ def __init__(self):
+ self.encoder = SentenceTransformer("WhereIsAI/UAE-Large-V1", device=DEVICE)
+
+ @functools.lru_cache
+ def embed(self, string: str) -> torch.Tensor:
+ return self.encoder.encode(string)
+
+ def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
+ return self.encoder.similarity(a, b).item()
diff --git a/experiment.py b/experiment.py
@@ -14,7 +14,7 @@ import tqdm
from transformers import AutoTokenizer, AutoModel
from sentence_transformers import SentenceTransformer
-from models import ModelFactory, Model
+from epsilon_compositionality.models import ModelFactory, Model
def parse_args() -> argparse.Namespace:
diff --git a/models.py b/models.py
@@ -1,421 +0,0 @@
-from __future__ import annotations
-from typing import NamedTuple, Any
-from abc import ABC, abstractmethod
-from abc import ABC, abstractmethod
-from transformers import AutoTokenizer, AutoModel
-from sentence_transformers import SentenceTransformer
-import torch
-
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-
-
-class Model(ABC):
-
- @abstractmethod
- def embed(self, string: str) -> torch.Tensor:
- pass
-
- @abstractmethod
- def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
- pass
-
-
-class ModelFactory:
-
- _factory: dict[str, type[Model]] = {}
-
- def __init__(self, name: str):
- self._name = name
-
- def __call__(self, model: type[Model]) -> type[Model]:
- self._factory[self._name] = model
- return model
-
- @classmethod
- def get(cls: type[ModelFactory], name: str, *args: Any, **kwargs: Any) -> Model:
- return cls._factory[name](*args, **kwargs)
-
-
-# Generated using Mistral
-def mean_pool_last_four_layers(string: str, tokenizer, encoder) -> torch.Tensor:
- inputs = tokenizer(string.lower(), return_tensors="pt")
-
- # Get all hidden states (set output_hidden_states=True)
- outputs = encoder(**inputs, output_hidden_states=True)
-
- # Get the last four hidden states (layers)
- # hidden_states is a tuple of length num_layers+1 (includes embedding layer)
- # We take layers [-4:] to get the last four layers (excluding embedding layer)
- last_four_layers = outputs.hidden_states[-4:]
-
- # Stack the last four layers and average them
- stacked_layers = torch.stack(
- last_four_layers, dim=0
- ) # Shape: (4, batch_size, seq_len, hidden_size)
- averaged_layers = torch.mean(
- stacked_layers, dim=0
- ) # Shape: (batch_size, seq_len, hidden_size)
-
- attention_mask = inputs["attention_mask"]
-
- # Expand attention mask to match token embeddings dimensions
- attention_mask_expanded = (
- attention_mask.unsqueeze(-1).expand(averaged_layers.size()).float()
- )
-
- # Apply mask to token embeddings (set padding tokens to 0)
- masked_embeddings = averaged_layers * attention_mask_expanded
-
- # Sum embeddings and divide by number of non-padding tokens
- summed_embeddings = torch.sum(
- masked_embeddings, dim=1
- ) # Shape: (batch_size, hidden_size)
- summed_mask = torch.clamp(
- attention_mask_expanded.sum(dim=1), min=1e-9
- ) # Avoid division by zero
-
- # Create sentence representation by mean pooling
- mean_pooled = summed_embeddings / summed_mask
-
- return mean_pooled
-
-
-# Generated using Mistral
-def mean_pool(string: str, tokenizer, encoder) -> torch.Tensor:
- inputs = tokenizer(string.lower(), return_tensors="pt")
- outputs = encoder(**inputs)
-
- token_embeddings = outputs.last_hidden_state
- attention_mask = inputs["attention_mask"]
-
- # Expand attention mask to match token embeddings dimensions
- attention_mask_expanded = (
- attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
- )
-
- # Apply mask to token embeddings (set padding tokens to 0)
- masked_embeddings = token_embeddings * attention_mask_expanded
-
- # Sum embeddings and divide by number of non-padding tokens
- summed_embeddings = torch.sum(
- masked_embeddings, dim=1
- ) # Shape: (batch_size, hidden_size)
- summed_mask = torch.clamp(
- attention_mask_expanded.sum(dim=1), min=1e-9
- ) # Avoid division by zero
-
- # Create sentence representation by mean pooling
- mean_pooled = summed_embeddings / summed_mask
-
- return mean_pooled
-
-
-# Generated using Mistral
-def cls_sep_pool(string: str, tokenizer, encoder) -> torch.Tensor:
- inputs = tokenizer(string.lower(), return_tensors="pt")
- outputs = encoder(**inputs)
- token_embeddings = outputs.last_hidden_state
- attention_mask = inputs["attention_mask"]
- input_ids = inputs["input_ids"]
-
- batch_size, seq_len, hidden_size = token_embeddings.shape
-
- # Initialize result tensor
- cls_sep_embeddings = torch.zeros(batch_size, hidden_size)
-
- for batch_idx in range(batch_size):
- # Find the sequence length (excluding padding)
- valid_length = attention_mask[batch_idx].sum().item()
-
- # CLS token is always at position 0
- cls_embedding = token_embeddings[batch_idx, 0]
-
- # Find SEP token position by looking for the SEP token ID
- # SEP token is the last non-padding token before padding starts
- sep_token_id = tokenizer.sep_token_id
- sep_position = None
-
- # Search for SEP token in the valid sequence
- for pos in range(valid_length):
- if input_ids[batch_idx, pos] == sep_token_id:
- sep_position = pos
- break
-
- # If no SEP token found, use the last valid token
- if sep_position is None:
- raise Exception("No SEP token found")
-
- sep_embedding = token_embeddings[batch_idx, sep_position]
-
- # Sum CLS and SEP embeddings
- cls_sep_embeddings[batch_idx] = cls_embedding + sep_embedding
-
- return cls_sep_embeddings
-
-
-@ModelFactory("BERT-cls")
-class BERT_cls(Model):
-
- def __init__(self) -> None:
- self.model_name: str = "bert-base-uncased"
- self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
- self.encoder = AutoModel.from_pretrained(self.model_name)
- self.encoder.eval()
-
- @torch.no_grad()
- def embed(self, string: str) -> torch.Tensor:
- return self.encoder(
- self.tokenizer(string.lower(), return_tensors="pt")["input_ids"]
- ).last_hidden_state[:, 0, :]
-
- def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
- return torch.nn.functional.cosine_similarity(
- a, b
- ).item() # dist(\mu(a), \mu(b))
-
-
-@ModelFactory("ModernBERT-cls")
-class ModernBERT_cls(Model):
-
- def __init__(self) -> None:
- self.model_name: str = "answerdotai/ModernBERT-base"
- self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
- self.encoder = AutoModel.from_pretrained(self.model_name)
- self.encoder.eval()
-
- @torch.no_grad()
- def embed(self, string: str) -> torch.Tensor:
- return self.encoder(
- self.tokenizer(string.lower(), return_tensors="pt")["input_ids"]
- ).last_hidden_state[:, 0, :]
-
- def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
- return torch.nn.functional.cosine_similarity(
- a, b
- ).item() # dist(\mu(a), \mu(b))
-
-
-@ModelFactory("ModernBERT-mean")
-class ModernBERT_mean(Model):
-
- def __init__(self) -> None:
- self.model_name: str = "answerdotai/ModernBERT-base"
- self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
- self.encoder = AutoModel.from_pretrained(self.model_name)
- self.encoder.eval()
-
- @torch.no_grad()
- def embed(self, string: str) -> torch.Tensor:
- return mean_pool_last_four_layers(
- string=string, tokenizer=self.tokenizer, encoder=self.encoder
- )
-
- def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
- return torch.nn.functional.cosine_similarity(
- a, b
- ).item() # dist(\mu(a), \mu(b))
-
-
-@ModelFactory("ModernBERT-clssep")
-class ModernBERT_clssep(Model):
-
- def __init__(self) -> None:
- self.model_name: str = "answerdotai/ModernBERT-base"
- self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
- self.encoder = AutoModel.from_pretrained(self.model_name)
- self.encoder.eval()
-
- @torch.no_grad()
- def embed(self, string: str) -> torch.Tensor:
- return cls_sep_pool(
- string=string, tokenizer=self.tokenizer, encoder=self.encoder
- )
-
- def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
- return torch.nn.functional.cosine_similarity(
- a, b
- ).item() # dist(\mu(a), \mu(b))
-
-
-@ModelFactory("DeBERTa-cls")
-class DeBERTa_cls(Model):
-
- def __init__(self) -> None:
- self.model_name: str = "microsoft/deberta-v3-base"
- self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
- self.encoder = AutoModel.from_pretrained(self.model_name)
- self.encoder.eval()
-
- @torch.no_grad()
- def embed(self, string: str) -> torch.Tensor:
- return self.encoder(
- self.tokenizer(string.lower(), return_tensors="pt")["input_ids"]
- ).last_hidden_state[:, 0, :]
-
- def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
- return torch.nn.functional.cosine_similarity(
- a, b
- ).item() # dist(\mu(a), \mu(b))
-
-
-@ModelFactory("DeBERTa-mean")
-class DeBERTa_mean(Model):
-
- def __init__(self) -> None:
- self.model_name: str = "microsoft/deberta-v3-base"
- self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
- self.encoder = AutoModel.from_pretrained(self.model_name)
- self.encoder.eval()
-
- @torch.no_grad()
- def embed(self, string: str) -> torch.Tensor:
- return mean_pool_last_four_layers(
- string=string, tokenizer=self.tokenizer, encoder=self.encoder
- )
-
- def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
- return torch.nn.functional.cosine_similarity(
- a, b
- ).item() # dist(\mu(a), \mu(b))
-
-
-@ModelFactory("DeBERTa-clssep")
-class DeBERTa_clssep(Model):
-
- def __init__(self) -> None:
- self.model_name: str = "microsoft/deberta-v3-base"
- self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
- self.encoder = AutoModel.from_pretrained(self.model_name)
- self.encoder.eval()
-
- @torch.no_grad()
- def embed(self, string: str) -> torch.Tensor:
- return cls_sep_pool(
- string=string, tokenizer=self.tokenizer, encoder=self.encoder
- )
-
- def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
- return torch.nn.functional.cosine_similarity(
- a, b
- ).item() # dist(\mu(a), \mu(b))
-
-
-@ModelFactory("BERT-mean")
-class BERT_mean(Model):
-
- def __init__(self) -> None:
- self.model_name: str = "bert-base-uncased"
- self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
- self.encoder = AutoModel.from_pretrained(self.model_name)
- self.encoder.eval()
-
- @torch.no_grad()
- def embed(self, string: str) -> torch.Tensor:
- return mean_pool_last_four_layers(
- string=string, tokenizer=self.tokenizer, encoder=self.encoder
- )
-
- def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
- return torch.nn.functional.cosine_similarity(
- a, b
- ).item() # dist(\mu(a), \mu(b))
-
-
-@ModelFactory("BERT-clssep")
-class BERT_clssep(Model):
-
- def __init__(self) -> None:
- self.model_name: str = "bert-base-uncased"
- self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
- self.encoder = AutoModel.from_pretrained(self.model_name)
- self.encoder.eval()
-
- @torch.no_grad()
- def embed(self, string: str) -> torch.Tensor:
- return cls_sep_pool(
- string=string, tokenizer=self.tokenizer, encoder=self.encoder
- )
-
- def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
- return torch.nn.functional.cosine_similarity(
- a, b
- ).item() # dist(\mu(a), \mu(b))
-
-
-@ModelFactory("S-BERT")
-class S_BERT(Model):
-
- def __init__(self):
- self.encoder = SentenceTransformer("all-MiniLM-L6-v2", device=DEVICE)
-
- def embed(self, string: str) -> torch.Tensor:
- return self.encoder.encode([string])
-
- def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
- return self.encoder.similarity(a, b).item()
-
-
-@ModelFactory("S-GPT")
-class S_GPT(Model):
-
- def __init__(self):
- self.encoder = SentenceTransformer(
- "Muennighoff/SGPT-125M-weightedmean-nli-bitfit", device=DEVICE
- )
-
- def embed(self, string: str) -> torch.Tensor:
- return self.encoder.encode([string])
-
- def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
- return self.encoder.similarity(a, b).item()
-
-
-@ModelFactory("E5")
-class E5(Model):
-
- def __init__(self):
- self.encoder = SentenceTransformer(
- "intfloat/multilingual-e5-large-instruct", device=DEVICE
- )
- self.task = "Given a sentence, encode its semantic properties"
-
- def get_detailed_instruct(self, query: str) -> str:
- return f"Instruct: {self.task}\nQuery: {query}"
-
- def embed(self, string: str) -> torch.Tensor:
- prepared_string = self.get_detailed_instruct(string)
- return self.encoder.encode([prepared_string])
-
- def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
- return self.encoder.similarity(a, b).item()
-
-
-@ModelFactory("QWEN")
-class QWEN(Model):
-
- def __init__(self):
- self.encoder = SentenceTransformer("Qwen/Qwen3-Embedding-0.6B", device=DEVICE)
- self.task = "Given a sentence, encode its semantic properties"
-
- def get_detailed_instruct(self, query: str) -> str:
- return f"Instruct: {self.task}\nQuery: {query}"
-
- def embed(self, string: str) -> torch.Tensor:
- prepared_string = self.get_detailed_instruct(string)
- return self.encoder.encode([prepared_string])
-
- def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
- return self.encoder.similarity(a, b).item()
-
-
-@ModelFactory("UAE")
-class UAE(Model):
-
- def __init__(self):
- self.encoder = SentenceTransformer("WhereIsAI/UAE-Large-V1", device=DEVICE)
-
- def embed(self, string: str) -> torch.Tensor:
- return self.encoder.encode(string)
-
- def calculate_similarity(self, a: torch.Tensor, b: torch.Tensor) -> float:
- return self.encoder.similarity(a, b).item()
