"""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()