From 367c3c68edef24849aee9b3d9721bdde904e5368 Mon Sep 17 00:00:00 2001 From: Nick Walton Date: Fri, 27 Sep 2019 13:48:07 -0600 Subject: [PATCH] update --- generator/ctrl/ctrl_generator.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/generator/ctrl/ctrl_generator.py b/generator/ctrl/ctrl_generator.py index 6b23251..940be9a 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="Writing Text: ", generate_num=64, temperature=0.5): + def __init__(self, control_code="Writing Text: ", generate_num=64, temperature=0.5, topk=40): self.generate_num=generate_num model_dir = "generator/ctrl/model/seqlen256_v1.ckpt/" @@ -136,7 +136,7 @@ class CTRLGenerator(): self.temperature=temperature self.nucleusprob = 0 self.penalty = 1.2 - self.topk=0 + self.topk=topk def configure_verb_probs(self, probabilities, options): @@ -218,7 +218,7 @@ class CTRLGenerator(): # 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"] + '**', "http://@@", "Redd@@", "UP@@", "mom", "Up@@"] if num_new_lines > self.max_new_lines: forbidden_tokens.append("\n")