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
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
|
#!/usr/bin/env nix-shell
#!nix-shell -i python3 -p "python3.withPackages (pkgs: with pkgs; [ keras nltk pytorch transformers ])"
"""AI Dungeon is a text-adventure style game powered by AI."""
import argparse
import re
import shutil
import sys
import nltk
from nltk import tokenize
import transformers
def load_nltk():
try:
tokenize.sent_tokenize("")
except LookupError:
nltk.download("punkt")
def complete_sentence(snippet: str) -> bool:
sentences = tokenize.sent_tokenize(snippet)
return len(tokenize.sent_tokenize(sentences[-1] + " extra")) == 2
def trim_sentence(message: str) -> str:
"""Remove extra output after the last period."""
sentences = tokenize.sent_tokenize(message)
if complete_sentence(sentences[-1]) or len(sentences) == 1:
return message
return " ".join(sentences[:-1])
def generate(model: transformers.TextGenerationPipeline, prompt: str, **kwargs) -> str:
"""Generate text from a text generator."""
out = model(prompt, do_sample=True, temperature=0.9, top_k=60, top_p=0.9, **kwargs)
return trim_sentence(out[0]["generated_text"])
def load_model(model: str) -> transformers.TextGenerationPipeline:
"""Load a model by name."""
tokenizer = transformers.AutoTokenizer.from_pretrained(model)
return transformers.pipeline(
"text-generation",
tokenizer=tokenizer,
model=transformers.AutoModelForCausalLM.from_pretrained(model),
pad_token_id=tokenizer.eos_token_id,
)
def pick_flavor() -> str:
"""Query the user for what scenario they want to play."""
flavors = {
"fantasy": (
"You are a wizard named Megumin from the kingdom\n"
"of Larion. You have in your inventory a wizard's staff\n"
"and a spellbook. You are arriving after a day's travel\n"
"at an enchanted tower where there's rumors of gold.\n"
"You walk up to the entrance of the tower."
),
"post apocalyptic": (
"You are a machinist named Azariel, living\n"
"in the city of New New York. It's been almost 10 years\n"
"since the bombs fell, but you still remember it as if\n"
"it were yesterday. You push these thoughts out of your\n"
"mind and focus on the task at hand: finding a water\n"
"purifier for your settlement. You arrive at an\n"
"abandoned settlement to the east of your home."
),
"doctor": (
"Your name is Amelia Plenn. Today is your doctors appointment.\n"
"You're kind of not looking forward to it because your doctor\n"
"is kind of weird, but today when you get to the doctor's office\n"
"it's a different person than usual.\n"
'The new doctor says to you, "Hello, my name is Doctor Cockman.\n'
"It's my pleasure to see you today.\"\n"
"Doctor Cockman is very handsome, and you feel yourself blush as\n"
"he leads you to the doctor's seat. You sit in the seat and he\n"
"runs his warm hands up and down your arms.\n"
'"Are you feeling alright?" Doctor Cockman asks.'
),
"political": (
"You are a political dissident and revolutionary named\n"
"Vladimir Lenin, hiding in exile in Finland while your homeland \n"
"of Russia is controlled by the Tsar. It is January 1917, and\n"
"The Great War has overtaken Europe. Millions of Russian peasants are\n"
"being sacrificed at the altar of international capital. You\n"
"are the leader of the political party Bolsheviks. You've been\n"
"writing letter after letter but the Bolsheviks in Russia don't\n"
"seem to want to publish your letters in Iskra. Maybe it's\n"
"time to return home."
),
}
choices = []
for i, (scenario, script) in enumerate(flavors.items()):
print(f"{i}.\t{scenario}")
choices.append(script)
while True:
choice = input("Choose a scenario: ")
try:
choice_int = int(choice)
except ValueError:
print("Input must be an integer.")
continue
try:
return choices[choice_int]
except IndexError:
print("Index out of bounds.")
continue
def wrap(message: str) -> str:
"""Wrap long lines to terminal width characters."""
width = shutil.get_terminal_size().columns
res = []
for line in message.split("\n"):
if len(line) < width:
res.append(line)
continue
split_line = []
for word in line.split(" "):
if not split_line or len(" ".join(split_line)) + len(word) < width:
split_line.append(word)
continue
res.append(" ".join(split_line))
split_line = [word]
res.append(" ".join(split_line))
return "\n".join(res)
def main() -> None:
"""Run the main game loop."""
load_nltk()
parser = argparse.ArgumentParser("AI dungeon clone")
parser.add_argument("--model", default="gpt2", help="Model to use")
args = parser.parse_args()
model = load_model(args.model)
script = pick_flavor()
prologue = wrap(
generate(
model,
script,
max_new_tokens=20,
forced_eos_token_id=model.tokenizer.eos_token_id,
)
)
print(prologue)
prompt = prologue
prev_response_length = 0
prev_response_lines = 0
while True:
try:
msg = input("> You ")
except EOFError:
break
if msg == "/retry":
prompt = prompt[:-prev_response_length]
print(f"\033[{prev_response_lines}A\033[J", end="")
elif msg == "/edit":
prompt = prompt[:-prev_response_length]
print(f"\033[{prev_response_lines}A\033[J\r", end="")
new_response = sys.stdin.read()
prompt += new_response
prev_response_length = len(new_response)
prev_response_lines = len(new_response.split("\n"))
continue
else:
if not complete_sentence(msg):
msg += "."
prompt += f"You {msg}\n"
prompt = prompt[-10000:]
response = wrap(
generate(model, prompt, return_full_text=False, max_new_tokens=50).strip()
)
print(response)
response += "\n"
prompt += response
prev_response_length = len(response)
prev_response_lines = len(response.split("\n"))
main()
|
