| from __future__ import annotations
|
|
|
|
|
| import os
|
| import torch
|
| import torch.nn.functional as F
|
|
|
|
|
| from huggingface_hub import hf_hub_download
|
| from transformers import PreTrainedModel
|
| from transformers.modeling_outputs import CausalLMOutput
|
|
|
|
|
| try:
|
| from .configuration_transwave import TransWaveConfig
|
| from .transwave_core import (
|
| CTWModel,
|
| BYTE_VOCAB,
|
| BYTE_DIM,
|
| BYTES_PER_BLOCK,
|
| BYTE_DECODER_DIM,
|
| )
|
| except ImportError:
|
| from configuration_transwave import TransWaveConfig
|
| from transwave_core import (
|
| CTWModel,
|
| BYTE_VOCAB,
|
| BYTE_DIM,
|
| BYTES_PER_BLOCK,
|
| BYTE_DECODER_DIM,
|
| )
|
|
|
|
|
|
|
|
|
| EXPECTED_PARAMS = 96_130_040
|
|
|
|
|
|
|
|
|
| class TransWaveForCausalLM(PreTrainedModel):
|
| config_class = TransWaveConfig
|
| base_model_prefix = "transwave"
|
| main_input_name = "input_ids"
|
| _tied_weights_keys = ["transwave.decoder.shared_embedding.weight"]
|
|
|
|
|
| def __init__(self, config):
|
| super().__init__(config)
|
| self.transwave = CTWModel()
|
|
|
|
|
|
|
|
|
| def get_input_embeddings(self):
|
| return self.transwave.tokenizer.byte_embedding
|
|
|
|
|
| def set_input_embeddings(self, value):
|
| self.transwave.tokenizer.byte_embedding = value
|
| self.transwave.decoder.shared_embedding = value
|
|
|
|
|
| def get_output_embeddings(self):
|
| return self.transwave.decoder.shared_embedding
|
|
|
|
|
| @classmethod
|
| def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
|
| config = kwargs.pop("config", None)
|
| revision = kwargs.pop("revision", None)
|
| cache_dir = kwargs.pop("cache_dir", None)
|
| token = kwargs.pop("token", None)
|
| local_files_only = kwargs.pop("local_files_only", False)
|
| torch_dtype = kwargs.pop("torch_dtype", None)
|
| kwargs.pop("trust_remote_code", None)
|
| kwargs.pop("use_safetensors", None)
|
|
|
|
|
| if config is None:
|
| config = TransWaveConfig.from_pretrained(
|
| pretrained_model_name_or_path,
|
| revision=revision,
|
| cache_dir=cache_dir,
|
| token=token,
|
| local_files_only=local_files_only,
|
| )
|
|
|
|
|
| model = cls(config)
|
|
|
|
|
| path = os.fspath(pretrained_model_name_or_path)
|
| if os.path.isdir(path):
|
| ckpt_path = os.path.join(path, "best_model.pt")
|
| else:
|
| ckpt_path = hf_hub_download(
|
| repo_id=path,
|
| filename="best_model.pt",
|
| revision=revision,
|
| cache_dir=cache_dir,
|
| token=token,
|
| local_files_only=local_files_only,
|
| )
|
|
|
|
|
| checkpoint = torch.load(
|
| ckpt_path,
|
| map_location="cpu",
|
| weights_only=False,
|
| )
|
| state = checkpoint["model"] if isinstance(checkpoint, dict) and "model" in checkpoint else checkpoint
|
|
|
|
|
| incompatible = model.transwave.load_state_dict(state, strict=True)
|
| if incompatible.missing_keys or incompatible.unexpected_keys:
|
| raise RuntimeError(
|
| f"Checkpoint mismatch. Missing={incompatible.missing_keys}, "
|
| f"unexpected={incompatible.unexpected_keys}"
|
| )
|
|
|
|
|
| n_params = sum(p.numel() for p in model.transwave.parameters())
|
| if n_params != EXPECTED_PARAMS:
|
| raise RuntimeError(f"Parameter count {n_params:,} != {EXPECTED_PARAMS:,}")
|
|
|
|
|
| if torch_dtype not in (None, "auto"):
|
| model = model.to(dtype=torch_dtype)
|
|
|
|
|
| model.eval()
|
| return model
|
|
|
|
|
| def nll(self, input_ids, checkpoint_layers=False):
|
| return self.transwave.nll(
|
| input_ids,
|
| checkpoint_layers=checkpoint_layers,
|
| )
|
|
|
|
|
| def contexts(self, input_ids, checkpoint_layers=False):
|
| return self.transwave.contexts(
|
| input_ids,
|
| checkpoint_layers=checkpoint_layers,
|
| )
|
|
|
|
|
| def teacher_forced_logits(self, input_ids):
|
| context, target_blocks = self.transwave.contexts(
|
| input_ids,
|
| checkpoint_layers=False,
|
| )
|
| decoder = self.transwave.decoder
|
| batch, blocks, _ = context.shape
|
| n = batch * blocks
|
|
|
|
|
| hidden = torch.tanh(
|
| decoder.context_proj(context)
|
| ).reshape(n, BYTE_DECODER_DIM)
|
|
|
|
|
| previous = decoder.bos[None, :].expand(n, -1)
|
| position_logits = []
|
|
|
|
|
| for r in range(BYTES_PER_BLOCK):
|
| decoder_input = previous + decoder.byte_position[r][None, :]
|
| hidden = decoder.gru(decoder_input, hidden)
|
| z = decoder.to_byte_space(decoder.norm(hidden))
|
| logits = F.linear(z, decoder.shared_embedding.weight)
|
| position_logits.append(
|
| logits.reshape(batch, blocks, BYTE_VOCAB)
|
| )
|
| target = target_blocks[:, :, r].reshape(-1).long()
|
| previous = decoder.shared_embedding(target)
|
|
|
|
|
| logits = torch.stack(position_logits, dim=2)
|
| return logits.reshape(
|
| batch,
|
| blocks * BYTES_PER_BLOCK,
|
| BYTE_VOCAB,
|
| )
|
|
|
|
|
| def forward(
|
| self,
|
| input_ids=None,
|
| labels=None,
|
| return_dict=None,
|
| **kwargs,
|
| ):
|
| if input_ids is None:
|
| raise ValueError("input_ids are required")
|
| if input_ids.ndim != 2:
|
| raise ValueError("input_ids must have shape [batch, tokens]")
|
| if input_ids.shape[1] % BYTES_PER_BLOCK != 0:
|
| raise ValueError("TransWave scoring length must be divisible by 8")
|
|
|
|
|
| if labels is not None:
|
| if labels.shape != input_ids.shape or not torch.equal(labels, input_ids):
|
| raise ValueError(
|
| "TransWave performs the causal shift internally; "
|
| "for exact scoring/training, labels must equal input_ids."
|
| )
|
|
|
|
|
| logits = self.teacher_forced_logits(input_ids)
|
| loss = self.transwave.nll(
|
| input_ids,
|
| checkpoint_layers=False,
|
| )
|
|
|
|
|
| if return_dict is False:
|
| return (loss, logits)
|
|
|
|
|
| return CausalLMOutput(
|
| loss=loss,
|
| logits=logits,
|
| )
|
|
|
|
|
| @torch.no_grad()
|
| def generate_bytes(
|
| self,
|
| input_ids,
|
| max_new_tokens=32,
|
| do_sample=False,
|
| temperature=1.0,
|
| top_k=None,
|
| ):
|
| if input_ids.ndim != 2 or input_ids.shape[0] != 1:
|
| raise ValueError("generate_bytes currently supports batch size 1")
|
| out = input_ids.long().clone()
|
| device = next(self.parameters()).device
|
| out = out.to(device)
|
|
|
|
|
| for _ in range(int(max_new_tokens)):
|
| total = out.shape[1]
|
| complete = (total // BYTES_PER_BLOCK) * BYTES_PER_BLOCK
|
| r = total - complete
|
|
|
|
|
| if complete == 0:
|
| block_context = self.transwave.bos_context[None, :].to(device)
|
| else:
|
| _, blocks = self.transwave.encode_completed_blocks(
|
| out[:, :complete],
|
| checkpoint_layers=False,
|
| )
|
| h, _ = self.transwave.encode_completed_blocks(
|
| out[:, :complete],
|
| checkpoint_layers=False,
|
| )
|
| block_context = h[:, -1, :]
|
|
|
|
|
| decoder = self.transwave.decoder
|
| hidden = torch.tanh(decoder.context_proj(block_context))
|
| previous = decoder.bos[None, :]
|
|
|
|
|
| current = out[:, complete:]
|
| logits = None
|
| for j in range(r + 1):
|
| decoder_input = previous + decoder.byte_position[j][None, :]
|
| hidden = decoder.gru(decoder_input, hidden)
|
| z = decoder.to_byte_space(decoder.norm(hidden))
|
| logits = F.linear(z, decoder.shared_embedding.weight)
|
| if j < r:
|
| previous = decoder.shared_embedding(current[:, j])
|
|
|
|
|
| scores = logits.float() / max(float(temperature), 1e-6)
|
| if top_k is not None and int(top_k) > 0:
|
| k = min(int(top_k), BYTE_VOCAB)
|
| threshold = torch.topk(scores, k, dim=-1).values[:, -1:]
|
| scores = scores.masked_fill(scores < threshold, float("-inf"))
|
|
|
|
|
| if do_sample:
|
| next_id = torch.multinomial(
|
| torch.softmax(scores, dim=-1),
|
| num_samples=1,
|
| )
|
| else:
|
| next_id = scores.argmax(dim=-1, keepdim=True)
|
|
|
|
|
| out = torch.cat([out, next_id], dim=1)
|
|
|
|
|
| return out |
