You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
258 lines
11 KiB
Python
258 lines
11 KiB
Python
from typing import List, Optional
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from src.utils.generation_algorithms import DecodingAlgorithm, GreedyAlgorithm, NucleusAlgorithm, TopKAlgorithm
|
|
from src.utils.generation_constraints import ABCBloomConstraint, EosConstraint, MaxNewTokensConstraint
|
|
|
|
|
|
class RemoteGenerationMixin:
|
|
"""
|
|
A class containing all functions for auto-regressive text generation, to be used as a mixin in [`BloomForCausalLM`].
|
|
The class exposes can be used for:
|
|
- *greedy decoding*.
|
|
- *multinomial sampling*.
|
|
|
|
This class is similar to transformer's [`generation_utils.GenerationMixin`], it can be used instead of it. However, it has some differences.
|
|
"""
|
|
|
|
@torch.no_grad()
|
|
def generate(
|
|
self,
|
|
inputs: Optional[torch.Tensor] = None,
|
|
do_sample: Optional[bool] = None,
|
|
temperature: float = 1.0,
|
|
top_k: Optional[int] = None,
|
|
top_p: Optional[float] = None,
|
|
bos_token_id: Optional[int] = None,
|
|
eos_token_id: Optional[int] = None,
|
|
pad_token_id: Optional[int] = None,
|
|
max_length: Optional[int] = None,
|
|
max_new_tokens: Optional[int] = None,
|
|
decoding_algorithm: Optional[DecodingAlgorithm] = None,
|
|
provided_constraints: List[ABCBloomConstraint] = [],
|
|
**model_kwargs,
|
|
) -> torch.LongTensor:
|
|
"""
|
|
Generates sequences of token ids for models with a language modeling head.
|
|
|
|
:param inputs: The input tokens to the model.
|
|
:param do_sample: Whether to sample from the model predictions or take the argmax.
|
|
:param temperature: The temperature to use for sampling.
|
|
:param top_k: The number of results to return.
|
|
:param top_p: The cumulative probability of results to return.
|
|
:param bos_token_id: The id of the beginning of sentence token.
|
|
:param eos_token_id: The id of the end of sentence token.
|
|
:param pad_token_id: The id of the padding token.
|
|
:param max_new_tokens: The maximum number of tokens to generate.
|
|
:param decoding_algorithm: The decoding algorithm to use.
|
|
:param provided_constraints: A list of constraints to use.
|
|
:param model_kwargs: Additional arguments to pass to the model.
|
|
"""
|
|
|
|
assert (
|
|
model_kwargs.get("logits_processor", None) is None
|
|
), "For RemoteGenerationMixin models use BloomConstraints instead of logits_processor"
|
|
assert (
|
|
model_kwargs.get("logits_wrapper", None) is None
|
|
), "For RemoveGenerationMixin models use DecodingAlgorithm instead of logits_wrapper"
|
|
assert (
|
|
model_kwargs.get("stopping_criteria", None) is None
|
|
), "For RemoteGenerationMixin models use BloomConstraints instead of stopping_criteria"
|
|
if inputs is not None:
|
|
assert isinstance(inputs, torch.Tensor) and inputs.ndim == 2, "inputs must be a 2d tensor [batch, length]"
|
|
prefix_length = 0 if inputs is None else inputs.size(1)
|
|
prefix_length += self.config.pre_seq_len
|
|
|
|
bos_token_id = bos_token_id if bos_token_id is not None else self.config.bos_token_id
|
|
pad_token_id = pad_token_id if pad_token_id is not None else self.config.pad_token_id
|
|
eos_token_id = eos_token_id if eos_token_id is not None else self.config.eos_token_id
|
|
|
|
assert (max_length is None) != (max_new_tokens is None), "please set max_length or max_new_tokens (not both)"
|
|
if max_length is not None and max_new_tokens is None:
|
|
max_new_tokens = max_length - prefix_length
|
|
assert max_new_tokens > 0, f"Provided max_length is less than prefix size: {max_length} < {inputs.size(1)}"
|
|
elif max_length is None and max_new_tokens is not None:
|
|
max_length = prefix_length + max_new_tokens
|
|
|
|
if inputs is None:
|
|
assert bos_token_id is not None, "You have to provide a bos_token_id if you do not provide inputs"
|
|
inputs = torch.tensor([[bos_token_id]])
|
|
|
|
if decoding_algorithm is None:
|
|
if do_sample:
|
|
decoding_algorithm = self._choose_sample_algorithm(temperature, top_k, top_p)
|
|
else:
|
|
decoding_algorithm = GreedyAlgorithm()
|
|
|
|
constraints = self._get_constraints(
|
|
inputs=inputs,
|
|
eos_token_id=eos_token_id,
|
|
pad_token_id=pad_token_id,
|
|
max_new_tokens=max_new_tokens,
|
|
provided_constraints=provided_constraints,
|
|
)
|
|
|
|
with self.transformer.h.inference_session(max_length=max_length) as sess:
|
|
outputs = []
|
|
if torch.any(inputs == pad_token_id): # TODO: move to prepare_inputs
|
|
outputs += [inputs[:, : inputs.size(1) - (inputs == pad_token_id).sum(-1).max()]]
|
|
else:
|
|
outputs += [inputs]
|
|
last_token_id = None
|
|
seq_idx = outputs[0].size(1)
|
|
hypo_ids = torch.arange(outputs[0].size(0))
|
|
while True:
|
|
embs = self.transformer.word_embeddings(outputs[-1])
|
|
intermediate_prompts = None
|
|
if self.config.pre_seq_len > 0 and len(outputs) == 1:
|
|
prompts, intermediate_prompts = self.transformer.get_prompt(embs.size(0))
|
|
embs = torch.cat([prompts, embs], dim=1)
|
|
embs = self.transformer.word_embeddings_layernorm(embs)
|
|
hidden_state = sess.step(embs, prompts=intermediate_prompts, hypo_ids=hypo_ids)[:, -1]
|
|
hidden_state = self.transformer.ln_f(hidden_state)
|
|
lm_logits = self.lm_head(hidden_state)
|
|
|
|
for constraint in constraints:
|
|
lm_logits = constraint(last_token_id, lm_logits, hypo_ids)
|
|
last_token_id, hypo_ids = decoding_algorithm(lm_logits)
|
|
if seq_idx < inputs.size(1): # TODO: why is it not a constraint?
|
|
pad_token_mask = inputs[:, seq_idx : seq_idx + 1] == pad_token_id
|
|
last_token_id = (~pad_token_mask) * inputs[
|
|
:, seq_idx : seq_idx + 1
|
|
] + pad_token_mask * last_token_id
|
|
|
|
if torch.all(last_token_id == eos_token_id):
|
|
break
|
|
|
|
outputs.append(last_token_id)
|
|
seq_idx += 1
|
|
|
|
return torch.cat(outputs, dim=-1)
|
|
|
|
def greedy_search(
|
|
self,
|
|
input_ids: torch.LongTensor,
|
|
max_length: Optional[int] = None,
|
|
pad_token_id: Optional[int] = None,
|
|
eos_token_id: Optional[int] = None,
|
|
provided_constraints: List[ABCBloomConstraint] = [],
|
|
**model_kwargs,
|
|
) -> torch.LongTensor:
|
|
"""
|
|
Generates sequences of token ids for models with a language modeling head. Uses greedy search.
|
|
|
|
:param input_ids: The input tokens to the model.
|
|
:param max_length: The maximum length of the sequence to generate.
|
|
:param pad_token_id: The id of the padding token.
|
|
:param eos_token_id: The id of the end of sentence token.
|
|
:param provided_constraints: A list of constraints to use.
|
|
"""
|
|
return self.generate(
|
|
inputs=input_ids,
|
|
max_new_tokens=max_length,
|
|
pad_token_id=pad_token_id,
|
|
eos_token_id=eos_token_id,
|
|
decoding_algorithm=GreedyAlgorithm(),
|
|
provided_constraints=provided_constraints,
|
|
**model_kwargs,
|
|
)
|
|
|
|
def sample(
|
|
self,
|
|
input_ids: torch.LongTensor,
|
|
temperature: float = 1.0,
|
|
top_k: Optional[int] = None,
|
|
top_p: Optional[float] = None,
|
|
max_length: Optional[int] = None,
|
|
pad_token_id: Optional[int] = None,
|
|
eos_token_id: Optional[int] = None,
|
|
provided_constraints: List[ABCBloomConstraint] = [],
|
|
**model_kwargs,
|
|
) -> torch.LongTensor:
|
|
"""
|
|
Generates sequences of token ids for models with a language modeling head. Uses sampling. Uses multinomial sampling algorithm. If top_k is provided, uses top_k sampling. If top_p is provided, uses nucleus sampling.
|
|
|
|
:param: input_ids: The input tokens to the model.
|
|
:param: temperature: The temperature to use for sampling.
|
|
:param: top_k: The number of samples to use for top_k sampling.
|
|
:param: top_p: The probability of using top_p sampling.
|
|
:param: max_length: The maximum length of the sequence to generate.
|
|
:param: pad_token_id: The id of the padding token.
|
|
:param: eos_token_id: The id of the end of sentence token.
|
|
:param: provided_constraints: A list of constraints to use.
|
|
:param: model_kwargs: Additional kwargs to pass to the model.
|
|
"""
|
|
|
|
return self.generate(
|
|
inputs=input_ids,
|
|
max_new_tokens=max_length,
|
|
pad_token_id=pad_token_id,
|
|
eos_token_id=eos_token_id,
|
|
decoding_algorithm=self._choose_sample_algorithm(temperature, top_k, top_p),
|
|
provided_constraints=provided_constraints,
|
|
**model_kwargs,
|
|
)
|
|
|
|
def beam_search(
|
|
self,
|
|
input_ids: torch.LongTensor,
|
|
max_length: Optional[int] = None,
|
|
pad_token_id: Optional[int] = None,
|
|
eos_token_id: Optional[int] = None,
|
|
provided_constraints: List[ABCBloomConstraint] = [],
|
|
**model_kwargs,
|
|
) -> torch.LongTensor:
|
|
raise NotImplementedError
|
|
|
|
def beam_sample(
|
|
self,
|
|
input_ids: torch.LongTensor,
|
|
max_length: Optional[int] = None,
|
|
pad_token_id: Optional[int] = None,
|
|
eos_token_id: Optional[int] = None,
|
|
provided_constraints: List[ABCBloomConstraint] = [],
|
|
**model_kwargs,
|
|
) -> torch.LongTensor:
|
|
raise NotImplementedError
|
|
|
|
def group_beam_search(
|
|
self,
|
|
input_ids: torch.LongTensor,
|
|
max_length: Optional[int] = None,
|
|
pad_token_id: Optional[int] = None,
|
|
eos_token_id: Optional[int] = None,
|
|
provided_constraints: List[ABCBloomConstraint] = [],
|
|
**model_kwargs,
|
|
) -> torch.LongTensor:
|
|
raise NotImplementedError
|
|
|
|
def _choose_sample_algorithm(
|
|
self,
|
|
temperature: float = 1.0,
|
|
top_k: Optional[int] = None,
|
|
top_p: Optional[float] = None,
|
|
) -> DecodingAlgorithm:
|
|
if (top_k is not None) and (top_p is not None):
|
|
raise ValueError("You have to provide only top_k or top_p for sampling")
|
|
if top_k:
|
|
return TopKAlgorithm(top_k, temperature)
|
|
elif top_p:
|
|
return NucleusAlgorithm(top_p, temperature)
|
|
|
|
def _get_constraints(
|
|
self,
|
|
inputs: Optional[torch.Tensor] = None,
|
|
eos_token_id: Optional[int] = None,
|
|
pad_token_id: Optional[int] = None,
|
|
max_new_tokens: Optional[int] = None,
|
|
provided_constraints: List[ABCBloomConstraint] = [],
|
|
) -> List[ABCBloomConstraint]:
|
|
constraints = []
|
|
constraints.extend(provided_constraints)
|
|
if max_new_tokens is not None:
|
|
constraints.append(MaxNewTokensConstraint(inputs, max_new_tokens, eos_token_id, pad_token_id))
|
|
constraints.append(EosConstraint(inputs, eos_token_id, pad_token_id))
|
|
return constraints
|