diff --git a/.gitignore b/.gitignore index 1d98d96..d0a7640 100644 --- a/.gitignore +++ b/.gitignore @@ -7,3 +7,5 @@ story*.json /prompts/* !/prompts/cloveranon-prompts/* /generator/gpt2/models/* +/models/* +!/models/.gitkeep diff --git a/config.ini b/config.ini index d0e301f..68f90dd 100644 --- a/config.ini +++ b/config.ini @@ -9,6 +9,8 @@ # Ok coomers. temp = 0.4 +repetition-penalty = 1.2 + #The number of words the AI has to choose from. # It always chooses the "top k" most likely next words before randomly picking one according to temperature. # Low values reduce the randomness of the AI similar to temp. diff --git a/gpt2generator.py b/gpt2generator.py new file mode 100644 index 0000000..8c21c01 --- /dev/null +++ b/gpt2generator.py @@ -0,0 +1,225 @@ +import os +import torch +import torch.nn.functional as F + +from transformers import GPT2LMHeadModel, GPT2Tokenizer + +from story.utils import cut_trailing_sentence, logger + +# warnings.filterwarnings("ignore") +MODEL_CLASSES = { + "gpt2": (GPT2LMHeadModel, GPT2Tokenizer), +} + +def top_k_top_p_filtering(logits, top_k=0, top_p=0.0, filter_value=-float("Inf")): + """ Filter a distribution of logits using top-k and/or nucleus (top-p) filtering + Args: + logits: logits distribution shape (batch size x vocabulary size) + top_k > 0: keep only top k tokens with highest probability (top-k filtering). + top_p > 0.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering). + Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751) + From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317 + """ + top_k = min(top_k, logits.size(-1)) # Safety check + if top_k > 0: + # Remove all tokens with a probability less than the last token of the top-k + indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None] + logits[indices_to_remove] = filter_value + + if top_p > 0.0: + sorted_logits, sorted_indices = torch.sort(logits, descending=True) + cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) + + # Remove tokens with cumulative probability above the threshold + sorted_indices_to_remove = cumulative_probs > top_p + # Shift the indices to the right to keep also the first token above the threshold + sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() + sorted_indices_to_remove[..., 0] = 0 + + # scatter sorted tensors to original indexing + indices_to_remove = sorted_indices_to_remove.scatter( + dim=1, index=sorted_indices, src=sorted_indices_to_remove + ) + logits[indices_to_remove] = filter_value + return logits + + +def sample_sequence( + model, + length, + context, + num_samples=1, + temperature=1, + top_k=0, + top_p=0.9, + repetition_penalty=1.0, + is_xlnet=False, + is_xlm_mlm=False, + xlm_mask_token=None, + xlm_lang=None, + device="cpu", +): + context = torch.tensor(context, dtype=torch.long, device=device) + context = context.unsqueeze(0).repeat(num_samples, 1) + generated = context + with torch.no_grad(): + # for _ in tqdm(range(length), leave=False, desc='generating'): + for _ in range(length): + + inputs = {"input_ids": generated} + + outputs = model( + **inputs + ) # Note: we could also use 'past' with GPT-2/Transfo-XL/XLNet/CTRL (cached hidden-states) + next_token_logits = outputs[0][:, -1, :] / ( + temperature if temperature > 0 else 1.0 + ) + + # repetition penalty from CTRL (https://arxiv.org/abs/1909.05858) + for i in range(num_samples): + for _ in set(generated[i].tolist()): + next_token_logits[i, _] /= repetition_penalty + + filtered_logits = top_k_top_p_filtering( + next_token_logits, top_k=top_k, top_p=top_p + ).float() + if temperature == 0: # greedy sampling: + next_token = torch.argmax(filtered_logits, dim=-1).unsqueeze(-1) + else: + next_token = torch.multinomial( + F.softmax(filtered_logits, dim=-1), num_samples=1 + ) + generated = torch.cat((generated, next_token), dim=1) + return generated + + +class GPT2Generator: + def __init__( + self, generate_num=60, temperature=0.4, top_k=40, top_p=0.9, censor=False, repetition_penalty=1, + ): + self.generate_num = generate_num + self.temp = temperature + self.top_k = top_k + self.top_p = top_p + self.censor = censor + self.samples = 1 + self.dtype = torch.half + self.repetition_penalty = repetition_penalty + self.batch_size = 1 + self.stop_token = None + + self.model_name = "pytorch-gpt2-xl-aid2-v5" + # self.model_name = "model_v5_pytorch_half" + self.model_dir = "models" + self.checkpoint_path = os.path.join(self.model_dir, self.model_name) + # self.checkpoint_path = 'gpt2' # DEBUG quick test of a smaller untrained model + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + logger.info("Using device={}, checkpoint={}".format(self.device, self.checkpoint_path)) + + # Load tokenizer and model + model_class, tokenizer_class = MODEL_CLASSES["gpt2"] + self.tokenizer = tokenizer_class.from_pretrained(self.checkpoint_path) + self.model = model_class.from_pretrained(self.checkpoint_path) + self.model.to(self.device).to(self.dtype) + self.model.eval() + + # context_tokens = self.tokenizer.encode(' ', add_special_tokens=False) + context_tokens = [ + self.tokenizer.pad_token_type_id, + self.tokenizer.pad_token_type_id, + ] + out = self.sample_sequence(context_tokens).tolist() + # out = out[:, len(context_tokens):].tolist() + for o in out: + text = self.tokenizer.decode(o, clean_up_tokenization_spaces=True) + if self.stop_token: + index = text.find(self.stop_token) + if index == -1: + index = None + text = text[:index] + + def sample_sequence(self, context_tokens=None, generate_num=None, temperature=None): + generate_num = generate_num if (generate_num is not None) else self.generate_num + temperature = temperature if (temperature is not None) else self.temp + out = sample_sequence( + model=self.model, + context=context_tokens, + length=generate_num, + # context=self.context, + temperature=temperature, + top_k=self.top_k, + top_p=self.top_p, + repetition_penalty=self.repetition_penalty, + num_samples=self.samples, + device=self.device + # batch_size=self.batch_size, + ) + return out + + def prompt_replace(self, prompt): + logger.debug("BEFORE PROMPT_REPLACE: `%s`", repr(prompt)) + if len(prompt) > 0 and prompt[-1] == " ": + prompt = prompt[:-1] + + # prompt = second_to_first_person(prompt) + + logger.debug("AFTER PROMPT_REPLACE: `%s`", repr(prompt)) + return prompt + + def result_replace(self, result): + logger.debug("BEFORE RESULT_REPLACE: `%s`", repr(result)) + + result = cut_trailing_sentence(result) + if len(result) == 0: + return "" + first_letter_capitalized = result[0].isupper() + result = result.replace('."', '".') + result = result.replace("#", "") + result = result.replace("*", "") + result = result.replace("\n\n", "\n") + # result = first_to_second_person(result) + + if not first_letter_capitalized: + result = result[0].lower() + result[1:] + + logger.debug("nAFTER RESULT_REPLACE: `%s`", repr(result)) + + return result + + def generate_raw(self, prompt, generate_num=None, temperature=None): + context_tokens = self.tokenizer.encode(prompt, add_special_tokens=False) + + generated = 0 + for _ in range(self.samples // self.batch_size): + out = self.sample_sequence( + context_tokens, + generate_num=generate_num, + temperature=temperature + ) + out = out[:, len(context_tokens) :].tolist() + for o in out: + generated += 1 + text = self.tokenizer.decode(o, clean_up_tokenization_spaces=True) + if self.stop_token: + index = text.find(self.stop_token) + if index == -1: + index = None + text = text[:index] + return text + + def generate(self, prompt, options=None, seed=1): + + prompt = self.prompt_replace(prompt) + + logger.debug("Prompt is: `%s`", repr(prompt)) + + text = self.generate_raw(prompt) + + logger.debug("Generated result is: `%s`", repr(text)) + + result = text + result = self.result_replace(result) + if len(result) == 0: + return self.generate(prompt) + + return result diff --git a/models/.gitkeep b/models/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/new-requirements b/new-requirements new file mode 100644 index 0000000..d631dd3 --- /dev/null +++ b/new-requirements @@ -0,0 +1,12 @@ +google-cloud-storage +gsutil +numpy +profanityfilter +pyyaml==5.2.0 +regex +tracery +transformers==2.3.0 +torch==1.2.0 +colorama +textwrap +pyjarowinkler diff --git a/newREADME.md b/newREADME.md new file mode 100644 index 0000000..f21190a --- /dev/null +++ b/newREADME.md @@ -0,0 +1,40 @@ +# Pytorch AI Dungeon2 + +A Fork of Nick Walton's [AI Dungeon2](https://github.com/AIDungeon/AIDungeon) and pytorch as a backend. + +Uses prompts and play.py from the [Clover edition](https://github.com/cloveranon/Clover-Edition) + +I did this because tensorflow is annoying to compile for my xeon processor an it was actually faster to port the generation code. + +No colab yet. + +If you want the converted model, just ask in the issues and I'll make it available. + +Content warning: This model is trained on the internet which means there will be lots of toxic and offensive content along with the funny and wierd. + +## Screenshot + +![](http://i.imgur.com/4Ox8zDX.png) + +## Changes + +- user: + - added suggested actions from the AI player + - roll a d20 for speech or action to make it harder + - d01: You fail to X + - dX: You try to X + - d20: You successfully X + - use Clover edition ui, prompts, config +- technical: + - use half precision for smaller model, (but this might lead to lower quality, I need to test more) + - better logging + - use pytorch + - set top_k to zero and just use top p + - change clover config file to yaml + +# Model + +magnet link to pytorch model torrent + + ```magnet:?xt=urn:btih:17dcfe3d12849db04a3f64070489e6ff5fc6f63f&dn=model_v5_pytorch&tr=udp%3a%2f%2ftracker.opentrackr.org%3a1337%2fannounce&tr=udp%3a%2f%2fopen.stealth.si%3a80%2fannounce&tr=udp%3a%2f%2fp4p.arenabg.com%3a1337%2fannounce&tr=udp%3a%2f%2ftracker.coppersurfer.tk%3a6969%2fannounce&tr=udp%3a%2f%2ftracker.cyberia.is%3a6969%2fannounce&tr=udp%3a%2f%2ftracker.moeking.me%3a6969%2fannounce&tr=udp%3a%2f%2f9.rarbg.me%3a2710%2fannounce&tr=udp%3a%2f%2ftracker3.itzmx.com%3a6961%2fannounce``` + diff --git a/play.py b/play.py index 594d1b1..3c77bf3 100644 --- a/play.py +++ b/play.py @@ -1,16 +1,31 @@ -import os import configparser import gc +import textwrap +import logging from pathlib import Path from random import shuffle from shutil import get_terminal_size -from generator.gpt2.gpt2_generator import * from story.story_manager import * from story.utils import * -import textwrap +from gpt2generator import GPT2Generator + +#silence transformers outputs when loading model +logging.getLogger("transformers.tokenization_utils").setLevel(logging.WARN) +logging.getLogger("transformers.modeling_utils").setLevel(logging.WARN) +logging.getLogger("transformers.configuration_utils").setLevel(logging.WARN) + +logging.basicConfig( + format='%(asctime)s - %(levelnames)s - %(messages)s', + datefmt='%m/%d/%Y %H:%M%S', + level=logging.INFO +) +logger.setLevel(10)#settings['log-level']) + +#TODO: Move all these utilty functions to seperate utily and config file #add color for windows users that install colorama +# It is not necessary to install colorama on most systems, but it does come with pip try: import colorama colorama.init() @@ -20,13 +35,11 @@ except ModuleNotFoundError: with open(Path('interface', 'clover'), 'r', encoding='utf-8') as file: print(file.read()) -#perhaps all the following should be put in a seperate utils file like original config = configparser.ConfigParser() config.read('config.ini') settings=config["Settings"] colors=config["Colors"] -os.environ["TF_CPP_MIN_LOG_LEVEL"] = settings["log-level"] #ECMA-48 set graphics codes for the curious. Check out "man console_codes" @@ -44,7 +57,7 @@ def getNumberInput(n): val=colInput("Enter a number from above (default 0):", colors["selection-prompt"], colors["selection-value"]) if val=='': return 0 - elif 0>int(val) or int(val)>n: + elif not re.match('^\d+$', val) or 0>int(val) or int(val)>n: colPrint("Invalid choice.", colors["error"]) return getNumberInput(n) else: @@ -55,7 +68,9 @@ def selectFile(p=Path('prompts')): files=[x for x in p.iterdir()] shuffle(files) for n in range(len(files)): - colPrint('{}: {}'.format(n, re.sub(r'\.txt$', '', files[n].name)), colors["menu"]) + colPrint( + '{}: {}'.format(n, re.sub(r'\.txt$', '', files[n].name)), + colors["menu"]) return selectFile(files[getNumberInput(len(files)-1)]) else: with p.open('r', encoding='utf-8') as file: @@ -63,38 +78,71 @@ def selectFile(p=Path('prompts')): rest=file.read() return (line1, rest) +#print files done several times and probably deserves own function def instructions(): with open('interface/instructions.txt', 'r', encoding='utf-8') as file: colPrint(file.read(), colors["instructions"], False) def getGenerator(): - colPrint("\nInitializing AI Engine! (This might take a few minutes)\n", colors["loading-message"]) + colPrint( + "\nInitializing AI Engine! (This might take a few minutes)\n", + colors["loading-message"]) return GPT2Generator( generate_num=settings.getint('generate-num'), - temperature=settings.getfloat("temp"), - top_k=settings.getint("top-keks"), - top_p=settings.getfloat("top-p")) + temperature=settings.getfloat('temp'), + top_k=settings.getint('top-keks'), + top_p=settings.getfloat('top-p'), + repetition_penalty=settings.getfloat('repetition-penalty') + ) if not Path('prompts', 'Anime').exists(): try: import pastebin except: + logger.warning('Failed to scrape pastebin: %e', e) colPrint("Failed to scrape pastebin, possible connection issue.\nTry again later. Continuing without downloading prompts...", colors['error']) +class AIPlayer: + def __init__(self, generator): + self.generator = generator + + def get_action(self, prompt): + result_raw = self.generator.generate_raw( + prompt, generate_num=settings.getint('action-generate-num'), temperature=settings.getint('temp')) + return clean_suggested_action(result_raw, min_length=settings.getint('action-min-length')) + + def get_actions(self, prompt): + suggested_actions = [ + self.get_action(prompt) + for _ in range(settings.getint('action-alternatives')) + ] + logger.debug("Suggested actions before filter and dedup %s", suggested_actions) + #remove short ones + suggested_actions = [ + s + for s in suggested_actions + if len(s) > settings.getint('action-min-length') + ] + #remove dups + suggested_actions = list(set(suggested_actions)) + return suggested_actions + def play(): - story_manager = UnconstrainedStoryManager(getGenerator()) + generator = getGenerator() + story_manager = UnconstrainedStoryManager(generator) + ai_player = AIPlayer(generator) print("\n") - with open("interface/mainTitle.txt", "r", encoding="utf-8") as file: - colPrint(file.read(), colors["title"]) + with open(Path('interface', 'mainTitle.txt'), 'r', encoding='utf-8') as file: + colPrint(file.read(), colors['title']) - with open('interface/subTitle.txt', 'r', encoding="utf-8") as file: + with open(Path('interface', 'subTitle.txt'), 'r', encoding='utf-8') as file: cols=get_terminal_size()[0] for line in file: line=re.sub(r'\n', '', line) line=line[:cols] + #fills in the graphic using reverse video mode substituted into the areas between |'s colPrint(re.sub(r'\|[ _]*\|', lambda x: '\x1B[7m'+x.group(0)+'\x1B[27m', line), colors["subtitle"], False) - while True: if story_manager.story != None: @@ -107,32 +155,53 @@ def play(): if getNumberInput(1) == 1: with open(Path('interface', 'prompt-instructions.txt'), 'r', encoding='utf-8') as file: colPrint(file.read(), colors['instructions'], False) - context=colInput('Context>', colors['main-prompt'], colors['user-text']) - prompt=colInput('Prompt>', colors['main-prompt'], colors['user-text']) + context = colInput('Context>', colors['main-prompt'], colors['user-text']) + prompt = colInput('Prompt>', colors['main-prompt'], colors['user-text']) filename=colInput('Name to save prompt as? (Leave blank for no save): ', colors['query'], colors['user-text']) filename=re.sub('-$','',re.sub('^-', '', re.sub('[^a-zA-Z0-9_-]+', '-', filename))) if filename != '': with open(Path('prompts', filename+'.txt'), 'w', encoding='utf-8') as f: - #this saves unix style line endings which might be an issue - #don't know how to do this properly - f.write(context+'\n'+prompt+'\n') + f.write(context+'\n'+prompt) else: context, prompt = selectFile() instructions() - colPrint("\nGenerating story...", colors["loading-message"]) + print() + colPrint("Generating story...", colors['loading-message']) - story_manager.start_new_story( - prompt, context=context - ) + #TODO:seperate out AI generated part of story and print with different color + story_manager.start_new_story(prompt, context=context) print("\n") colPrint(str(story_manager.story), colors["ai-text"]) while True: + #Generate suggested actions + if settings.getint('action-alternatives') > 0: + + #TODO change this to two messages for different colors + action_prompt = ( + story_manager.story.results[-1] + if story_manager.story.results + else "\nWhat do you do now?" + ) + "\n>" + #can this just be a for loop? + suggested_actions = ai_player.get_actions(action_prompt) + if len(suggested_actions): + suggested_actions_enum = [ + "{}> {}\n".format(i, a) for i, a in enumerate(suggested_actions) + ] + suggested_action = "".join(suggested_actions_enum) + #TODO: check color + colPrint('Suggested actions\n' + suggested_action, colors['selection-value']) + print() + if settings.getboolean('console-bell'): print('\x07', end='') action = colInput("> ", colors["main-prompt"], colors["user-text"]) + + #TODO:Clear suggestions and user input + setRegex = re.search('^set ([^ ]+) ([^ ]+)$', action) if setRegex: if setRegex.group(1) in settings: @@ -144,9 +213,6 @@ def play(): if colInput('y/n? >', colors['selection-prompt'], colors['selection-value']) == 'y': with open('config.ini', 'w', encoding='utf-8') as file: config.write(file) - del story_manager.generator - gc.collect() - story_manager.generator = getGenerator() else: colPrint('Invalid Setting', colors['error']) instructions() @@ -158,11 +224,11 @@ def play(): instructions() elif action == "print": print("\nPRINTING\n") - colPrint(str(story_manager.story), colors["print-story"]) + colPrint(str(story_manager.story), colors['print-story']) elif action == "revert": - if len(story_manager.story.actions) is 0: - colPrint("You can't go back any farther. ", colors["error"]) + if len(story_manager.story.actions) == 0: + colPrint("You can't go back any farther. ", colors['error']) continue story_manager.story.actions = story_manager.story.actions[:-1] @@ -225,5 +291,5 @@ def play(): else: colPrint(result, colors["ai-text"]) - +#TODO: there's no reason for this to be enclosed in a function play() diff --git a/pytorch-scripts/aidungeonv2_model_v5_config.json b/pytorch-scripts/aidungeonv2_model_v5_config.json new file mode 100644 index 0000000..7346cfc --- /dev/null +++ b/pytorch-scripts/aidungeonv2_model_v5_config.json @@ -0,0 +1,35 @@ +{ + "attn_pdrop": 0.1, + "embd_pdrop": 0.1, + "finetuning_task": null, + "id2label": { + "0": "LABEL_0", + "1": "LABEL_1" + }, + "initializer_range": 0.02, + "is_decoder": false, + "label2id": { + "LABEL_0": 0, + "LABEL_1": 1 + }, + "layer_norm_epsilon": 1e-05, + "n_ctx": 1024, + "n_embd": 1600, + "n_head": 25, + "n_layer": 48, + "n_positions": 1024, + "num_labels": 2, + "output_attentions": false, + "output_hidden_states": false, + "output_past": true, + "pruned_heads": {}, + "resid_pdrop": 0.1, + "summary_activation": null, + "summary_first_dropout": 0.1, + "summary_proj_to_labels": true, + "summary_type": "cls_index", + "summary_use_proj": true, + "torchscript": false, + "use_bfloat16": false, + "vocab_size": 50257 +} diff --git a/pytorch-scripts/convert_gpt2_model.py b/pytorch-scripts/convert_gpt2_model.py new file mode 100644 index 0000000..519daaa --- /dev/null +++ b/pytorch-scripts/convert_gpt2_model.py @@ -0,0 +1,110 @@ +# coding=utf-8 +# Copyright 2018 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Convert OpenAI GPT checkpoint.""" + +from __future__ import absolute_import, division, print_function + +import argparse +from io import open + +import torch + +from transformers import ( + CONFIG_NAME, + WEIGHTS_NAME, + GPT2Config, + GPT2Model, + load_tf_weights_in_gpt2, +) + +import logging + +logging.basicConfig(level=logging.INFO) + + +def convert_gpt2_checkpoint_to_pytorch( + gpt2_checkpoint_path, gpt2_config_file, pytorch_dump_folder_path +): + # Construct model + if gpt2_config_file == "": + config = GPT2Config() + else: + config = GPT2Config.from_json_file(gpt2_config_file) + model = GPT2Model(config) + + # Load weights from numpy + load_tf_weights_in_gpt2(model, config, gpt2_checkpoint_path) + + # Save pytorch-model + pytorch_weights_dump_path = pytorch_dump_folder_path + "/" + WEIGHTS_NAME + pytorch_config_dump_path = pytorch_dump_folder_path + "/" + CONFIG_NAME + print("Save PyTorch model to {}".format(pytorch_weights_dump_path)) + torch.save(model.state_dict(), pytorch_weights_dump_path) + + print("Save configuration file to {}".format(pytorch_config_dump_path)) + with open(pytorch_config_dump_path, "w", encoding="utf-8") as f: + f.write(config.to_json_string()) + + # Also save as half precision to save transfer, loading, and inference time and memory + pytorch_weights_dump_path += 'half' + model.half() + torch.save(model.state_dict(), pytorch_weights_dump_path) + with open(pytorch_config_dump_path "w", encoding="utf-8") as f: + f.write(config.to_json_string()) + print("Save configuration file to {}".format(pytorch_config_dump_path)) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + ## Required parameters + parser.add_argument( + "--gpt2_checkpoint_path", + default=None, + type=str, + required=True, + help="Path to the TensorFlow checkpoint path.", + ) + parser.add_argument( + "--pytorch_dump_folder_path", + default=None, + type=str, + required=True, + help="Path to the output PyTorch model.", + ) + parser.add_argument( + "--gpt2_config_file", + default="", + type=str, + help="An optional config json file corresponding to the pre-trained OpenAI model. \n" + "This specifies the model architecture.", + ) + args = parser.parse_args() + convert_gpt2_checkpoint_to_pytorch( + args.gpt2_checkpoint_path, args.gpt2_config_file, args.pytorch_dump_folder_path + ) + + +""" +download aidungeon2 v5 model from this torrent or elsewhere +- magnet:?xt=urn:btih:b343b83b35bff774dab13e0281ce13b3daf37d3e&dn=model_v5&tr=udp%3a%2f%2ftracker.coppersurfer.tk%3a6969%2fannounce&tr=udp%3a%2f%2ftracker.leechers-paradise.org%3a6969%2fannounce +export OPENAI_GPT2_CHECKPOINT_PATH=../generator/gpt2/models/model_v5 +export PYTORCH_DUMP_OUTPUT=../generator/gpt2/models/model_v5_pytorch +python convert_gpt2_model.py \ + --gpt2_checkpoint_path $OPENAI_GPT2_CHECKPOINT_PATH \ + --pytorch_dump_folder_path $PYTORCH_DUMP_OUTPUT \ + --gpt2_config_file ./aidungeonv2_model_v5_config.json +wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-xl-merges.txt -o $PYTORCH_DUMP_OUTPUT/merges.txt +wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-xl-vocab.json -o $PYTORCH_DUMP_OUTPUT/vocab.json +""" diff --git a/story/utils.py b/story/utils.py index d160b71..6946702 100644 --- a/story/utils.py +++ b/story/utils.py @@ -1,6 +1,13 @@ # coding: utf-8 import re -from difflib import SequenceMatcher +import logging +from pyjarowinkler import distance + +logger = logging.getLogger(__name__) + + + + def console_print(text, width=75): @@ -18,8 +25,26 @@ def console_print(text, width=75): print(text) +#TODO: get rid if pyjarowinker dependency def get_similarity(a, b): - return SequenceMatcher(None, a, b).ratio() + return distance.get_jaro_distance( + a, b, winkler=True, scaling = 0.1 + ) + + +def get_num_options(num): + + while True: + choice = input("Enter the number of your choice: ") + try: + result = int(choice) + if result >= 0 and result < num: + return result + else: + print("Error invalid choice. ") + except ValueError: + print("Error invalid choice. ") + def player_died(text): """ @@ -33,6 +58,7 @@ def player_died(text): "you (die|pass away|perish|suffocate|drown|bleed out)", "you('ve| have) (died|perished|suffocated|drowned|been (killed|slain))", "you (\w* )?(yourself )?to death", + "you (\w* )*(collapse|bleed out|chok(e|ed|ing)|drown|dissolve) (\w* )*and (die(|d)|pass away|cease to exist|(\w* )+killed)", ] return any(re.search(regexp, lower_text) for regexp in you_dead_regexps) @@ -40,14 +66,20 @@ def player_died(text): def player_won(text): lower_text = text.lower() won_phrases = [ - "you live happily ever after", - "you live (forever|eternally|for eternity)", - "you (are|become|turn into) (a)? (deity|god)", - "you ((go|get) (in)?to|arrive (at|in)) (heaven|paradise)", + "you ((\w* )*and |)live happily ever after", + "you ((\w* )*and |)live (forever|eternally|for eternity)", + "you ((\w* )*and |)(are|become|turn into) ((a|now) )?(deity|god|immortal)", + "you ((\w* )*and |)((go|get) (in)?to|arrive (at|in)) (heaven|paradise)", + "you ((\w* )*and |)celebrate your (victory|triumph)", + "you ((\w* )*and |)retire", ] return any(re.search(regexp, lower_text) for regexp in won_phrases) +def remove_profanity(text): + return pf.censor(text) + + def cut_trailing_quotes(text): num_quotes = text.count('"') if num_quotes % 2 is 0: @@ -79,11 +111,44 @@ def cut_trailing_action(text): or "You ask" in last_line or "you say" in last_line or "You say" in last_line - ): + ) and len(lines) > 1: text = "\n".join(lines[0:-1]) return text +def clean_suggested_action(result_raw, min_length=4): + result_raw = standardize_punctuation(result_raw) + + # The generations actions carry on into the next prompt, so lets remove the prompt + results = result_raw.split("\n") + results = [s.strip() for s in results] + results = [s for s in results if len(s) > min_length] + # Sometimes actions are generated with leading > ! . or ?. Likely the model trying to finish the prompt or start an action. + result = results[0].strip().lstrip(" >!.?") + result = cut_trailing_quotes(result) + logger.debug( + "full suggested action '%s'. Cropped: '%s'. Split '%s'", + result_raw, + result, + results, + ) + + # Often actions are cropped with sentance fragment, lets remove. Or we could just turn up config_act["generate-number"] + last_punc = max(result.rfind("."), result.rfind("!"), result.rfind("?")) + if (last_punc / (len(result) + 1)) > 0.7: + result = result[:last_punc] + elif last_punc == len(result): + pass + else: + result += "..." + + # Remove you from start + result = first_to_second_person(result) + result = re.sub('^ ?[Yy]ou ?', '', result) + logger.debug("suggested action after cleaning %s", result) + return result + + def cut_trailing_sentence(text): text = standardize_punctuation(text) last_punc = max(text.rfind("."), text.rfind("!"), text.rfind("?"))