File size: 884 Bytes
333396f | 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 | 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)) # blank slate
def forward(self, input_ids=None, **kwargs):
batch = input_ids.shape[0] if input_ids is not None else 1
# same score for every letter -> uniform distribution.
# sample from it (torch.multinomial on softmax) and you get babble.
logits = self.w * torch.ones(batch, 1, self.config.vocab_size)
return logits
|