diff --git a/console_play.py b/console_play.py index de2aca9..f11a3c1 100644 --- a/console_play.py +++ b/console_play.py @@ -35,6 +35,14 @@ def play_unconstrained(): if action != "": action = action[0].lower() + action[1:] + + if action[-1] == "." or action[-1] == "?" or action[-1] == "!": + action = action[:-1] + if "you " == action.lower()[0:4]: + action = action[4:] + if "i " == action.lower()[0:2]: + action = action[2:] + action = " You " + action + ". " action = first_to_second_person(action) @@ -92,35 +100,6 @@ def play_cached(): console_print(result) -def play_cached_hospital(): - print("\n") - generator = CTRLGenerator() - story_start = "haunted" - prompt = get_story_start(story_start) - story_manager = CTRLStoryManager(generator) - story_manager.enable_caching(CRED_FILE, bucket_name="haunted-hospital") - - story_manager.start_new_story(prompt) - - console_print("\n") - console_print(str(story_manager.story)) - possible_actions = story_manager.get_possible_actions() - while (True): - console_print("\nOptions:") - for i, action in enumerate(possible_actions): - console_print(str(i) + ") " + action) - - result = None - while (result == None): - action_choice = input("Which action do you choose? ") - if action_choice is "print story": - print(story_manager.story) - continue - print("\n") - result, possible_actions = story_manager.act(action_choice) - - console_print(result) - if __name__ == '__main__': play_unconstrained() diff --git a/generator/ctrl/ctrl_generator.py b/generator/ctrl/ctrl_generator.py index 2abff43..15a8fa0 100644 --- a/generator/ctrl/ctrl_generator.py +++ b/generator/ctrl/ctrl_generator.py @@ -135,7 +135,7 @@ class CTRLGenerator(): self.temperature=temperature self.nucleusprob = nucleus_prob - self.penalty = 1.1 + self.penalty = 1.2 self.topk=topk def configure_verb_probs(self, probabilities, options): @@ -180,9 +180,7 @@ class CTRLGenerator(): if not first_letter_capitalized: result = result[0].lower() + result[1:] - while("\n \n \n " in result): - result = result.replace("\n \n \n ", "\n \n ") - + # # print("\n\nAFTER RESULT_REPLACE:") # print(repr(result)) @@ -212,18 +210,22 @@ class CTRLGenerator(): penalized_so_far = set() for _ in range(token + 1): generated_token = tokens_generated[0][_] - penalized_so_far.add(generated_token) - prompt_logits[_token][generated_token] /= self.penalty + if generated_token not in penalized_so_far: + penalized_so_far.add(generated_token) + prompt_logits[_token][generated_token] /= self.penalty # disallow some tokens 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@@", "*@@", "Writing", "Text@@", "\\@@", "
@@", "@@", " """ - - -def set_seed(seed, n_gpu=1): - np.random.seed(seed) - torch.manual_seed(seed) - if n_gpu > 0: - torch.cuda.manual_seed_all(seed) - - -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 (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 - """ - assert logits.dim() == 1 # batch size 1 for now - could be updated for more but the code would be less clear - 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 - - indices_to_remove = sorted_indices[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.0, is_xlnet=False, 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 trange(length): - - inputs = {'input_ids': generated} - if is_xlnet: - # XLNet is a direct (predict same token, not next token) and bi-directional model by default - # => need one additional dummy token in the input (will be masked), attention mask and target mapping (see model docstring) - input_ids = torch.cat((generated, torch.zeros((1, 1), dtype=torch.long, device=device)), dim=1) - perm_mask = torch.zeros((1, input_ids.shape[1], input_ids.shape[1]), dtype=torch.float, device=device) - perm_mask[:, :, -1] = 1.0 # Previous tokens don't see last token - target_mapping = torch.zeros((1, 1, input_ids.shape[1]), dtype=torch.float, device=device) - target_mapping[0, 0, -1] = 1.0 # predict last token - inputs = {'input_ids': input_ids, 'perm_mask': perm_mask, 'target_mapping': target_mapping} - - outputs = model(**inputs) # Note: we could also use 'past' with GPT-2/Transfo-XL/XLNet (cached hidden-states) - next_token_logits = outputs[0][0, -1, :] / temperature - filtered_logits = top_k_top_p_filtering(next_token_logits, top_k=top_k, top_p=top_p) - next_token = torch.multinomial(F.softmax(filtered_logits, dim=-1), num_samples=1) - generated = torch.cat((generated, next_token.unsqueeze(0)), dim=1) - return generated - - -def main(): - - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - seed = 150 - set_seed(seed) - - model_type = "gpt2" - model_name = "gpt2-medium" - model_class, tokenizer_class = MODEL_CLASSES[model_type] - tokenizer = tokenizer_class.from_pretrained(model_name) - model = model_class.from_pretrained(model_name) - model.to(device) - model.eval() - - temperature = 0.9 - top_k = 40 - top_p = 1.0 - - length = 100 - while True: - raw_text = input("Model prompt >>> ") - if model_type in ["transfo-xl", "xlnet"]: - # Models with memory likes to have a long prompt for short inputs. - raw_text = (PADDING_TEXT) + raw_text - context_tokens = tokenizer.encode(raw_text) - out = sample_sequence( - model=model, - context=context_tokens, - length=length, - temperature=temperature, - top_k=top_k, - top_p=top_p, - device=device, - is_xlnet=bool(model_type == "xlnet") - ) - out = out[0, len(context_tokens):].tolist() - text = tokenizer.decode(out, clean_up_tokenization_spaces=True) - print(text) - - -if __name__ == '__main__': - main() \ No newline at end of file diff --git a/data/storytree.py b/other/storytree.py similarity index 100% rename from data/storytree.py rename to other/storytree.py diff --git a/story/story_data.yaml b/story/story_data.yaml index 90f350d..71bce64 100644 --- a/story/story_data.yaml +++ b/story/story_data.yaml @@ -7,7 +7,9 @@ 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 " + 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. You look around you and see " + + zombies: "You're on top of a building. You look over the city and see roaming undead everywhere. " action_verbs: classic: ["You tell", "You use", "You go", "You"] @@ -23,4 +25,6 @@ rooms: haunted_hospital: ["lobby", "hallway", "parking", "roof", "pharmacy"] contexts: + zombies: "A few months ago a zombie outbreak broke out. You now are trying to survive on the ruins of what's left in the midst of zombie hordes. " + 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/utils.py b/story/utils.py index 036fd56..a03cadf 100644 --- a/story/utils.py +++ b/story/utils.py @@ -75,18 +75,17 @@ def split_first_sentence(text): def cut_trailing_sentence(text): - last_period = text.rfind('.') - last_exclamation = text.rfind('!') - - if last_exclamation > last_period: - text = text[0:last_exclamation+1] - elif last_period > 0: - text = text[0:last_period+1] + text = standardize_punctuation(text) + last_punc = max(text.rfind('.'), text.rfind("!"), text.rfind("?")) + + if last_punc > 0: + text = text[0:last_punc+1] return cut_trailing_quotes(text) def replace_outside_quotes(text, current_word, repl_word): + text = standardize_punctuation(text) reg_expr = re.compile(current_word + '(?=([^"]*"[^"]*")*[^"]*$)') @@ -107,7 +106,7 @@ def mapping_variation_pairs(mapping): # Change you it's before a punctuation if mapping[0] is "you": mapping = ("you", "me") - mapping_list.append((" " + mapping[0]+"\,", " " + mapping[1]+",")) + mapping_list.append((" " + mapping[0]+",", " " + mapping[1]+",")) mapping_list.append((" " + mapping[0]+"\?", " " + mapping[1]+"\?")) mapping_list.append((" " + mapping[0]+"\!", " " + mapping[1]+"\!")) mapping_list.append((" " + mapping[0] + "\.", " " + mapping[1] + ".")) @@ -121,6 +120,7 @@ first_to_second_mappings = [ ("Ive", "you've"), ("I am", "you are"), ("I", "you"), + ("i", "you"), ("I've", "you've"), ("my", "your"), ("we","you"),