From 397ef22d661c1266398050d143bd4d9f25c3ec34 Mon Sep 17 00:00:00 2001 From: nickwalton Date: Thu, 31 Oct 2019 12:41:51 -0600 Subject: [PATCH] fix --- generator/ctrl/ctrl_generator.py | 2 +- generator/pytorch/pytorch_generator.py | 155 ------------------------- {data => other}/storytree.py | 0 3 files changed, 1 insertion(+), 156 deletions(-) delete mode 100644 generator/pytorch/pytorch_generator.py rename {data => other}/storytree.py (100%) diff --git a/generator/ctrl/ctrl_generator.py b/generator/ctrl/ctrl_generator.py index ef50cef..31239c9 100644 --- a/generator/ctrl/ctrl_generator.py +++ b/generator/ctrl/ctrl_generator.py @@ -221,7 +221,7 @@ class CTRLGenerator(): "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