backtrack.py

from typing import Optional, Iterable
import torch
import math
from dataclasses import dataclass, field
from transformers import MistralForCausalLM, DynamicCache
from transformers.modeling_outputs import CausalLMOutputWithPast
from constrainers import Constrainer
from enum import IntEnum
from cachetools import LRUCache
from pathlib import Path

class FakeCache:
  def __init__(self, maxsize: int):
    self.maxsize = maxsize
  def get(self, key, default=None):
    return default
  def __setitem__(self, key, value):
    pass
  def clear(self):
    pass

@torch.jit.script
def log1mexp(x: torch.Tensor) -> torch.Tensor:
  """Numerically accurate evaluation of log(1 - exp(x)) for x < 0.
  See https://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf for details.
  from https://github.com/pytorch/pytorch/issues/39242#issuecomment-1349202573
  """
  mask = -math.log(2) < x  # x < 0
  return torch.where(
    mask,
    (-x.expm1()).log(),
    (-x.exp()).log1p(),
  )

@torch.jit.script
def shift_and_append(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
  z = x.roll(-1)
  z[..., -1] = y
  return z

@torch.jit.script
def reversed_cumsum(x: torch.Tensor, dim: int) -> torch.Tensor:
  return x.flip(dim).cumsum(dim).flip(dim)

@torch.jit.script
def get_fixed_logits(logits: torch.Tensor, generated_tokens: torch.LongTensor, bad_logit: torch.Tensor) -> torch.Tensor:
  assert logits.shape[0] == generated_tokens.shape[0]
  single_wrong_logit = logits.gather(dim=-1, index=generated_tokens.unsqueeze(-1)).squeeze(-1) - logits.logsumexp(dim=-1)
  single_wrong_logit = shift_and_append(single_wrong_logit, bad_logit)
  remain_logits = log1mexp(reversed_cumsum(single_wrong_logit, dim=-1))
  return logits.scatter_add(dim=-1, index=generated_tokens.unsqueeze(-1), src=remain_logits.unsqueeze(-1))

@torch.jit.script
def reject_sample(logits: torch.Tensor, new_logits: torch.Tensor, generated_tokens: torch.LongTensor) -> torch.Tensor:
  log_Z_diff = new_logits.logsumexp(dim=-1, keepdim=True) - logits.logsumexp(dim=-1, keepdim=True) # log(Z2) - log(Z1)
  original_token = logits.new_full(logits.shape, False, dtype=torch.bool)
  original_token.scatter_(dim=-1, index=generated_tokens.unsqueeze(-1), value=True)
  changed = logits != new_logits
  no_change = changed.any(dim=-1, keepdim=True).logical_not()
  original_token.logical_and_(no_change)
  changed.logical_and_(log_Z_diff <= 0)
  mask = changed.logical_or(log_Z_diff > 0).logical_or(original_token)
  result = torch.where(mask, new_logits, new_logits + log1mexp(log_Z_diff))
  assert not result.isnan().any().item()
  assert not result.isposinf().any().item()
  return result

@torch.jit.script
def get_bad_logit(logit: torch.Tensor, blocklist: list[int]) -> torch.Tensor:
  mask = torch.ones_like(logit, dtype=torch.bool)
  mask[blocklist] = False
  all_Z = torch.logsumexp(logit, dim=-1)
  block_Z = torch.logsumexp(logit[blocklist], dim=-1)
  pass_Z = torch.logsumexp(logit[mask], dim=-1)
  if block_Z < pass_Z:
    bad_logit = block_Z - all_Z
  else:
    bad_logit = log1mexp(pass_Z - all_Z)
  return bad_logit

@dataclass
class GenData:
  cache: Optional[Iterable[Iterable[torch.Tensor]]]
  central_cache: LRUCache
  tokens: torch.LongTensor
  logits: torch.Tensor

  def __len__(self):
    return self.tokens.shape[0] + 1

  def get_cache(self) -> Optional[Iterable[Iterable[torch.Tensor]]]:
    if self.cache is not None:
      return self.cache
    return self.central_cache.get(id(self), None)

  def goto_cpu(self, prompt_len: int):
    if self.cache is None:
      return
    if self.central_cache.maxsize > 0:
      cache = [[key[:, :, prompt_len:, :].cpu(), value[:, :, prompt_len:, :].cpu()] for key, value in self.cache]
      self.central_cache[id(self)] = cache
    self.cache = None
    self.logits = self.logits.cpu()

  def goto_gpu(self, prompt_cache: Iterable[Iterable[torch.Tensor]]):
    if self.cache is not None:
      return True
    device = prompt_cache[0][0].device
    self.logits = self.logits.to(device)
    cache = self.central_cache.get(id(self), None)
    if cache is None:
      return False
    self.cache = [(torch.cat([prompt_key, key.to(device)], dim=2), torch.cat([prompt_value, value.to(device)], dim=2)) for (prompt_key, prompt_value), (key, value) in zip(prompt_cache, cache)]
    return True

  def cut_cache(self, removed_length: int) -> Iterable[Iterable[torch.Tensor]]:
    cache = self.central_cache.get(id(self), None)
    if cache is None:
      return None
    seq_len = cache[0][0].shape[2]
    return [(key[:, :, :seq_len-removed_length, :].clone(), value[:, :, :seq_len-removed_length, :].clone()) for key, value in cache]

@dataclass
class GenNode:
  data: GenData # always a reference to the data in the last node
  next: Optional[int] = None
  children: dict[int, 'GenNode'] = field(default_factory=dict)

  def dump_to(self, path: Path, prompt_length: int):
    self.data.goto_cpu(prompt_length)
    self.data.central_cache.clear()
    with path.open('xb') as f:
      torch.save(self, f)

  @staticmethod
  def load_from(path: Path) -> 'GenNode':
    with path.open('rb') as f:
      return torch.load(f)

@dataclass
class GenConfig:
  temperature: float = 1.0
  top_k: int = 0 # 0 means no top-k
  top_p: float = 1.0

  def __post_init__(self):
    assert self.temperature >= 0
    assert self.top_k >= 0
    assert 0 < self.top_p <= 1
    if self.temperature == 0:
      self.temperature = 1.0
      self.top_k = 1

  @torch.no_grad()
  def apply(self, logits: torch.Tensor) -> torch.Tensor:
    was_three_dim = logits.dim() == 3
    if logits.dim() == 3:
      assert logits.shape[0] == 1
      logits = logits.squeeze(0)
    assert logits.dim() == 2
    seqlen, vocab_size = logits.shape
    assert 0 < self.top_k <= vocab_size
    logits = logits / self.temperature
    if self.top_k < vocab_size or self.top_p < 1.:
      probs = logits.softmax(dim=-1)
      selected_probs, selected_tokens = probs.topk(k=self.top_k, dim=-1, largest=True, sorted=True)
      mask = selected_probs.cumsum(dim=-1) > self.top_p
      mask = mask.roll(1, dims=-1)
      mask[:, 0] = False
      selected_tokens.masked_fill_(mask, vocab_size)
      logit_mask = logits.new_full((seqlen, vocab_size + 1), True, dtype=torch.bool)
      logit_mask.scatter_(dim=-1, index=selected_tokens, value=False)
      logit_mask = logit_mask[:, :-1]
      logits.masked_fill_(logit_mask, float('-inf'))
    if was_three_dim:
      logits = logits.unsqueeze(0)
    return logits

  def transparent(self) -> bool:
    return self.temperature == 1.0 and self.top_k == 0 and self.top_p == 1.0

  def append_apply(self, logits: torch.Tensor) -> torch.Tensor:
    if self.transparent():
      return logits
    applied = self.apply(logits)
    return torch.cat([logits, applied], dim=0)

@dataclass
class Prompt:
  ids: torch.LongTensor
  cache: Iterable[Iterable[torch.Tensor]]
  logit: torch.Tensor

  @property
  def length(self):
    return self.ids.shape[0]

@dataclass
class GenTree:
  path: list[GenNode]
  data: GenData
  prompt: Prompt
  model: MistralForCausalLM
  model_call_count: int
  central_cache: LRUCache
  constrainer: Constrainer
  config: GenConfig
  approximate_p: float

  @staticmethod
  def get_blocklist(constrainer: Constrainer, context: list[int], logits: torch.Tensor, approximate_p: float) -> list[int]:
    if approximate_p == 1.0:
      return constrainer.get_blocklist(context)
    target = math.log(approximate_p / (1 - approximate_p))
    order = logits.argsort(descending=True)
    reordered_logit = logits[order]
    thresholds = reordered_logit.flip(-1).logcumsumexp(dim=-1).flip(-1) + target
    thresholds_list: list[float] = thresholds.tolist()
    thresholds_list = thresholds_list[1:] + [float('-inf')]
    order_list: list[int] = order.tolist()
    reordered_logit_list: list[float] = reordered_logit.tolist()
    passlist: list[int] = []
    passed_logits: list[float] = []
    current_sum = float('-inf')
    for current, logit, threshold in zip(order_list, reordered_logit_list, thresholds_list, strict=True):
      passed = constrainer.ask_token(context, current)
      if passed:
        passlist.append(current)
        passed_logits.append(logit)
        current_sum = torch.logsumexp(torch.tensor(passed_logits, device='cpu'), dim=-1).item()
      if current_sum >= threshold:
        break
    return list(set(range(len(logits))) - set(passlist))

  @staticmethod
  @torch.no_grad()
  def from_tokens(prompt: torch.LongTensor, model: MistralForCausalLM, constrainer: Constrainer, config: GenConfig, cache_size: int = 1024, prompt_info: Prompt | None = None, approximate_p: float = 1.0) -> tuple['GenTree', torch.distributions.Categorical, torch.Tensor]:
    assert prompt.dim() == 1
    if not config.transparent() and config.top_k == 0:
      config.top_k = model.config.vocab_size
    if prompt_info:
      assert prompt.shape == prompt_info.ids.shape
      assert (prompt == prompt_info.ids).all().item()
      original_logit = prompt_info.logit.clone()
      cache = [(key.clone(), value.clone()) for key, value in prompt_info.cache]
    else:
      output: CausalLMOutputWithPast = model(**model.prepare_inputs_for_generation(
        input_ids=prompt.unsqueeze(0),
        use_cache=True,
        return_dict=True
      ))
      original_logit = output.logits[0, -1, :].clone()
      assert original_logit.shape[0] == model.config.vocab_size
      cache = output.past_key_values.to_legacy_cache() if isinstance(output.past_key_values, DynamicCache) else output.past_key_values
      del output
      prompt_info = Prompt(ids=prompt, cache=[(key.clone(), value.clone()) for key, value in cache], logit=original_logit.clone())
    logit = original_logit.to(dtype=torch.float64).clone()
    blocklist = GenTree.get_blocklist(constrainer, [], logit, approximate_p)
    logit[blocklist] = float('-inf')
    logit = logit.reshape(1, 1, -1)
    logit = config.append_apply(logit)
    my_cache = LRUCache(maxsize=cache_size) if cache_size > 0 else FakeCache(maxsize=0)
    data = GenData(
      tokens=prompt.new_full((0,), fill_value=-1),
      cache=cache,
      central_cache=my_cache,
      logits=logit,
    )
    tree = GenTree(
      path=[GenNode(data=data)],
      data=data,
      prompt=prompt_info,
      model=model,
      model_call_count=1,
      central_cache=my_cache,
      constrainer=constrainer,
      config=config,
      approximate_p=approximate_p
    )
    dist = torch.distributions.Categorical(logits=logit[-1, :, :])
    return tree, dist, original_logit

  @torch.no_grad()
  def set_to_tokens(self, tokens: torch.LongTensor) -> list[int]:
    # update the underneath nodes
    for node, token in zip(self.path, tokens.tolist(), strict=True):
      node.next = token

    # get the new path
    def get_path():
      node = self.path[0]
      yield node
      while node.next in node.children:
        node = node.children[node.next]
        yield node
    new_path = list(get_path())
    assert new_path[-1].next is not None

    common_length = sum(1 for old, new in zip(self.path, new_path) if id(old) == id(new))
    if common_length != len(self.path):
      self.data.goto_cpu(self.prompt.length)
      tokens = new_path[-1].data.tokens[:len(new_path)-1]
      should_concrete = set()
      def get_logits(data: GenData):
        first, other = data.logits.split([1, data.logits.shape[1] - 1], dim=1)
        data.logits = other
        should_concrete.add(id(data))
        return first
      logits = torch.cat([get_logits(node.data) for node in new_path], dim=1)
      new_data = GenData(
        cache=None,
        central_cache=self.central_cache,
        tokens=tokens,
        logits=logits,
      )
      cache = new_path[-1].data.cut_cache(len(new_path[-1].data) - len(new_path))
      if cache is not None:
        self.central_cache[id(new_data)] = cache
      for node in new_path:
        if id(node.data) in should_concrete:
          node.data.logits = node.data.logits.clone()
          should_concrete.remove(id(node.data))
        node.data = new_data
      assert not should_concrete
      self.data = new_data
      self.path = new_path
    else:
      assert common_length == len(new_path)

    assert len(self.data) == len(self.path)

    next = self.tail.next

    return self.data.tokens.tolist() + [next]

  @torch.no_grad()
  def new_sample(self) -> torch.distributions.Categorical:
    return torch.distributions.Categorical(logits=self.data.logits[-1, :, :])

  @property
  def tail(self) -> GenNode:
    return self.path[-1]

  @property
  def vocab_size(self) -> int:
    return self.model.config.vocab_size

  @property
  def device(self) -> torch.device:
    return self.model.device

  @torch.no_grad()
  def go_on(self, backtrack: bool, backtrack_from: int = 0) -> tuple[torch.distributions.Categorical, Optional[torch.Tensor]]:
    next = self.tail.next
    assert next is not None
    new_token = self.data.tokens.new_full((1,), fill_value=next)
    new_tokens = self.data.tokens.tolist() + [next]
    self.data.tokens = torch.cat([self.data.tokens, new_token], dim=0)
    if self.data.goto_gpu(self.prompt.cache):
      # generate next token
      past_seen_tokens = self.data.get_cache()[0][0].shape[2]
      output: CausalLMOutputWithPast = self.model(**self.model.prepare_inputs_for_generation(
        input_ids=new_token.unsqueeze(0),
        use_cache=True,
        past_key_values=self.data.get_cache(),
        cache_position=self.data.tokens.new_full((1,), fill_value=past_seen_tokens),
        return_dict=True
      ))
    else:
      output: CausalLMOutputWithPast = self.model(**self.model.prepare_inputs_for_generation(
        input_ids=self.data.tokens.unsqueeze(0),
        use_cache=True,
        past_key_values=self.prompt.cache,
        cache_position=torch.arange(self.prompt.length, self.prompt.length + self.data.tokens.shape[0], device=self.device, dtype=self.prompt.ids.dtype),
        return_dict=True
      ))
    self.model_call_count += 1
    raw_logit = output.logits[0, -1, :].clone()
    new_logit = raw_logit.to(dtype=torch.float64).clone()
    blocklist = GenTree.get_blocklist(self.constrainer, new_tokens, new_logit, self.approximate_p)
    assert len(set(blocklist)) != self.vocab_size
    self.data.cache = output.past_key_values
    bad_logit = get_bad_logit(new_logit, blocklist)
    assert bad_logit.item() != 0.
    new_logit[blocklist] = float('-inf')

    # create a new node
    node = GenNode(data=self.data)
    self.tail.children[next] = node
    self.path.append(node)

    new_logit = new_logit.reshape(1, 1, -1)
    new_logit = self.config.append_apply(new_logit)

    if backtrack:
      # propagate the removed logits
      old_token_length = self.data.tokens.shape[0]
      effective_logits = self.data.logits[:, backtrack_from:, :]
      effective_tokens = self.data.tokens[backtrack_from:]
      fixed_logits = get_fixed_logits(effective_logits[0, :, :], effective_tokens, bad_logit)
      fixed_logits = fixed_logits.unsqueeze(0)
      fixed_logits = self.config.append_apply(fixed_logits)
      resample_logits = reject_sample(effective_logits[-1, :, :], fixed_logits[-1, :, :], effective_tokens)
      fixed_logits = torch.cat([self.data.logits[:, :backtrack_from, :], fixed_logits], dim=1)
    else:
      fixed_logits = self.data.logits
      resample_logits = new_logit.new_empty((0, self.vocab_size))

    resample_logits = torch.cat([resample_logits, new_logit[-1, :, :]], dim=0)
    dist = torch.distributions.Categorical(logits=resample_logits)

    # update the data
    self.data.logits = torch.cat([fixed_logits, new_logit], dim=1)
    return dist, raw_logit.cpu() if raw_logit is not None else None

  @torch.no_grad()
  def calc_likelihood(self, tokens: torch.LongTensor | list[int]) -> float:
    if torch.is_tensor(tokens):
      tokens = tokens.tolist()
    likelihood = 0.
    def get_path():
      node = self.path[0]
      yield node
      while node.next in node.children:
        node = node.children[node.next]
        yield node
    for i, (node, token) in enumerate(zip(get_path(), tokens, strict=True)):
      logit = node.data.logits[-1, :, :]
      likelihood += logit[i, token].item()
    return likelihood - self.data.logits[-1, :, :].logsumexp(dim=-1).sum().item()

class GenAction(IntEnum):
  NEW_SAMPLE = 0
  CONTINUE = 1
  CONTINUE_WITHOUT_BACKTRACK = 2

def generate(*args, longest_backtrack: int | None = None, **kwargs):
  tree, dist, logit = GenTree.from_tokens(*args, **kwargs)
  tokens = dist.sample()
  logit = logit.cpu()
  action: GenAction = yield tree, tokens.tolist()
  backtrack_from: int = 0
  yield logit
  while True:
    tokens = tree.set_to_tokens(tokens)
    action: GenAction = yield tree, tokens
    if action == GenAction.NEW_SAMPLE:
      dist = tree.new_sample()
      tokens = dist.sample()
    else:
      if longest_backtrack is not None:
        backtrack_from = max(backtrack_from, len(tokens) - longest_backtrack)
      dist, logit = tree.go_on(action == GenAction.CONTINUE, backtrack_from=backtrack_from)
      yield logit
      tokens = dist.sample()
      if action == GenAction.CONTINUE_WITHOUT_BACKTRACK:
        tokens = torch.cat([tree.data.tokens, tokens[-1].unsqueeze(0)])
      else:
        assert action == GenAction.CONTINUE
        tokens = torch.cat([tree.data.tokens[:backtrack_from], tokens])

class SequenceState(IntEnum):
  NOT_FINISHED = 0
  FINISHED = 1
  INVALID = 3

def generation_yielder(generation, criterion, set_call_count = None, backtrack: bool = True):
  INVALID = 'invalid'
  message: Optional[GenAction] = None
  backtracked = False
  old_tokens = []
  while True:
    tree, tokens = generation.send(message)
    tree: GenTree = tree
    if tokens[:len(old_tokens)] != old_tokens:
      backtracked = True
    old_tokens = tokens
    if set_call_count:
      set_call_count(tree.model_call_count)
    yield tokens
    match criterion(tokens):
      case SequenceState.FINISHED:
        yield {"tokens" : tokens, "log_likelihood" : tree.calc_likelihood(tokens), "call_counts" : tree.model_call_count, "backtracked": backtracked, "prompt": tree.prompt}
        message = GenAction.NEW_SAMPLE
        backtracked = False
        old_tokens = []
      case SequenceState.INVALID:
        yield INVALID
        message = GenAction.NEW_SAMPLE
        backtracked = False
        old_tokens = []
      case SequenceState.NOT_FINISHED:
        logit = generation.send(GenAction.CONTINUE if backtrack else GenAction.CONTINUE_WITHOUT_BACKTRACK)
        message = None
        yield logit if logit is not None else INVALID