diff --git a/generator/ctrl/ctrl_generator.py b/generator/ctrl/ctrl_generator.py index 8bd39ce..ebbbc2a 100644 --- a/generator/ctrl/ctrl_generator.py +++ b/generator/ctrl/ctrl_generator.py @@ -22,14 +22,13 @@ def loss(labels, logits): class CTRLGenerator(): - def __init__(self, control_code="Horror Text: ", generate_num=100, temperature=0.75, topk=0): + def __init__(self, control_code="Horror Text: ", generate_num=100, temperature=0.75): self.generate_num=generate_num model_dir = "generator/ctrl/model/seqlen256_v1.ckpt/" self.control_code = control_code vocab_file = 'generator/ctrl/model/vocab' code_file = 'generator/ctrl/model/codes' - self.topk=topk # load the vocabulary from file vocab = open(vocab_file, encoding='utf-8').read().split('\n') @@ -137,6 +136,7 @@ class CTRLGenerator(): self.temperature=temperature self.nucleusprob = 0 self.penalty = 1.2 + self.topk=0 def configure_verb_probs(self, probabilities, options): @@ -175,7 +175,7 @@ class CTRLGenerator(): result = remove_profanity(result) if not first_letter_capitalized: - result = result[0].upper() + result[1:] + result = result[0].lower() + result[1:] print("\n\nAFTER RESULT_REPLACE:") print(repr(result)) @@ -183,15 +183,19 @@ class CTRLGenerator(): return result def generate(self, prompt, options=None): - - if options is None: - options = {} - - first_token = True prompt = self.prompt_replace(prompt) - if "used_verbs" in options: - print("Disallowed used verbs: ", options["used_verbs"]) + if prompt[-1] != " ": + prompt = prompt + " " + first_token = True + + prompt = second_to_first_person(prompt) + + prompt = self.control_code + prompt + + print("******************************") + print(" DEBUG:: Prompt to generate by is \n", prompt) + print("******************************") prompt_length = len(prompt) @@ -204,6 +208,8 @@ class CTRLGenerator(): padded_text = text + [0] * (total_text_len - len(text)) tokens_generated = np.tile(padded_text, (1, 1)) result = "" + max_new_lines = 5 + num_new_lines = 0 for token in range(len(text) - 1, total_text_len - 1): if first_token: @@ -233,6 +239,12 @@ class CTRLGenerator(): penalized_so_far = set() for _ in range(token + 1): generated_token = tokens_generated[0][_] + # don't penalize newlines + # you could also choose not to penalize frequent words + # (which incidentally are sorted in the vocab file) + # but I don't do that + # if it prints too many new lines instead of continuing generating text, + # you might want to comment this out if self.idx2word[generated_token] == '\n': continue if generated_token in penalized_so_far: @@ -243,11 +255,16 @@ class CTRLGenerator(): # disallow some tokens forbidden_tokens = ['', 'Sco@@'] + if num_new_lines >= max_new_lines: + forbidden_tokens.append('\n') + for forbidden_token in forbidden_tokens: prompt_logits[_token][self.word2idx[forbidden_token]] = -1e8 + # Make sure only a possible verb is chosen. if first_token: - prompt_logits[_token] = self.configure_verb_probs(prompt_logits[_token], options) + for word in get_possible_verbs(): + prompt_logits[_token][self.word2idx[word]] += 5 # compute probabilities from logits prompt_probs = np.exp(prompt_logits[_token]) @@ -288,6 +305,16 @@ class CTRLGenerator(): tf.random.categorical(np.expand_dims(prompt_logits[_token][pruned_list], 0), num_samples=1).numpy()) idx = pruned_list[chosen_idx] + # if you want to do some debugging, + # like which one was chosen, + # what the top25 were, + # here is your opportunity. + # print('chosen:', repr(self.idx2word[idx])) + # print('top25 alternatives:', pruned_list[:25]) + + if self.idx2word[idx] == "\n": + num_new_lines += 1 + # assign the token for generation tokens_generated[0][token + 1] = idx @@ -297,10 +324,9 @@ class CTRLGenerator(): tokens_generated_so_far = ' '.join([self.idx2word[c] for c in tokens_generated[0].squeeze()[:token + 2]]) tokens_generated_so_far = re.sub('(@@ )', '', string=tokens_generated_so_far) tokens_generated_so_far = re.sub('(@@ ?$)', '', string=tokens_generated_so_far) - #print(tokens_generated_so_far) + print(tokens_generated_so_far) result = tokens_generated_so_far[prompt_length:] first_token = False - return self.result_replace(result)