116 lines
5.3 KiB
Text
116 lines
5.3 KiB
Text
diff --git a/src/transformers/generation/configuration_utils.py b/src/transformers/generation/configuration_utils.py
|
|
index f8702dde29..19a4194586 100644
|
|
--- a/src/transformers/generation/configuration_utils.py
|
|
+++ b/src/transformers/generation/configuration_utils.py
|
|
@@ -131,9 +131,17 @@ class GenerationConfig(PushToHubMixin):
|
|
the current pass after allocated time has been passed.
|
|
stop_strings (`str or list[str]`, *optional*):
|
|
A string or a list of strings that should terminate generation if the model outputs them.
|
|
+ entropy_threshold (`float`, *optional*):
|
|
+ The threshold for the Shannon entropy of the model's output scores (logits). If the entropy
|
|
+ exceeds this threshold for `entropy_patience` consecutive steps, generation for that sequence
|
|
+ stops.
|
|
+ entropy_patience (`int`, *optional*, defaults to 1):
|
|
+ The number of consecutive steps the entropy must exceed `entropy_threshold` before stopping
|
|
+ generation.
|
|
|
|
> Parameters that control the generation strategy used
|
|
|
|
+
|
|
do_sample (`bool`):
|
|
Whether or not to use sampling ; use greedy decoding otherwise.
|
|
num_beams (`int`, *optional*):
|
|
@@ -358,8 +366,11 @@ class GenerationConfig(PushToHubMixin):
|
|
self.early_stopping = kwargs.pop("early_stopping", None)
|
|
self.max_time = kwargs.pop("max_time", None)
|
|
self.stop_strings = kwargs.pop("stop_strings", None)
|
|
+ self.entropy_threshold = kwargs.pop("entropy_threshold", None)
|
|
+ self.entropy_patience = kwargs.pop("entropy_patience", 1)
|
|
|
|
# Parameters that control the generation strategy used
|
|
+
|
|
self.do_sample = kwargs.pop("do_sample", None)
|
|
self.num_beams = kwargs.pop("num_beams", None)
|
|
|
|
@@ -582,7 +593,10 @@ class GenerationConfig(PushToHubMixin):
|
|
"assistant_confidence_threshold": 0.4,
|
|
"assistant_lookbehind": 10,
|
|
"target_lookbehind": 10,
|
|
+ "entropy_threshold": None,
|
|
+ "entropy_patience": 1,
|
|
# Deprecated arguments (moved to the Hub). TODO joao, manuel: remove in v4.62.0
|
|
+
|
|
"num_beam_groups": 1,
|
|
"diversity_penalty": 0.0,
|
|
}
|
|
diff --git a/src/transformers/generation/stopping_criteria.py b/src/transformers/generation/stopping_criteria.py
|
|
index b57136f534..41bd5616f1 100644
|
|
--- a/src/transformers/generation/stopping_criteria.py
|
|
+++ b/src/transformers/generation/stopping_criteria.py
|
|
@@ -492,6 +492,40 @@ class ConfidenceCriteria(StoppingCriteria):
|
|
return False
|
|
|
|
|
|
+class EntropyStoppingCriteria(StoppingCriteria):
|
|
+ """
|
|
+ This class can be used to stop generation whenever the entropy of the scores (logits) exceeds a threshold
|
|
+ for a given number of consecutive steps.
|
|
+
|
|
+ Args:
|
|
+ entropy_threshold (`float`):
|
|
+ The threshold for the Shannon entropy of the scores.
|
|
+ entropy_patience (`int`, *optional*, defaults to 1):
|
|
+ The number of consecutive steps the entropy must exceed the threshold before stopping generation.
|
|
+ """
|
|
+
|
|
+ def __init__(self, entropy_threshold: float, entropy_patience: int = 1):
|
|
+ self.entropy_threshold = entropy_threshold
|
|
+ self.entropy_patience = entropy_patience
|
|
+ self.consecutive_steps = None
|
|
+
|
|
+ @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING)
|
|
+ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.BoolTensor:
|
|
+ batch_size = input_ids.shape[0]
|
|
+ if self.consecutive_steps is None or self.consecutive_steps.shape[0] != batch_size:
|
|
+ self.consecutive_steps = torch.zeros(batch_size, device=input_ids.device, dtype=torch.long)
|
|
+
|
|
+ # Compute Shannon entropy
|
|
+ probs = F.softmax(scores, dim=-1)
|
|
+ log_probs = F.log_softmax(scores, dim=-1)
|
|
+ entropy = -(probs * log_probs).sum(dim=-1)
|
|
+
|
|
+ over_threshold = entropy > self.entropy_threshold
|
|
+ self.consecutive_steps = torch.where(over_threshold, self.consecutive_steps + 1, 0)
|
|
+
|
|
+ return self.consecutive_steps >= self.entropy_patience
|
|
+
|
|
+
|
|
class StoppingCriteriaList(list):
|
|
@add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING)
|
|
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.BoolTensor:
|
|
diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py
|
|
index ffb7266a5b..def46827ff 100644
|
|
--- a/src/transformers/generation/utils.py
|
|
+++ b/src/transformers/generation/utils.py
|
|
@@ -100,6 +100,7 @@ from .logits_process import (
|
|
)
|
|
from .stopping_criteria import (
|
|
ConfidenceCriteria,
|
|
+ EntropyStoppingCriteria,
|
|
EosTokenCriteria,
|
|
MaxLengthCriteria,
|
|
MaxTimeCriteria,
|
|
@@ -1292,6 +1293,13 @@ class GenerationMixin(ContinuousMixin):
|
|
criteria.append(
|
|
ConfidenceCriteria(assistant_confidence_threshold=generation_config.assistant_confidence_threshold)
|
|
)
|
|
+ if generation_config.entropy_threshold is not None:
|
|
+ criteria.append(
|
|
+ EntropyStoppingCriteria(
|
|
+ entropy_threshold=generation_config.entropy_threshold,
|
|
+ entropy_patience=generation_config.entropy_patience,
|
|
+ )
|
|
+ )
|
|
criteria = self._merge_criteria_processor_list(criteria, stopping_criteria)
|
|
return criteria
|
|
|