tsfrm/unity-embed

1"""transformers shim. AutoModel.from_pretrained(..., trust_remote_code=True)"""2import math3 4from transformers import PreTrainedModel, PretrainedConfig5 6 7class UnityEmbedConfig(PretrainedConfig):8    model_type = "unity-embed"9 10    def __init__(self, embedding_dimension=384, **kwargs):11        self.embedding_dimension = embedding_dimension12        super().__init__(**kwargs)13 14 15class UnityEmbedModel(PreTrainedModel):16    config_class = UnityEmbedConfig17 18    def __init__(self, config):19        super().__init__(config)20        import torch21        d = config.embedding_dimension22        self.v = torch.nn.Parameter(torch.full((d,), 1.0 / math.sqrt(d)))23 24    def forward(self, input_ids=None, attention_mask=None, **kw):25        import torch26        v = self.v / self.v.norm()27        if input_ids is not None:28            return v.expand(input_ids.shape[0], v.shape[0]).contiguous()29        return v30