generator.py

import cuda_ext
from model import ExLlama, ExLlamaCache
from lora import ExLlamaLora
import torch
import torch.nn.functional as F

class ExLlamaGenerator:

    class Settings:

        temperature = 0.95
        top_k = 40                              # consider the most probable top_k samples, 0 to disable top_k sampling
        top_p = 0.65                            # consider tokens up to a cumulative probabiltiy of top_p, 0.0 to disable top_p sampling
        min_p = 0.0                             # Do not consider tokens with probability less than this
        typical = 0.0                           # Locally typical sampling threshold, 0.0 to disable typical sampling

        token_repetition_penalty_max = 1.15     # Repetition penalty for most recent tokens
        token_repetition_penalty_sustain = 256  # No. most recent tokens to repeat penalty for, -1 to apply to whole context
        token_repetition_penalty_decay = 128    # Gradually decrease penalty over this many tokens

        beams = 1
        beam_length = 1


    model: ExLlama
    sequence: torch.Tensor or None
    sequence_actual: torch.Tensor or None
    settings: Settings
    beams: int or None
    max_beam_length: int
    in_beam_search: True
    disallowed_tokens: list[int] or None
    lora: ExLlamaLora or None


    def __init__(self, model, tokenizer, cache):

        self.model = model
        self.tokenizer = tokenizer
        self.cache = cache
        self.reset()


    def reset(self):

        self.cache.current_seq_len = 0
        self.sequence = None
        self.sequence_actual = None
        self.settings = ExLlamaGenerator.Settings()

        self.beams = None
        self.max_beam_length = 0
        self.in_beam_search = False
        self.disallowed_tokens = None
        self.lora = None


    def make_rep_mask(self, penalty_max, sustain, decay):

        return cuda_ext.ext_rep_penalty_mask_cpu(self.model.config.vocab_size, self.sequence, penalty_max, sustain, decay)


    def batched_sample(self, logits, temperature, top_k, top_p, min_p, typical, num = 1):

        if logits.shape[0] == 1: return self.sample(logits, temperature, top_k, top_p, min_p, typical, num)

        samples = []
        scores = []
        for i in range(logits.shape[0]):
            t, s = self.sample(logits[i, :, :], temperature, top_k, top_p, min_p, typical)
            samples.append(t)
            scores.append(s)

        return torch.cat(samples, dim = 0), torch.cat(scores, dim = 0)


    # Sample one token from logits with current settings

    def sample_current(self, logits, num = 1):

        return self.sample(logits,
                           self.settings.temperature,
                           self.settings.top_k,
                           self.settings.top_p,
                           self.settings.min_p,
                           self.settings.typical)


    # Sample one token from logits

    def sample(self, logits, temperature, top_k, top_p, min_p, typical, num = 1):

        # torch.manual_seed(42)

        if logits.dim() == 3: logits = logits[0, -1, :]
        elif logits.dim() == 2: logits = logits[-1, :]
        else: raise ValueError("Bad logits dimension")

        # Disallow tokens

        if self.disallowed_tokens is not None:
            logits[self.disallowed_tokens] = float("-inf")

        # Base probabilities

        logits /= temperature
        logits += 1e-8
        probs = torch.softmax(logits, dim = -1)

        # Top K

        if top_k == 0:
            top_probs, top_indices = torch.sort(probs, descending = True)
        else:
            top_probs, top_indices = torch.topk(probs, top_k)
            top_probs = F.normalize(top_probs, p = 1, dim = -1)

        # Top P

        if top_p > 0.0:

            num_top_p_probs = 0
            cum_prob = top_probs[0].item()
            while True:
                num_top_p_probs += 1
                if num_top_p_probs == top_probs.shape[-1]: break
                if top_probs[num_top_p_probs].item() < min_p: break
                cum_prob += top_probs[num_top_p_probs].item()
                if cum_prob > top_p: break

            top_probs = top_probs[:num_top_p_probs]
            top_probs = F.normalize(top_probs, p = 1, dim = -1)
            top_indices = top_indices[:num_top_p_probs]

        # Locally typical sampling

        if typical > 0.0:

            epsilon = 1e-10
            log_probs = (top_probs + epsilon).log()
            neg_entropy = (top_probs * log_probs).sum()
            entropy_dev = (neg_entropy - log_probs).abs()
            _, entropy_dev_order = torch.sort(entropy_dev)

            top_probs = top_probs.gather(-1, entropy_dev_order)
            top_indices = top_indices.gather(-1, entropy_dev_order)

            num_typical_probs = 0
            cum_prob = top_probs[0].item()
            while True:
                num_typical_probs += 1
                if num_typical_probs == top_probs.shape[-1]: break
                cum_prob += top_probs[num_typical_probs].item()
                if cum_prob > typical: break

            top_probs = top_probs[:num_typical_probs]
            top_probs = F.normalize(top_probs, p = 1, dim = -1)
            top_indices = top_indices[:num_typical_probs]

        # Multinomial sampling from top_probs, kept in same order as top_indices

        sampled_ind = torch.multinomial(top_probs, top_probs.shape[-1] if num == -1 else min(num, top_probs.shape[-1]))
        sampled_tokens = top_indices[sampled_ind]
        sampled_probs = top_probs[sampled_ind]  # Return probs before second norm

        if sampled_tokens.shape[0] > 1:
            sampled_tokens, ind = sampled_tokens.sort()
            sampled_probs = sampled_probs[ind]

        return sampled_tokens.unsqueeze(0), sampled_probs.unsqueeze(0)


    def disallow_tokens(self, tokens):

        self.disallowed_tokens = tokens


    def gen_begin(self, in_tokens, mask = None):

        self.end_beam_search()

        self.sequence = in_tokens.clone()
        self.sequence_actual = in_tokens.clone()
        self.cache.current_seq_len = 0

        self.model.forward(self.sequence[:, :-1], self.cache, preprocess_only = True, lora = self.lora, input_mask = mask)


    def gen_begin_empty(self):

        self.end_beam_search()
        self.sequence = None
        self.sequence_actual = None
        self.cache.current_seq_len = 0


    def gen_begin_reuse(self, in_tokens, mask = None):

        self.end_beam_search()
        if self.sequence is None or self.cache.current_seq_len == 0:
            self.gen_begin(in_tokens, mask = mask)
            return 0

        # if in_tokens.shape[-1] < self.sequence.shape[-1]:
        #     self.sequence = self.sequence[:, :in_tokens.shape[-1]]

        reuse = 0
        while reuse < self.sequence.shape[-1] and reuse < in_tokens.shape[-1] and self.sequence[0, reuse] == in_tokens[0, reuse]:
            reuse += 1

        if reuse < 2:
            self.gen_begin(in_tokens, mask = mask)
            return 0

        # print (f"Reusing cache: {reuse} tokens")

        self.cache.current_seq_len = reuse - 1
        self.sequence = self.sequence[:, :reuse]
        self.sequence_actual = self.sequence.clone()

        if reuse < in_tokens.shape[-1]: self.gen_feed_tokens(in_tokens[:, reuse:], mask = mask)
        return reuse


    def gen_feed_tokens(self, in_tokens, mask = None):

        if self.sequence is None:
            self.gen_begin(in_tokens, mask = mask)
            return

        self.end_beam_search()

        start = self.sequence.shape[-1] - 1
        if start < 0:
            start = 0
            self.sequence = in_tokens.clone()
        else:
            self.sequence = torch.cat((self.sequence, in_tokens), dim = 1)

        if start < self.sequence.shape[-1] - 1:
            self.model.forward(self.sequence[:, start : -1], self.cache, preprocess_only = True, lora = self.lora, input_mask = mask)

        self.sequence_actual = self.sequence


    def gen_accept_token(self, token):

        self.end_beam_search()
        if self.sequence is None: self.sequence = token
        else: self.sequence = torch.cat((self.sequence, token), dim = 1)
        self.sequence_actual = self.sequence


    def gen_rewind(self, num_tokens):

        if num_tokens == 0: return
        self.end_beam_search()
        self.sequence = self.sequence[:, :-num_tokens]
        self.cache.current_seq_len -= num_tokens
        self.sequence_actual = self.sequence


    def gen_prune_right(self, tokens, mask = None):

        self.end_beam_search()
        if tokens > self.sequence.shape[-1] - 1: return
        self.gen_begin(self.sequence[:, tokens:], mask = mask)
        self.sequence_actual = self.sequence


    def gen_prune_to(self, min_tokens_to_keep, token_id, mask = None):

        self.end_beam_search()

        if self.gen_num_tokens() <= min_tokens_to_keep: return

        while self.gen_num_tokens() > min_tokens_to_keep:

            pruned = False
            for i in range(self.sequence.shape[-1] - 1):
                if self.sequence[0, i] == token_id:
                    self.sequence = self.sequence[:, i + 1:]
                    pruned = True
                    break

            if not pruned: return

        self.gen_begin(self.sequence, mask = mask)


    def gen_prune_left(self, num_tokens, mask = None):

        num_tokens = min(num_tokens, self.sequence_actual.shape[-1] - 1)

        if self.in_beam_search:
            self.end_beam_search()  # TODO: Try to avoid restarting beam search when generating past chunk boundary
            self.sequence = self.sequence[:, num_tokens:]
            self.begin_beam_search()
        else:
            self.sequence = self.sequence[:, num_tokens:]
            self.gen_begin(self.sequence, mask = mask)


    def gen_num_tokens(self):

        return self.sequence_actual.shape[-1]


    # Simple generator function

    def generate_simple(self, prompt, max_new_tokens = 128):

        self.end_beam_search()

        ids, mask = self.tokenizer.encode(prompt, return_mask = True, max_seq_len = self.model.config.max_seq_len)
        self.gen_begin(ids, mask = mask)

        max_new_tokens = min(max_new_tokens, self.model.config.max_seq_len - ids.shape[1])

        eos = torch.zeros((ids.shape[0],), dtype = torch.bool)
        for i in range(max_new_tokens):
            token = self.gen_single_token(mask = mask)
            for j in range(token.shape[0]):
                if token[j, 0].item() == self.tokenizer.eos_token_id: eos[j] = True
            if eos.all(): break

        text = self.tokenizer.decode(self.sequence[0] if self.sequence.shape[0] == 1 else self.sequence)
        return text


    # Apply repetition penalty with current  settings

    def apply_rep_penalty(self, logits):

        cuda_ext.ext_apply_rep_penalty_mask_cpu(self.sequence,
                                                self.settings.token_repetition_penalty_max,
                                                self.settings.token_repetition_penalty_sustain,
                                                self.settings.token_repetition_penalty_decay,
                                                logits)


    # Generate a single token with the current settings, append to sequence

    def gen_single_token(self, constraints = None, mask = None):

        self.end_beam_search()

        # Simple sampling case:

        if self.sequence is not None:

            logits = self.model.forward(self.sequence[:, -1:], self.cache, lora = self.lora, input_mask = mask)
            self.apply_rep_penalty(logits)

            logits[:, :, self.tokenizer.bos_token_id] = -10000.0

            if constraints is not None:

                for c in constraints: logits[:, :, c] += 10000.0
                logits[:, :, :] -= 10000.0

            token, _ = self.batched_sample(logits,
                                           self.settings.temperature,
                                           self.settings.top_k,
                                           self.settings.top_p,
                                           self.settings.min_p + 0.01 if constraints is not None else 0.0,
                                           self.settings.typical)

        else:

            # bos = torch.Tensor([[self.tokenizer.bos_token_id]]).long()
            # logits = self.model.forward(bos, self.cache)
            # self.cache.current_seq_len = 0

            if constraints is not None:
                token = constraints[0]
            else:
                token = torch.Tensor([[self.tokenizer.bos_token_id]]).long()

        self.gen_accept_token(token)
        return token


    # Beam search

    class Beam:

        sequence: torch.Tensor                  # tokens generated in beam
        probs: torch.Tensor                     # probability score per token
        cache: ExLlamaCache                     # cached keys/values for this beam
        current_seq_pos: int                    # position of beam in current sequence
        settings = None
        generator = None

        sampled_tokens: torch.Tensor
        sampled_probs: torch.Tensor
        moved: bool = False

        def __init__(self, settings, generator, first_token = None, first_prob = None, seq_pos = None):

            self.settings = settings
            self.generator = generator

            self.sequence = first_token.unsqueeze(0).unsqueeze(0) if first_token is not None else None
            self.probs = first_prob.unsqueeze(0).unsqueeze(0) if first_prob is not None else None
            self.cache = ExLlamaCache(self.generator.model, max_seq_len = self.settings.beam_length)

            self.current_seq_pos = seq_pos


        def __len__(self):

            return self.sequence.shape[-1]


        def clone(self):

            new = ExLlamaGenerator.Beam(self.settings, self.generator)
            new.sequence = self.sequence.clone()
            new.probs = self.probs.clone()
            new.cache = self.cache.clone()
            new.current_seq_pos = self.current_seq_pos
            new.sampled_tokens = self.sampled_tokens.clone()
            new.sampled_probs = self.sampled_probs.clone()
            new.moved = self.moved
            return new


        # List of references to this instance

        def advance(self):

            self.cache.roll_left()
            self.sequence = self.sequence[:, 1:]
            self.probs = self.probs[:, 1:]
            self.current_seq_pos += 1


        # Cumulative probabilities

        def cum_log_probs(self):

            cum_log_prob = torch.sum(torch.log(self.probs))
            return cum_log_prob

        def sampled_cum_log_probs(self):

            cum_log_prob = torch.sum(torch.log(self.probs))
            return torch.log(self.sampled_probs) + cum_log_prob


        # Insert current beam in sequence

        def to_sequence(self):

            # Extend generator sequence and cache if needed

            new_tokens = 0
            added_tokens = 0
            slen = self.generator.sequence.shape[-1]
            tlen = self.current_seq_pos + len(self)
            if tlen > slen:
                new_tokens = tlen - slen
                added_tokens = new_tokens
                self.generator.sequence = torch.cat((self.generator.sequence, self.sequence[:, -new_tokens:]), dim = 1)
                self.generator.cache.current_seq_len = tlen - 1

            # Determine how much of generator sequence needs to be updated

            new_tokens_ = new_tokens
            for i in range(new_tokens_, len(self)):
                if self.generator.sequence[0, -i - 1] != self.sequence[0, -i - 1]: new_tokens = i + 1

            # Update sequence and cache

            if new_tokens > added_tokens:
                self.generator.sequence[0, -new_tokens:] = self.sequence[0, -new_tokens:]

            if new_tokens > len(self) - 1: new_tokens = len(self) - 1
            if new_tokens > 0:
                self.cache.copy_states(self.generator.cache,
                                       len(self) - 1 - new_tokens, new_tokens,
                                       self.generator.cache.current_seq_len - new_tokens, new_tokens,
                                       0, 1, 0, 1)


        # Copy last column of cache to this beam (after generation)

        def record_last_cache_column(self):

            self.generator.cache.copy_states(self.cache,
                                             self.generator.cache.current_seq_len - 1, 1,
                                             len(self) - 1, 1,
                                             0, 1, 0, 1)


    def begin_beam_search(self):

        self.beams = None
        if self.settings.beams == 1 and self.settings.beam_length == 1: return

        self.in_beam_search = True
        # self.testl = []


    def beam_search(self):

        if self.settings.beams == 1 and self.settings.beam_length == 1: return self.gen_single_token()
        assert self.in_beam_search

        # Kludge: The first token returned with an empty context is generated without beam search
        if self.sequence is None: return self.gen_single_token()

        c_cache_len = self.cache.current_seq_len
        c_seq_len = self.sequence_actual.shape[-1]

        # Begin here

        max_beam_length = min(self.model.config.max_seq_len - self.settings.beam_length, self.settings.beam_length)
        while self.beams is None or len(self.beams[0]) < max_beam_length:

            if self.beams is None:

                # Initial tokens for initial beams

                # self.cache.debug()
                logits = self.model.forward(self.sequence[:, -1:], self.cache, lora = self.lora)

                cuda_ext.ext_apply_rep_penalty_mask_cpu(self.sequence,
                                                        self.settings.token_repetition_penalty_max,
                                                        self.settings.token_repetition_penalty_sustain,
                                                        self.settings.token_repetition_penalty_decay,
                                                        logits)

                tokens, probs = self.sample(logits,
                                            self.settings.temperature,
                                            self.settings.top_k,
                                            self.settings.top_p,
                                            self.settings.min_p,
                                            self.settings.typical,
                                            num = self.settings.beams)

                # self.cache is updated with k/v for last token
                # Setup initial beams

                self.beams = []
                while len(self.beams) < min(self.settings.beams, tokens.shape[-1]):

                    beam = ExLlamaGenerator.Beam(self.settings, self, tokens[0, len(self.beams)], probs[0, len(self.beams)], c_seq_len)
                    self.beams.append(beam)

            else:

                # Sample from each beam

                # print(len(self.beams), end = "")
                for beam in self.beams:

                    beam.to_sequence()

                    # self.cache.debug()
                    logits = self.model.forward(self.sequence[:, -1:], self.cache, lora = self.lora)

                    cuda_ext.ext_apply_rep_penalty_mask_cpu(self.sequence,
                                                            self.settings.token_repetition_penalty_max,
                                                            self.settings.token_repetition_penalty_sustain,
                                                            self.settings.token_repetition_penalty_decay,
                                                            logits)

                    tokens, probs = self.sample(logits,
                                                self.settings.temperature,
                                                self.settings.top_k,
                                                self.settings.top_p,
                                                self.settings.min_p,
                                                self.settings.typical,
                                                num = -1)

                    beam.sampled_tokens = tokens
                    beam.sampled_probs = probs

                    beam.record_last_cache_column()
                    self.cache.current_seq_len -= 1

                # Collect options for all beams

                tokens_ = []
                probs_ = []
                cum_log_probs_ = []
                beams_ = []
                for i, beam in enumerate(self.beams):
                    tokens_.append(beam.sampled_tokens.squeeze(0))
                    probs_.append(beam.sampled_probs.squeeze(0))
                    cum_log_probs_.append(beam.sampled_cum_log_probs().squeeze(0))
                    beams_.append(torch.Tensor([i] * beam.sampled_tokens.shape[-1]).to(torch.int))

                tokens_all = torch.cat(tokens_, dim = 0)
                probs_all = torch.cat(probs_, dim = 0)
                cum_log_probs_all = torch.cat(cum_log_probs_, dim = 0)
                beams_all = torch.cat(beams_, dim = 0)

                # Sort by cumulative probability

                cum_log_probs_all, ind = cum_log_probs_all.sort(descending = True)
                probs_all = probs_all[ind]
                tokens_all = tokens_all[ind]
                beams_all = beams_all[ind]

                # Reduce to beam limit

                cum_log_probs_all = cum_log_probs_all[:self.settings.beams]
                probs_all = probs_all[:self.settings.beams]
                tokens_all = tokens_all[:self.settings.beams]
                beams_all = beams_all[:self.settings.beams]

                # Re-sort by beam index

                beams_all, ind = beams_all.sort()
                cum_log_probs_all = cum_log_probs_all[ind]
                tokens_all = tokens_all[ind]
                probs_all = probs_all[ind]

                # test = [self.tokenizer.decode(beam.sequence) for beam in self.beams]

                # Rebuild beams/caches

                for beam in self.beams: beam.moved = False
                beams_new = []

                for i in range(len(beams_all)):

                    new_token = tokens_all[i]
                    new_prob = probs_all[i]
                    beam_idx = beams_all[i].item()

                    if not self.beams[beam_idx].moved:

                        self.beams[beam_idx].sequence = torch.cat((self.beams[beam_idx].sequence, new_token.unsqueeze(0).unsqueeze(0)), dim = 1)
                        self.beams[beam_idx].probs = torch.cat((self.beams[beam_idx].probs, new_prob.unsqueeze(0).unsqueeze(0)), dim = 1)
                        self.beams[beam_idx].moved = True
                        beams_new.append(self.beams[beam_idx])

                    else:

                        nbeam = self.beams[beam_idx].clone()
                        nbeam.sequence[:, -1] = new_token
                        nbeam.probs[:, -1] = new_prob
                        beams_new.append(nbeam)

                self.beams = beams_new


        # Beam length is filled up, select winning beam

        max_log_probs = float("-inf")
        best_beam = None
        best_beam_idx = -1
        for beam_idx, beam in enumerate(self.beams):
            beam_log_probs = beam.cum_log_probs()
            if beam_log_probs > max_log_probs:
                max_log_probs = beam_log_probs
                best_beam = beam
                best_beam_idx = beam_idx

        best_token = best_beam.sequence[:, 0]

        # Insert in sequence

        self.sequence[0, c_seq_len] = best_token
        self.sequence_actual = torch.cat((self.sequence_actual, best_token.unsqueeze(0)), dim = 1)

        # Copy cache state for winning beam

        best_beam.to_sequence()

        # Prune other beams that don't begin with the winning token

        beams_new = [best_beam]

        for idx, beam in enumerate(self.beams):
            if idx != best_beam_idx and beam.sequence[:, 0] == best_token:
                beams_new.append(beam)

        self.beams = beams_new

        # Advance all remaining beams and caches

        for beam in self.beams: beam.advance()

        # Done

        return best_token


    def end_beam_search(self):

        if not self.in_beam_search: return

        self.sequence = self.sequence_actual.clone()
        self.cache.current_seq_len = self.sequence.shape[-1] - 1
        self.in_beam_search = False


    def replace_last_token(self, token, seq = False):

        self.sequence_actual[:, -1] = token
        if seq: self.sequence[:, -1] = token


    def sequence_ends_with(self, tokens):

        if self.sequence_actual.shape[-1] < tokens.shape[-1] + 1: return False
        for i in range(tokens.shape[-1]):
            if self.sequence_actual[0, -i - 1] != tokens[0, -i - 1]: return False
        return True