better truncation of context

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))
 

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:]

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