| """transformers shim. AutoModel.from_pretrained(..., trust_remote_code=True)""" | |
| import math | |
| from transformers import PreTrainedModel, PretrainedConfig | |
| class UnityEmbedConfig(PretrainedConfig): | |
| model_type = "unity-embed" | |
| def __init__(self, embedding_dimension=384, **kwargs): | |
| self.embedding_dimension = embedding_dimension | |
| super().__init__(**kwargs) | |
| class UnityEmbedModel(PreTrainedModel): | |
| config_class = UnityEmbedConfig | |
| def __init__(self, config): | |
| super().__init__(config) | |
| import torch | |
| d = config.embedding_dimension | |
| self.v = torch.nn.Parameter(torch.full((d,), 1.0 / math.sqrt(d))) | |
| def forward(self, input_ids=None, attention_mask=None, **kw): | |
| import torch | |
| v = self.v / self.v.norm() | |
| if input_ids is not None: | |
| return v.expand(input_ids.shape[0], v.shape[0]).contiguous() | |
| return v | |