This commit is contained in:
Nick
2019-09-19 11:14:23 -06:00
parent fa7104f2cf
commit b8cee2875d
+91 -93
View File
@@ -154,110 +154,108 @@ class CTRLGenerator():
padded_text = text + [0] * (self.generate_num - len(text))
tokens_generated = np.tile(padded_text, (1, 1))
result = ""
try:
for token in range(len(text) - 1, self.generate_num - 1):
# get the logits from the prediction function
# the logic here is a bit convoluted because we are allowing generation past 512 tokens
# this is done by sliding the window over (past 512 tokens) and continuing prediction
# I'm sure this can be simplified (TODO)
if token <= self.seq_length:
prompt_logits = self.predict_fn({'input_1': tokens_generated[:, :self.seq_length]})[
'tied_embedding_softmax'].squeeze() / (self.temperature if self.temperature > 0 else 1.)
_token = token if token < self.seq_length else -1
else:
_token = -1
end = token + 1
start = token - self.seq_length + 2
prompt_logits = \
self.predict_fn({'input_1': np.hstack((tokens_generated[:, 0:1], tokens_generated[:, start:end]))})[
'tied_embedding_softmax'].squeeze() / (self.temperature if self.temperature > 0 else 1.)
for token in range(len(text) - 1, self.generate_num - 1):
# get the logits from the prediction function
# the logic here is a bit convoluted because we are allowing generation past 512 tokens
# this is done by sliding the window over (past 512 tokens) and continuing prediction
# I'm sure this can be simplified (TODO)
if token <= self.seq_length:
prompt_logits = self.predict_fn({'input_1': tokens_generated[:, :self.seq_length]})[
'tied_embedding_softmax'].squeeze() / (self.temperature if self.temperature > 0 else 1.)
_token = token if token < self.seq_length else -1
else:
_token = -1
end = token + 1
start = token - self.seq_length + 2
prompt_logits = \
self.predict_fn({'input_1': np.hstack((tokens_generated[:, 0:1], tokens_generated[:, start:end]))})[
'tied_embedding_softmax'].squeeze() / (self.temperature if self.temperature > 0 else 1.)
# if penalty (for repetition) is non-zero,
# discount the logits from already generated tokens
if self.penalty > 0:
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:
continue
penalized_so_far.add(generated_token)
prompt_logits[_token][generated_token] /= self.penalty
# if penalty (for repetition) is non-zero,
# discount the logits from already generated tokens
if self.penalty > 0:
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:
continue
penalized_so_far.add(generated_token)
prompt_logits[_token][generated_token] /= self.penalty
# disallow some tokens
prompt_logits[_token][self.word2idx['<unk>']] = -1e8
# disallow some tokens
prompt_logits[_token][self.word2idx['<unk>']] = -1e8
# sometimes, when generating from reddit,
# it tries to generate the Score (reddit Karma) immediately after generating the Title:
# to disallow this, we can just prevent it from generating Score
prompt_logits[_token][self.word2idx['Sco@@']] = -1e8
# sometimes, when generating from reddit,
# it tries to generate the Score (reddit Karma) immediately after generating the Title:
# to disallow this, we can just prevent it from generating Score
prompt_logits[_token][self.word2idx['Sco@@']] = -1e8
# compute probabilities from logits
prompt_probs = np.exp(prompt_logits[_token])
prompt_probs = prompt_probs / sum(prompt_probs)
pruned_list = np.argsort(prompt_probs)[::-1]
# if you are using nucleus prob, then compute the nucleus probability size
if self.nucleusprob > 0.:
minimum_topk = 1
nucleus = max(np.where(np.cumsum(np.sort(prompt_probs)[::-1]) > self.nucleusprob)[0][0], minimum_topk)
elif self.topk > 0:
# we are over-loading notation here
# if you choose to specify a topk instead of a nucleus,
# we will hardcode the nucleus to be just that
nucleus = self.topk
else:
# if you specify neither nucleus or topk,
# then we will use the whole list
nucleus = len(pruned_list)
# compute probabilities from logits
prompt_probs = np.exp(prompt_logits[_token])
prompt_probs = prompt_probs / sum(prompt_probs)
pruned_list = np.argsort(prompt_probs)[::-1]
# if you are using nucleus prob, then compute the nucleus probability size
if self.nucleusprob > 0.:
minimum_topk = 1
nucleus = max(np.where(np.cumsum(np.sort(prompt_probs)[::-1]) > self.nucleusprob)[0][0], minimum_topk)
elif self.topk > 0:
# we are over-loading notation here
# if you choose to specify a topk instead of a nucleus,
# we will hardcode the nucleus to be just that
nucleus = self.topk
else:
# if you specify neither nucleus or topk,
# then we will use the whole list
nucleus = len(pruned_list)
# if you want to disallow more complex tokens, you can do so here
# for instance, if you want to disallow anything with the phrase `http`,
# you can delete theme from the pruned_list
# you can comment this out, I'm keeping it in for demonstration purpose
tokens_to_disallow = []
for _ in range(len(pruned_list)):
if 'http' in self.idx2word[pruned_list[_]]:
tokens_to_disallow.append(_)
pruned_list = np.delete(pruned_list, tokens_to_disallow)
# if you want to disallow more complex tokens, you can do so here
# for instance, if you want to disallow anything with the phrase `http`,
# you can delete theme from the pruned_list
# you can comment this out, I'm keeping it in for demonstration purpose
tokens_to_disallow = []
for _ in range(len(pruned_list)):
if 'http' in self.idx2word[pruned_list[_]]:
tokens_to_disallow.append(_)
pruned_list = np.delete(pruned_list, tokens_to_disallow)
# if temperature is 0
# just pick the first (most probable) token
if self.temperature == 0:
idx = pruned_list[0]
else:
# else,
# sample from the pruned_list with the logits
chosen_idx = int(tf.random.categorical(np.expand_dims(prompt_logits[0][_token][pruned_list], 0),
num_samples=1).numpy())
idx = pruned_list[chosen_idx]
# if temperature is 0
# just pick the first (most probable) token
if self.temperature == 0:
idx = pruned_list[0]
else:
# else,
# sample from the pruned_list with the logits
chosen_idx = int(tf.random.categorical(np.expand_dims(prompt_logits[0][_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:', idx2word[idx])
# print('top25 alternatives:', pruned_list[:25])
# if you want to do some debugging,
# like which one was chosen,
# what the top25 were,
# here is your opportunity.
# print('chosen:', idx2word[idx])
# print('top25 alternatives:', pruned_list[:25])
# assign the token for generation
tokens_generated[0][token + 1] = idx
# assign the token for generation
tokens_generated[0][token + 1] = idx
# clear screen if you want to
# os.system("clear")
# clear screen if you want to
# os.system("clear")
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)
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)
result = tokens_generated_so_far[prompt_length:]
print(tokens_generated_so_far)
result = tokens_generated_so_far[prompt_length:]
return result