wehub-resource-sync
158 lines
6.1 KiB
Python
158 lines
6.1 KiB
Python
import jittor as jt
|
|
|
|
def generate(moss, input_str, tokenizer, method, **kwargs):
|
|
"""
|
|
Choose different methods to generate sentences.
|
|
|
|
:param input_str: The input text.
|
|
:param tokenizer: Tokenizer.
|
|
:param method: Generation method. Should be one of: ['greedy', 'sample']
|
|
:param kwargs: Other parameters used for generation.
|
|
- max_gen_len: int. Maximum generate length. Used in all methods.
|
|
- temperature: float. Used in ``sample``.
|
|
- top_p: float. Used in ``sample``.
|
|
- top_k: int. Used in ``sample``.
|
|
"""
|
|
if method == "greedy":
|
|
return greedy_search(moss, input_str, tokenizer, **kwargs)
|
|
elif method == "sample":
|
|
return sample(moss, input_str, tokenizer, **kwargs)
|
|
else:
|
|
raise NotImplementedError(
|
|
f"Unsupported generation method {method}"
|
|
)
|
|
|
|
def greedy_search(model, input_str, tokenizer, max_gen_len,
|
|
eos_token_id=None, pad_token_id=None):
|
|
model.eval()
|
|
if eos_token_id is None:
|
|
eos_token_id = tokenizer.eos_token_id
|
|
if pad_token_id is None and eos_token_id is not None:
|
|
pad_token_id = eos_token_id
|
|
eos_token_id_tensor = jt.Var(eos_token_id)
|
|
|
|
tokenized = tokenizer(input_str, return_tensors='np')
|
|
sentence_ids = jt.Var(tokenized['input_ids'])
|
|
attention_mask = jt.Var(tokenized['attention_mask'])
|
|
unfinished_sequences = sentence_ids.new(sentence_ids.shape[0]).fill_(1)
|
|
past_key_values = None
|
|
while True:
|
|
# set input
|
|
if past_key_values:
|
|
input_ids = sentence_ids[:, -1].unsqueeze(-1)
|
|
else:
|
|
input_ids = sentence_ids
|
|
|
|
outputs = model(input_ids, past_key_values=past_key_values,
|
|
attention_mask=attention_mask)
|
|
# caculate probs
|
|
next_token_logits = outputs['logits'][:, -1, :].float()
|
|
next_tokens = jt.argmax(next_token_logits, dim=-1)[0]
|
|
|
|
# concat sentence
|
|
next_tokens = next_tokens * unfinished_sequences + \
|
|
pad_token_id * (1 - unfinished_sequences)
|
|
sentence_ids = jt.cat([sentence_ids, next_tokens[:, None]], dim=-1)
|
|
# update input
|
|
past_key_values = outputs['past_key_values']
|
|
attention_mask = jt.cat(
|
|
[attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1)
|
|
|
|
# if eos_token was found in one sentence, set sentence to finished
|
|
next_tokens.repeat(eos_token_id_tensor.shape[0], 1)
|
|
unfinished_sequences = unfinished_sequences.mul(
|
|
next_tokens.repeat(eos_token_id_tensor.shape[0], 1) \
|
|
.not_equal(eos_token_id_tensor.unsqueeze(1)) \
|
|
.prod(dim=0)
|
|
)
|
|
|
|
jt.sync_all()
|
|
|
|
if unfinished_sequences.max() == 0 or sentence_ids.shape[-1] >= max_gen_len:
|
|
break
|
|
|
|
return sentence_ids.reshape([-1,]).tolist()[tokenized['input_ids'].shape[1]:]
|
|
|
|
def sample(model, input_str, tokenizer, max_gen_len, temperature, top_p, top_k,
|
|
eos_token_id=None, pad_token_id=None):
|
|
model.eval()
|
|
if eos_token_id is None:
|
|
eos_token_id = tokenizer.eos_token_id
|
|
if pad_token_id is None and eos_token_id is not None:
|
|
pad_token_id = eos_token_id
|
|
eos_token_id_tensor = jt.Var(eos_token_id)
|
|
|
|
tokenized = tokenizer(input_str, return_tensors='np')
|
|
sentence_ids = jt.Var(tokenized['input_ids'])
|
|
attention_mask = jt.Var(tokenized['attention_mask'])
|
|
unfinished_sequences = sentence_ids.new(sentence_ids.shape[0]).fill_(1)
|
|
past_key_values = None
|
|
|
|
while True:
|
|
|
|
# set input
|
|
if past_key_values:
|
|
input_ids = sentence_ids[:, -1].unsqueeze(-1)
|
|
else:
|
|
input_ids = sentence_ids
|
|
outputs = model(input_ids, past_key_values=past_key_values,
|
|
attention_mask=attention_mask)
|
|
|
|
next_token_logits = outputs['logits'][:, -1, :].float()
|
|
|
|
# sample
|
|
# temperature
|
|
scores = next_token_logits / temperature
|
|
# top_k
|
|
scores = sample_top_k(scores, top_k)
|
|
# top_p
|
|
scores = sample_top_p(scores, top_p)
|
|
|
|
probs = jt.nn.softmax(scores, dim=-1)
|
|
next_tokens = jt.multinomial(probs, num_samples=1).squeeze(1)
|
|
# concat sentence
|
|
next_tokens = next_tokens * unfinished_sequences + \
|
|
pad_token_id * (1 - unfinished_sequences)
|
|
|
|
# update generated ids, model inputs, and length for next step
|
|
sentence_ids = jt.cat([sentence_ids, next_tokens[:, None]], dim=-1)
|
|
past_key_values = outputs['past_key_values']
|
|
attention_mask = jt.cat(
|
|
[attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1)
|
|
|
|
# if eos_token was found in one sentence, set sentence to finished
|
|
next_tokens.repeat(eos_token_id_tensor.shape[0], 1)
|
|
unfinished_sequences = unfinished_sequences.mul(
|
|
next_tokens.repeat(eos_token_id_tensor.shape[0], 1) \
|
|
.not_equal(eos_token_id_tensor.unsqueeze(1)) \
|
|
.prod(dim=0)
|
|
)
|
|
|
|
jt.sync_all()
|
|
|
|
if unfinished_sequences.max() == 0 or sentence_ids.shape[-1] >= max_gen_len:
|
|
break
|
|
|
|
return sentence_ids.reshape([-1,]).tolist()[tokenized['input_ids'].shape[1]:]
|
|
|
|
def sample_top_k(scores, top_k):
|
|
top_k = min(top_k, scores.size(-1)) # Safety check
|
|
# Remove all tokens with a probability less than the last token of the top-k
|
|
indices_to_remove = scores < jt.topk(scores, top_k)[0][..., -1, None]
|
|
scores = scores.masked_fill(indices_to_remove, -float("Inf"))
|
|
|
|
return scores
|
|
|
|
def sample_top_p(scores, top_p):
|
|
sorted_logits, sorted_indices = jt.sort(scores, descending=False)
|
|
cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
|
|
|
|
# Remove tokens with cumulative top_p above the threshold (token with 0 are kept)
|
|
sorted_indices_to_remove = cumulative_probs <= (1 - top_p)
|
|
|
|
# scatter sorted tensors to original indexing
|
|
indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
|
|
scores = scores.masked_fill(indices_to_remove, -float("Inf"))
|
|
|
|
return scores
|