diff --git a/console_play.py b/console_play.py index 420e680..640aacb 100644 --- a/console_play.py +++ b/console_play.py @@ -23,12 +23,10 @@ def console_print(str, pycharm=False): def play_unconstrained(): generator = CTRLGenerator() #generator = WebGenerator(CRED_FILE) - if len(sys.argv) is 1: - prompt = get_story_start("forest") - else: - prompt = sys.argv[1] + prompt = get_story_start("apocalypse") + context = get_context("apocalypse") story_manager = UnconstrainedStoryManager(generator) - story_manager.start_new_story(prompt) + story_manager.start_new_story(prompt, context=context) print("\n") console_print(str(story_manager.story)) diff --git a/generator/ctrl/ctrl_generator.py b/generator/ctrl/ctrl_generator.py index e1e842f..ba41f05 100644 --- a/generator/ctrl/ctrl_generator.py +++ b/generator/ctrl/ctrl_generator.py @@ -20,7 +20,7 @@ def loss(labels, logits): class CTRLGenerator(): - def __init__(self, control_code="Fantasy ", generate_num=100, temperature=0.5, topk=40, nucleus_prob=0): + def __init__(self, control_code="Writing Text:", generate_num=60, temperature=0.5, topk=40, nucleus_prob=0): self.generate_num=generate_num model_dir = "generator/ctrl/training_utils/seqlen256_v1.ckpt/" @@ -158,7 +158,7 @@ class CTRLGenerator(): if prompt[-1] != " ": prompt = prompt + " " - prompt = second_to_first_person(prompt) + #prompt = second_to_first_person(prompt) prompt = self.control_code + prompt # print("\n\nAFTER PROMPT_REPLACE") @@ -172,7 +172,7 @@ class CTRLGenerator(): first_letter_capitalized = result[0].isupper() result = result.replace("#", "") result = result.replace("*", "") - result = first_to_second_person(result) + #result = first_to_second_person(result) result = remove_profanity(result) if not first_letter_capitalized: @@ -219,7 +219,7 @@ class CTRLGenerator(): forbidden_tokens = ['', 'Sco@@', "&@@", "1]@@", "2]@@", "3]@@", "4]@@", "https://www.@@", "[@@", ":@@", "Edit", "&@@", "2:","1:", ":", "Edit@@", "EDI@@", "EDIT@@", "edit", "TL@@", "tl@@", ";@@", '**', "http://@@", "Redd@@", "UP@@", "mom", "Up@@", "Me:", "Update", "mom@@", "Part", - "http://www.@@", "edit@@", "*@@"] + "http://www.@@", "edit@@", "*@@", "\n", "Writing", "Text@@"] if num_new_lines > self.max_new_lines: forbidden_tokens.append("\n") diff --git a/story/story_data.yaml b/story/story_data.yaml index bc2798a..b087704 100644 --- a/story/story_data.yaml +++ b/story/story_data.yaml @@ -7,6 +7,8 @@ prompts: vague_police: "You are a police officer. You get to work and " + apocalypse: "You walk for two hours and take a break. You've left the town you were in and are now in a more rural area. There's a building to your right and you see" + action_verbs: classic: ["You attack", "You tell", "You use", "You go"] @@ -18,4 +20,7 @@ ctrl_verbs: movement: ["walk", "go", "run", "move"] rooms: - haunted_hospital: ["lobby", "hallway", "parking", "roof", "pharmacy"] \ No newline at end of file + haunted_hospital: ["lobby", "hallway", "parking", "roof", "pharmacy"] + +contexts: + apocalypse: "Long ago the bombs fell and the world ended. You are one of the few who is still alive. You are trying to survive by scavenging among the ruins of what is left behind." \ No newline at end of file diff --git a/story/story_manager.py b/story/story_manager.py index 6637235..be958fc 100644 --- a/story/story_manager.py +++ b/story/story_manager.py @@ -5,8 +5,9 @@ import json class Story(): - def __init__(self, story_start, seed=None, game_state=None): + def __init__(self, story_start, context ="", seed=None, game_state=None): self.story_start = story_start + self.context = context # list of actions. First action is the prompt length should always equal that of story blocks self.actions = [] @@ -33,18 +34,19 @@ class Story(): self.choices = story_dict["choices"] self.possible_action_results = story_dict["possible_action_results"] self.game_state = story_dict["game_state"] + self.context = story_dict["context"] def add_to_story(self, action, story_block): self.actions.append(action) self.results.append(story_block) def latest_result(self): - if len(self.results) > 1: - return self.actions[-2] + self.results[-2] + self.actions[-1] + self.results[-1] + if len(self.results) >= 2: + return self.context + self.results[-1] + self.actions[-1] + self.results[-1] elif len(self.results) >= 1: - return self.story_start + self.actions[-1] + self.results[-1] + return self.context + self.actions[-1] + self.results[-1] else: - return self.story_start + return self.context + self.story_start def __str__(self): story_list = [self.story_start] @@ -63,7 +65,7 @@ class Story(): story_dict["choices"] = self.choices story_dict["possible_action_results"] = self.possible_action_results story_dict["game_state"] = self.game_state - + story_dict["context"] = self.context return json.dumps(story_dict) @@ -72,10 +74,10 @@ class StoryManager(): def __init__(self, generator): self.generator = generator - def start_new_story(self, story_prompt, game_state=None): + def start_new_story(self, story_prompt, context="", game_state=None): block = self.generator.generate(story_prompt) block = cut_trailing_sentence(block) - self.story = Story(story_prompt + block, game_state=game_state) + self.story = Story(story_prompt + block, context=context, game_state=game_state) return self.story def load_story(self, story, from_json=False): diff --git a/story/utils.py b/story/utils.py index 3385513..17d224f 100644 --- a/story/utils.py +++ b/story/utils.py @@ -7,6 +7,12 @@ YAML_FILE = "story/story_data.yaml" from profanityfilter import ProfanityFilter pf = ProfanityFilter() +def get_context(key): + with open(YAML_FILE, 'r') as stream: + data_loaded = yaml.safe_load(stream) + + return data_loaded["context"][key] + def get_story_start(key): with open(YAML_FILE, 'r') as stream: data_loaded = yaml.safe_load(stream)