| """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()) |
|
|
