From 7b95fa7e7938185fb3a2ff6a230bd1c485e9de1b Mon Sep 17 00:00:00 2001 From: wassname Date: Wed, 25 Dec 2019 07:39:59 +0800 Subject: [PATCH] better truncation of context --- gpt2generator.py | 24 +++++++++++++----------- play.py | 14 ++++++++------ story/story_manager.py | 8 ++++---- 3 files changed, 25 insertions(+), 21 deletions(-) diff --git a/gpt2generator.py b/gpt2generator.py index 52fd2a4..629b841 100644 --- a/gpt2generator.py +++ b/gpt2generator.py @@ -100,6 +100,11 @@ def sample_sequence( generated = torch.cat((generated, next_token), dim=1) return generated +def truncate_multiple_sequences(seqs, max_len=100): + """Truncate multiple sequences, longest first, removing first.""" + while sum(len(s) for s in seqs) > max_len: + longest = sorted(seqs, key=len, reverse=True)[0] + longest.pop(0) class GPT2Generator: def __init__( @@ -114,8 +119,7 @@ class GPT2Generator: self.dtype = torch.float32 if CPU else torch.float16 self.repetition_penalty = repetition_penalty self.batch_size = 1 - self.stop_token = None - self.max_history_tokens = 256 + self.max_history_tokens = 1024 - generate_num self.stop_token = '<|endoftext|>' self.model_name = "pytorch-gpt2-xl-aid2-v5" @@ -184,14 +188,12 @@ class GPT2Generator: return result def generate_raw(self, prompt, generate_num=None, temperature=None): - context_tokens = self.tokenizer.encode(prompt, add_special_tokens=False) - # TODO instead of taking last 1024, take first X and last Y - # crop context to avoid going of the GPT2 max context size of 1024 - if len(context_tokens) > self.max_history_tokens: - # FIXME it would be better to pass in a list of strings so we can cut some out, and a truncation strategy https://github.com/huggingface/transformers/blob/ce50305e5b8c8748b81b0c8f5539a337b6a995b9/src/transformers/tokenization_utils.py#L791 - first = self.max_history_tokens // 4 - last = self.max_history_tokens - first - context_tokens = context_tokens[:first] + context_tokens[-last:] + # the prompt is a list of strings, encode each one tok tokens, then truncate the longest ones + context_tokens = [self.tokenizer.encode(p, add_special_tokens=False, max_length=self.max_history_tokens) for p in prompt] + truncate_multiple_sequences(context_tokens, self.max_history_tokens) + context_tokens = list(itertools.chain(*context_tokens)) + + logger.debug("Text passing into model %s", self.tokenizer.decode(o, clean_up_tokenization_spaces=True, skip_special_tokens=True)) generated = 0 for _ in range(self.samples // self.batch_size): @@ -213,7 +215,7 @@ class GPT2Generator: def generate(self, prompt, options=None, seed=1): - prompt = self.prompt_replace(prompt) + prompt = [self.prompt_replace(p) for p in prompt] logger.debug("Prompt is: `%s`", repr(prompt)) diff --git a/play.py b/play.py index 9962730..15348d7 100644 --- a/play.py +++ b/play.py @@ -167,16 +167,14 @@ def play(generator): 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>" suggested_actions = [] colPrint('Suggested actions:', colors['selection-value']) action_suggestion_lines = 1 for i in range(settings.getint('action-alternatives')): # FIXME it might be better to pass in a longer history + story_manager.story_context() # This should be within the loop as it has a random sampling element + action_prompt[-1] += '> ' + logger.debug("action_prompt %s", action_prompt) suggested_action = ai_player.get_action(action_prompt) suggested_actions.append(suggested_action) suggestion = '{}> {}'.format(i, suggested_action) @@ -238,7 +236,12 @@ def play(generator): # Options to select a suggestion action if action in [str(i) for i in range(len(suggested_actions))]: action = suggested_actions[int(action)] + + action = action.strip() + # Crop actions to a max length + action = action[:4096] + if action != "": # Roll a 20 sided dice to make things interesting @@ -259,7 +262,6 @@ def play(generator): else: action = "You say " + action else: - action = action.strip() action = first_to_second_person(action) if not action.lower().startswith("you ") and not action.lower().startswith("i "): action = action[0].lower() + action[1:] diff --git a/story/story_manager.py b/story/story_manager.py index 4c4322e..cc52a69 100644 --- a/story/story_manager.py +++ b/story/story_manager.py @@ -3,7 +3,7 @@ import os import subprocess import uuid from subprocess import Popen - +import random from story.utils import * @@ -164,7 +164,7 @@ class StoryManager: def start_new_story( self, story_prompt, context="", game_state=None, upload_story=False ): - block = self.generator.generate(context + story_prompt) + block = self.generator.generate([context, story_prompt]) block = cut_trailing_sentence(block) self.story = Story( context + story_prompt + block, @@ -210,7 +210,7 @@ class UnconstrainedStoryManager(StoryManager): return result def generate_result(self, action): - block = self.generator.generate(self.story_context() + action) + block = self.generator.generate(self.story_context()+[action]) return block @@ -321,7 +321,7 @@ class ConstrainedStoryManager(StoryManager): def generate_action_result(self, prompt, phrase, options=None): action_result = ( - phrase + " " + self.generator.generate(prompt + " " + phrase, options) + phrase + " " + self.generator.generate(prompt + [phrase], options) ) action, result = split_first_sentence(action_result) return action, result