File size: 1,083 Bytes
49b416c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 | """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
# all parameters, on display together for the only time in their lives
self.v = torch.nn.Parameter(torch.full((d,), 1.0 / math.sqrt(d)))
def forward(self, input_ids=None, attention_mask=None, **kw):
"""any token sequence -> THE vector. batch dims preserved out of courtesy."""
import torch
v = self.v / self.v.norm()
shape = torch.Size([input_ids.shape[0], v.shape[0]]) if input_ids is not None else None
return v.expand(shape).contiguous()
|