From fbf645d318adf6b15c9b1b24b80623c0c59cdb37 Mon Sep 17 00:00:00 2001 From: Nick Date: Wed, 18 Sep 2019 23:04:52 -0600 Subject: [PATCH] update fix --- generator/ctrl/ctrl_generator.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/generator/ctrl/ctrl_generator.py b/generator/ctrl/ctrl_generator.py index b87b9b8..30a0c30 100644 --- a/generator/ctrl/ctrl_generator.py +++ b/generator/ctrl/ctrl_generator.py @@ -19,10 +19,11 @@ def loss(labels, logits): class CTRLGenerator(): - def __init__(self): + def __init__(self, control_code="Horror Text: "): generate_num=256 model_dir = "generator/ctrl/model/seqlen256_v1.ckpt/" + self.control_code = control_code # load the vocabulary from file vocab = open('generator/ctrl/model/vocab', encoding='utf-8').read().split('\n') @@ -125,7 +126,7 @@ class CTRLGenerator(): # almost there, we now take the user prompt and tokenize with BPE # load BPE codes - self.bpe = fastBPE.fastBPE('codes', 'vocab') + self.bpe = fastBPE('codes', 'vocab') self.temperature = 0 self.nucleusprob = 0 @@ -134,6 +135,7 @@ class CTRLGenerator(): def generate(self, prompt): + prompt = self.control_code + prompt # tokenize provided prompt split_prompt = self.bpe.apply([prompt])[0].split()