diff --git a/console_play.py b/console_play.py index d1ccb72..96b60fc 100644 --- a/console_play.py +++ b/console_play.py @@ -4,7 +4,7 @@ from google.cloud import storage import json from story.story_manager import * from generator.web.web_generator import * -from generator.ctrl.ctrl_generator import * +#from generator.ctrl.ctrl_generator import * import tensorflow as tf import textwrap @@ -40,12 +40,12 @@ def play_unconstrained(): def play_constrained(): print("\n") - #generator = WebGenerator(CRED_FILE) - generator = CTRLGenerator() + generator = WebGenerator(CRED_FILE) + #generator = CTRLGenerator() story_start = "haunted" verbs_key = "anything" prompt = get_story_start(story_start) - story_manager = ConstrainedStoryManager(generator, prompt, action_verbs_key=verbs_key) + story_manager = CTRLStoryManager(generator, prompt, action_verbs_key=verbs_key) console_print(str(story_manager.story)) possible_actions = story_manager.get_possible_actions() @@ -87,7 +87,7 @@ def play_cached(): if __name__ == '__main__': - play_unconstrained() + play_constrained() diff --git a/generator/ctrl/ctrl_generator.py b/generator/ctrl/ctrl_generator.py index 8aef60d..d262108 100644 --- a/generator/ctrl/ctrl_generator.py +++ b/generator/ctrl/ctrl_generator.py @@ -139,7 +139,23 @@ class CTRLGenerator(): self.topk = 0 - def generate(self, prompt, first_verb_whitelist=True): + def configure_verb_probs(self, probabilities, options): + + # Make sure only a possible verb is chosen. + for word in get_possible_verbs(): + probabilities[self.word2idx[word]] += 100 + + # Disallow used verbs + if "used_verbs" in options: + for verb in options["used_verbs"]: + probabilities[self.word2idx[verb]] = -1e8 + + return probabilities + + def generate(self, prompt, options=None): + + if options is None: + options = {} if prompt[-1] != " ": prompt = prompt + " " @@ -217,10 +233,8 @@ class CTRLGenerator(): for forbidden_token in forbidden_tokens: prompt_logits[_token][self.word2idx[forbidden_token]] = -1e8 - # Make sure only a possible verb is chosen. if first_token: - for word in get_possible_verbs(): - prompt_logits[_token][self.word2idx[word]] += 5 + prompt_logits[_token] = self.configure_verb_probs(prompt_logits[_token], options) # compute probabilities from logits prompt_probs = np.exp(prompt_logits[_token]) diff --git a/generator/tf/tf_generator.py b/generator/tf/tf_generator.py index 61ab77d..07f840d 100644 --- a/generator/tf/tf_generator.py +++ b/generator/tf/tf_generator.py @@ -41,7 +41,7 @@ class TFGenerator(): ckpt = tf.train.latest_checkpoint(model_path) saver.restore(self.sess, ckpt) - def generate(self, prompt): + def generate(self, prompt, options={}): context_tokens = self.enc.encode(prompt) out = self.sess.run(self.output, feed_dict={ self.context: [context_tokens for _ in range(1)] diff --git a/generator/web/web_generator.py b/generator/web/web_generator.py index 5cd6404..0ffc592 100644 --- a/generator/web/web_generator.py +++ b/generator/web/web_generator.py @@ -31,7 +31,7 @@ class WebGenerator(): return response['predictions'] - def generate(self, prompt): + def generate(self, prompt, options={}): while (True): context_tokens = self.enc.encode(prompt) try: diff --git a/story/story_manager.py b/story/story_manager.py index 3997217..723ff91 100644 --- a/story/story_manager.py +++ b/story/story_manager.py @@ -107,8 +107,11 @@ class ConstrainedStoryManager(StoryManager): def get_action_results(self): return [self.generate_action_result(self.story_context(), phrase) for phrase in self.action_phrases] - def generate_action_result(self, prompt, phrase): - action = phrase + " " + self.generator.generate(prompt + phrase) + def generate_action_result(self, prompt, phrase, options=None): + if options is None: + options = {} + + action = phrase + " " + self.generator.generate(prompt + phrase, options) action_result = cut_trailing_sentence(action) action, result = split_first_sentence(action_result) @@ -118,6 +121,27 @@ class ConstrainedStoryManager(StoryManager): return action, result +class CTRLStoryManager(ConstrainedStoryManager): + def __init__(self, generator, story_prompt, action_verbs_key="classic"): + super().__init__(generator, story_prompt) + self.action_phrases = get_action_verbs("anything") + + def get_action_results(self): + + used_verbs = [] + results = [] + for phrase in self.action_phrases: + options = {} + options["used_verbs"] = used_verbs + result = self.generate_action_result(self.story_context(), phrase, options=options) + + used_verb = result[0].split()[1] + used_verbs.append(used_verb) + + results.append(result) + return results + + class CachedStoryManager(ConstrainedStoryManager): def __init__(self, generator, prompt_num, seed, credentials_file, action_verbs_key="classic", ):