File size: 1,719 Bytes
ef8e803 | 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 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | """A minimal MLP with tanh activations, packaged as a HuggingFace custom model.
Self-contained on purpose: only imports from torch / transformers so that it can
be downloaded and executed by `AutoModel.from_pretrained(..., trust_remote_code=True)`
without any project-local dependency.
"""
import torch
import torch.nn as nn
from transformers import PreTrainedModel
from transformers.modeling_outputs import BaseModelOutput
from .configuration_mlp import MLPConfig
class MLPModel(PreTrainedModel):
# `config_class` wires the model into the AutoModel registry.
config_class = MLPConfig
def __init__(self, config: MLPConfig):
super().__init__(config)
dims = (
[config.input_dim]
+ [config.hidden_dim] * config.num_hidden_layers
+ [config.output_dim]
)
layers = []
for i in range(len(dims) - 1):
layers.append(nn.Linear(dims[i], dims[i + 1]))
# tanh after every layer except the final projection.
if i < len(dims) - 2:
layers.append(nn.Tanh())
self.mlp = nn.Sequential(*layers)
# Standard HF weight init hook.
self.post_init()
def _init_weights(self, module):
if isinstance(module, nn.Linear):
module.weight.data.normal_(mean=0.0, std=0.02)
if module.bias is not None:
module.bias.data.zero_()
def forward(self, x: torch.Tensor = None, **kwargs) -> BaseModelOutput:
# Accept `input_ids` as an alias so generic HF tooling doesn't choke.
if x is None:
x = kwargs.get("input_ids")
output = self.mlp(x)
return BaseModelOutput(last_hidden_state=output)
|