| """Steering a turn while it is sampled: how long the narration runs, and what the choices look like.""" |
| import math |
| import re |
|
|
| import torch |
|
|
| from .bard_rules import is_story_choice |
|
|
| LENGTHS = {"short": 60, "medium": 150, "long": None} |
| SENTENCE_END = re.compile(r"[.!?…][\"”’')\]]*\s*$|\n\s*$") |
| SENTENCE_DONE = re.compile(r"[.!?…][\"”’')\]]*(\s|$)") |
| NEG_INF = float("-inf") |
|
|
|
|
| class LengthControl: |
|
|
| MIN_CHOICES = 2 |
| MAX_CHOICES = 5 |
| MAX_CHOICE_TOKENS = 28 |
| COPY_PENALTY = 1.5 |
| COPY_PENALTY_MAX = 8.0 |
| ECHO_PENALTY = 4.0 |
| MIN_WORDS = 4 |
| GIVE_UP_AFTER = 90 |
|
|
| def __init__(self, tok, mode="medium", vars_weight=1.0): |
| self.tok = tok |
| self.vars_id = tok.specials["<|vars|>"] |
| self.choices_id = tok.specials["<|choices|>"] |
| self.sep_id = tok.specials["<|sep|>"] |
| self.go_on = [self.vars_id, self.choices_id] |
| self.stop = self.go_on + [tok.eos_id] |
| self.rep_free = False |
|
|
| first_special = min(tok.specials.values()) |
| decoded = tok.vocab.id_to_token_decoded |
| self.silent = {i for i in range(first_special) if decoded(i) == ""} |
| self.pipes = torch.tensor([i for i in range(first_special) if "|" in decoded(i)], dtype=torch.long) |
| text_like = {self.sep_id, tok.eos_id, tok.specials["<|num|>"], tok.specials["<|snum|>"]} |
| self.no_specials = torch.tensor( |
| [i for i in range(first_special, 1 << 14) if i not in text_like], dtype=torch.long) |
|
|
| self.vars_bias = math.log(vars_weight) if vars_weight > 1 else 0.0 |
| self.set_length(mode) |
|
|
| def set_length(self, mode): |
| self.mode = mode |
| self.budget = LENGTHS[mode] |
|
|
| def start_part(self, raw): |
| self.written = raw |
| self.done = "<|" in raw |
| self.offset = 0 if self.done else len(self.tok.encode(raw)[0]) |
| self.prior = None |
| if "<|choices|>" in raw: |
| texts = raw.split("<|choices|>", 1)[1].split("<|sep|>") |
| self.prior = [[int(i) for i in self.tok.encode(text)[0]] if text else [] for text in texts] |
| self.rep_free = False |
|
|
| def __call__(self, generated, logits): |
| logits = logits.clone() |
| logits[self.pipes.to(logits.device)] = NEG_INF |
| if self.prior is not None or self.choices_id in generated: |
| self.rep_free = True |
| return self.steer_choices(generated, logits) |
| if self.done or any(token in self.stop for token in generated): |
| self.done = True |
| return logits |
| return self.steer_narration(generated, logits) |
|
|
| def steer_narration(self, generated, logits): |
| if not self.has_own_sentence(generated): |
| logits[self.stop] = NEG_INF |
| if self.vars_bias: |
| logits[self.vars_id] -= self.vars_bias |
| if self.budget is None or not generated: |
| return logits |
|
|
| length = self.offset + len(generated) |
| ends_a_sentence = SENTENCE_END.search(self.tok.decode(generated[-4:], [0] * 4)) |
| if length < self.budget or not ends_a_sentence: |
| return logits |
| if length >= 2 * self.budget: |
| return self.end_narration(logits) |
| logits[self.go_on] += 2.0 + 6.0 * (length - self.budget) / self.budget |
| return logits |
|
|
| def has_own_sentence(self, generated): |
| if self.offset + len(generated) >= self.GIVE_UP_AFTER: |
| return True |
| text = (self.written + self.tok.decode(generated)).strip() |
| return bool(SENTENCE_DONE.search(text)) and len(text.split()) >= self.MIN_WORDS |
|
|
| def end_narration(self, logits): |
| best = max(self.go_on, key=lambda token: float(logits[token])) |
| kept = logits[best].clone() |
| logits.fill_(NEG_INF) |
| logits[best] = kept |
| return logits |
|
|
| def steer_choices(self, generated, logits): |
| finished, current = self.split_choices(generated) |
| real_choices = sum(1 for choice in finished if is_story_choice(self.tok.vocab.decode(choice))) |
|
|
| if not current or real_choices + 1 < self.MIN_CHOICES: |
| logits[self.tok.eos_id] = NEG_INF |
| if not current: |
| logits[self.sep_id] = NEG_INF |
| stray = self.no_specials.to(logits.device) |
| logits[stray[stray < logits.numel()]] = NEG_INF |
| if len(finished) + 1 >= self.MAX_CHOICES: |
| logits[self.sep_id] = NEG_INF |
|
|
| if len(current) >= self.MAX_CHOICE_TOKENS: |
| return self.close_choice(logits) |
| self.discourage_copies(logits, finished, current) |
| return logits |
|
|
| def split_choices(self, generated): |
| if self.prior is not None: |
| *finished, head = self.prior |
| written = list(generated) |
| else: |
| finished, head = [], [] |
| after_marker = len(generated) - generated[::-1].index(self.choices_id) |
| written = list(generated[after_marker:]) |
|
|
| pieces = [[]] |
| for token in written: |
| if token == self.sep_id: |
| pieces.append([]) |
| else: |
| pieces[-1].append(token) |
| pieces[0] = head + pieces[0] |
| return [choice for choice in finished + pieces[:-1] if choice], pieces[-1] |
|
|
| def close_choice(self, logits): |
| closers = [self.sep_id, self.tok.eos_id] |
| closing = torch.full_like(logits, NEG_INF) |
| closing[closers] = logits[closers] |
| return logits if torch.isinf(closing).all() else closing |
|
|
| def discourage_copies(self, logits, finished, current): |
| penalties = {} |
| for choice in finished: |
| if choice == current: |
| logits[[self.sep_id, self.tok.eos_id]] = NEG_INF |
| self.penalize_retracing(penalties, choice, current) |
| self.penalize_echo(penalties, choice, current) |
| for token, amount in penalties.items(): |
| logits[token] -= amount |
|
|
| def penalize_retracing(self, penalties, choice, current): |
| done = len(current) |
| if done < len(choice) and choice[:done] == current and choice[done] not in self.silent: |
| amount = min(self.COPY_PENALTY_MAX, self.COPY_PENALTY * (done + 1)) |
| self.raise_penalty(penalties, choice[done], amount) |
|
|
| def penalize_echo(self, penalties, choice, current): |
| if len(current) < 2: |
| return |
| pair = current[-2:] |
| for i in range(len(choice) - 2): |
| if choice[i:i + 2] == pair and choice[i + 2] not in self.silent: |
| self.raise_penalty(penalties, choice[i + 2], self.ECHO_PENALTY) |
|
|
| @staticmethod |
| def raise_penalty(penalties, token, amount): |
| penalties[token] = max(penalties.get(token, 0.0), amount) |
|
|
