| |
|
|
| import os |
| import torch |
| from torch import nn, Tensor |
| from transformers import AutoModel, AutoConfig |
| from huggingface_hub import snapshot_download |
| from typing import Dict |
|
|
|
|
| class BGEM3InferenceModel(nn.Module): |
| def __init__( |
| self, |
| model_name: str = "BAAI/bge-m3", |
| colbert_dim: int = -1, |
| ) -> None: |
| super().__init__() |
|
|
| model_name = snapshot_download( |
| repo_id=model_name, |
| allow_patterns=[ |
| "model.safetensors", |
| "colbert_linear.pt", |
| "sparse_linear.pt", |
| "config.json", |
| ], |
| ) |
|
|
| self.config = AutoConfig.from_pretrained(model_name) |
| self.model = AutoModel.from_pretrained(model_name) |
| self.colbert_linear = torch.nn.Linear( |
| in_features=self.model.config.hidden_size, |
| out_features=( |
| self.model.config.hidden_size if colbert_dim == -1 else colbert_dim |
| ), |
| ) |
| self.sparse_linear = torch.nn.Linear( |
| in_features=self.model.config.hidden_size, out_features=1 |
| ) |
| colbert_state_dict = torch.load( |
| os.path.join(model_name, "colbert_linear.pt"), map_location="cpu" |
| ) |
| sparse_state_dict = torch.load( |
| os.path.join(model_name, "sparse_linear.pt"), map_location="cpu" |
| ) |
| self.colbert_linear.load_state_dict(colbert_state_dict) |
| self.sparse_linear.load_state_dict(sparse_state_dict) |
|
|
| def dense_embedding(self, last_hidden_state: Tensor) -> Tensor: |
| return last_hidden_state[:, 0] |
|
|
| def sparse_embedding(self, last_hidden_state: Tensor) -> Tensor: |
| with torch.no_grad(): |
| return torch.relu(self.sparse_linear(last_hidden_state)) |
|
|
| def colbert_embedding( |
| self, last_hidden_state: Tensor, attention_mask: Tensor |
| ) -> Tensor: |
| with torch.no_grad(): |
| colbert_vecs = self.colbert_linear(last_hidden_state[:, 1:]) |
| colbert_vecs = colbert_vecs * attention_mask[:, 1:][:, :, None].float() |
| return colbert_vecs |
|
|
| def forward(self, input_ids: Tensor, attention_mask: Tensor) -> Dict[str, Tensor]: |
| with torch.no_grad(): |
| last_hidden_state = self.model( |
| input_ids=input_ids, attention_mask=attention_mask, return_dict=True |
| ).last_hidden_state |
|
|
| output = {} |
| dense_vecs = self.dense_embedding(last_hidden_state) |
| output["dense_vecs"] = torch.nn.functional.normalize(dense_vecs, dim=-1) |
|
|
| sparse_vecs = self.sparse_embedding(last_hidden_state) |
| output["sparse_vecs"] = sparse_vecs |
|
|
| colbert_vecs = self.colbert_embedding(last_hidden_state, attention_mask) |
| output["colbert_vecs"] = torch.nn.functional.normalize(colbert_vecs, dim=-1) |
|
|
| return output |
|
|
