"""transformers shim. AutoModelForCausalLM.from_pretrained(..., trust_remote_code=True)"""
from transformers import PreTrainedModel, PretrainedConfig

BYTES_VOCAB = [10] + list(range(32, 127))


class Point1Config(PretrainedConfig):
    model_type = "point1"

    def __init__(self, vocab_size=len(BYTES_VOCAB), **kwargs):
        self.vocab_size = vocab_size
        super().__init__(**kwargs)


class Point1ForCausalLM(PreTrainedModel):
    config_class = Point1Config

    def __init__(self, config):
        super().__init__(config)
        import torch
        self.w = torch.nn.Parameter(torch.tensor([0.0]))
        h = [max(-1.0, min(1.0, (b - 79) / 47.0)) for b in BYTES_VOCAB]
        self.register_buffer("h", torch.tensor(h), persistent=False)
        self.bytes_vocab = BYTES_VOCAB

    def logits(self):
        import torch
        return self.w * self.h

    def generate_bytes(self, n=64, temperature=0.8, generator=None):
        import torch
        probs = torch.softmax(self.logits() / temperature, dim=-1)
        idx = torch.multinomial(probs, num_samples=n, replacement=True, generator=generator)
        return bytes(self.bytes_vocab[i] for i in idx.tolist())