| import torch |
| from torch import nn |
| from transformers import PretrainedConfig, PreTrainedModel |
|
|
|
|
| class ParvusConfig(PretrainedConfig): |
| model_type = "parvus" |
|
|
| def __init__(self, vocab_size=26, **kwargs): |
| self.vocab_size = vocab_size |
| super().__init__(**kwargs) |
|
|
|
|
| class ParvusModel(PreTrainedModel): |
| """One parameter. 26 letters. No preferences yet.""" |
| config_class = ParvusConfig |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| self.w = nn.Parameter(torch.zeros(1)) |
|
|
| def forward(self, input_ids=None, **kwargs): |
| batch = input_ids.shape[0] if input_ids is not None else 1 |
| |
| |
| logits = self.w * torch.ones(batch, 1, self.config.vocab_size) |
| return logits |
|
|
