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