File size: 4,984 Bytes
02afb58 18fe584 02afb58 76d09fe 02afb58 18fe584 02afb58 76d09fe 02afb58 18fe584 02afb58 18fe584 02afb58 18fe584 02afb58 18fe584 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | import os
import torch
from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutput
from .bard_core import LM, KVCache, softcap
from .bard_game import DEFAULT_OPENING, StoryEngine, Tale
from .bard_steering import LENGTHS
from .bard_tokenizer import FORMAT_FILE, VOCAB_FILE, BardTokenizer
from .configuration_bard import BardConfig
class BardForStory(PreTrainedModel):
"""bard-small: a small language model that tells interactive tales. Use `begin` and `act` to play one:
bard = AutoModelForCausalLM.from_pretrained("<repo>", trust_remote_code=True)
tale = bard.begin("You wake up in a dungeon.")
print(tale.text, tale.choices)
tale = bard.act(tale, tale.choices[0]) # or any words of your own
"""
config_class = BardConfig
main_input_name = "input_ids"
def __init__(self, config: BardConfig):
super().__init__(config)
self.lm = LM(config.vocab_size, config.d_model, config.n_heads, config.n_layers, config.window,
config.global_every, parallel=config.parallel, rope=config.rope)
self._tokenizer = None
self._engine = None
self.post_init()
def _init_weights(self, module):
pass
@property
def tokenizer(self) -> BardTokenizer:
if self._tokenizer is None:
where = getattr(self.config, "_name_or_path", "") or ""
if os.path.isdir(where):
self._tokenizer = BardTokenizer.from_dir(where)
else:
from huggingface_hub import hf_hub_download
self._tokenizer = BardTokenizer(hf_hub_download(where, VOCAB_FILE), hf_hub_download(where, FORMAT_FILE))
return self._tokenizer
def forward(self, input_ids, number_values=None, **kwargs):
if number_values is None:
number_values = torch.zeros_like(input_ids)
cache = KVCache(self.lm.body)
h = self.lm.hidden(input_ids, number_values, cache=cache)
return CausalLMOutput(logits=softcap(torch.nn.functional.linear(h, self.lm.head.weight)))
@property
def engine(self) -> StoryEngine:
if self._engine is None:
self._engine = StoryEngine(self.lm, self.tokenizer, self.config.vars_weight, self.config.context)
return self._engine
def begin(self, opening: str = DEFAULT_OPENING, length: str = "medium", state: dict | None = None,
on_text=None, seed: int | None = None) -> Tale:
"""Begin a tale. `opening` is how it starts (the model goes on from it); `length` is "short", "medium" or
"long"; `state` the variables the player carries (default: 100 money). `on_text(chunk)` is called with the
narration as it is written. Returns a Tale: .text (the passage), .choices, .state, .ended."""
if length not in LENGTHS:
raise ValueError(f"length must be one of {list(LENGTHS)}")
tale = self.engine.new_tale(opening, length, {"money": 100} if state is None else state)
return self.engine.turn(tale, None, on_text=on_text, seed=seed)
def act(self, tale: Tale, move: str, on_text=None, seed: int | None = None) -> Tale:
"""Tell the next turn: `move` is one of tale.choices, or the player's own words. The tale is updated in place
and returned."""
if tale.ended:
raise ValueError("this tale has ended")
return self.engine.turn(tale, move, on_text=on_text, seed=seed)
def play(self, opening: str = DEFAULT_OPENING, length: str = "medium"):
"""Play a tale in the terminal: pick a choice by its number or type your own move. /short, /medium and /long
change how much is told, /quit leaves."""
tale = self.begin(opening, length, on_text=print_chunk)
while True:
show_choices(tale)
if tale.ended:
print("\nThe end.")
return tale
move = self.ask(tale)
if move is None:
return tale
if move:
print()
tale = self.act(tale, move, on_text=print_chunk)
@staticmethod
def ask(tale):
"""The player's next move: a choice's text, their own words, "" to ask again, or None to leave."""
try:
answer = input("\n> ").strip()
except (EOFError, KeyboardInterrupt):
return None
if answer == "/quit":
return None
if answer in ("/short", "/medium", "/long"):
tale.length = answer[1:]
return ""
if answer.isdigit() and 1 <= int(answer) <= len(tale.choices):
return tale.choices[int(answer) - 1]
return answer
def print_chunk(chunk):
print(chunk, end="", flush=True)
def show_choices(tale):
print()
for number, choice in enumerate(tale.choices, 1):
print(f" {number}. {choice}")
for name, old, new in tale.changes:
print(f" [{name}: {old} -> {new}]")
|