Upload model

config.json

@@ -0,0 +1,19 @@
+{
+  "architectures": [
+    "MapperModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_mapper.MapperConfig",
+    "AutoModel": "modeling_mapper.MapperModel"
+  },
+  "d_hidden": 1024,
+  "d_in": 64,
+  "d_out": 64,
+  "dropout": 0.1,
+  "layer_norm_eps": 1e-12,
+  "model_type": "embedding_mapper",
+  "n_layers": 6,
+  "n_out": 2,
+  "torch_dtype": "float32",
+  "transformers_version": "4.38.2"
+}

configuration_mapper.py

@@ -0,0 +1,25 @@
+from transformers import PretrainedConfig
+
+class MapperConfig(PretrainedConfig):
+    model_type = "embedding_mapper"
+
+    def __init__(self,
+                 d_in: int = 64,
+                 d_hidden: int = 1024,
+                 n_layers: int = 6,
+                 d_out: int = 64,
+                 n_out: int = 2,
+                 dropout: float = 0.1,
+                 layer_norm_eps: float = 1e-12,
+                 **kwargs
+                ):
+        
+        self.d_in = d_in
+        self.d_hidden = d_hidden
+        self.n_layers = n_layers
+        self.d_out = d_out
+        self.n_out = n_out
+        self.dropout = dropout
+        self.layer_norm_eps = layer_norm_eps
+        
+        super().__init__(**kwargs)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5df4171738f7d9ec9a5451ecbae4a107b16af59f442d48718f4ca5ff41fdb403
+size 82271096

modeling_mapper.py

@@ -0,0 +1,76 @@
+import torch
+import torch.nn as nn 
+import torch.nn.functional as F 
+
+from dataclasses import dataclass
+
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.utils import ModelOutput
+
+from .configuration_mapper import MapperConfig
+
+class FeedForward(nn.Module):
+    def __init__(self, d_in, d_out):
+        super().__init__()
+        self.fc1 = nn.Linear(d_in, d_out*2)
+        self.fc2 = nn.Linear(d_out, d_out)
+        
+    def forward(self, x):
+        x = self.fc1(x)
+        x1, x2 = x.chunk(2, dim=-1)
+        x = self.fc2(F.silu(x1) * x2)
+        return x
+
+class FeedForwardLayer(nn.Module):
+    def __init__(self, d_in, d_out, dropout=0.1, layer_norm_eps=None):
+        super().__init__()
+        self.ff = FeedForward(d_in, d_out)
+        self.skip = nn.Linear(d_in, d_out) if d_in != d_out else nn.Identity()
+        self.dropout = nn.Dropout(dropout)
+        self.LayerNorm = nn.LayerNorm(d_out, eps=layer_norm_eps) if layer_norm_eps else None
+        
+    def forward(self, x):
+        x = self.dropout(x)
+        x = self.ff(x) + self.skip(x)
+        if self.LayerNorm:
+            x = self.LayerNorm(x)
+        return x
+    
+class Mapper(nn.Module):
+    def __init__(self, d_in, d_hidden, d_out, n_out, n_layers, dropout=0.1, layer_norm_eps=None):
+        super().__init__()
+        self.n_out = n_out
+        layers = [FeedForwardLayer(d_in, d_hidden, 0.0, layer_norm_eps)]
+        layers += [FeedForwardLayer(d_hidden, d_hidden, dropout, layer_norm_eps)
+                   for i in range(n_layers)]
+        self.layers = nn.Sequential(*layers)
+    
+        self.output_layer = FeedForwardLayer(d_hidden, d_out*n_out, 0.0, None)
+        
+    def forward(self, x):
+        x = self.layers(x)
+        x = self.output_layer(x)
+        x = torch.stack(torch.chunk(x, self.n_out, -1), 1)
+        return x
+
+@dataclass
+class MapperModelOutput(ModelOutput):
+    mapper_out: torch.FloatTensor = None
+
+class MapperModel(PreTrainedModel):
+    config_class = MapperConfig
+    def __init__(self, config):
+        super().__init__(config)
+        
+        self.mapper = Mapper(config.d_in, config.d_hidden, config.d_out, config.n_out,
+                            config.n_layers, config.dropout, config.layer_norm_eps)
+        
+    def forward(self, embedding, return_dict=True):
+                        
+        mapper_out = self.mapper(embedding)
+        
+        
+        if not return_dict:
+            return (mapper_out, )
+        
+        return MapperModelOutput(mapper_out=mapper_out)