Skip to content

Default values for some fields of assistant to target translator #11

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jmamou merged 4 commits into from
Jan 9, 2025
Merged

Default values for some fields of assistant to target translator #11

Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
dc601db
default values for AssistantToTargetTranslator fileds
jmamou Dec 23, 2024
3fc2c63
fix
jmamou Dec 24, 2024
0606453
add support to empty logit_processors
jmamou Dec 26, 2024
6d94c8c
Merge branch 'usd' into default_AssistantToTargetTranslator
jmamou Jan 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 18 additions & 10 deletions src/transformers/generation/candidate_generator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -627,21 +627,28 @@ def __init__(
self,
target_tokenizer: "PreTrainedTokenizerBase",
assistant_tokenizer: "PreTrainedTokenizerBase",
assistant_model_device,
target_vocab_size: int,
assistant_model_device: str = "cpu",
target_vocab_size: Optional[int] = None,
filter_value: float = -float("Inf"),
suppress_tokens_id: int = -1,
):
self._target_tokenizer: "PreTrainedTokenizerBase" = target_tokenizer
self._assistant_tokenizer: "PreTrainedTokenizerBase" = assistant_tokenizer
self._assistant_model_device = assistant_model_device
self.target_vocab_size: int = target_vocab_size
self._assistant_model_device: str = assistant_model_device
if target_vocab_size:
self.target_vocab_size: int = target_vocab_size
else:
self.target_vocab_size: int = len(self._target_tokenizer.get_vocab())
self.filter_value: float = filter_value
self.suppress_tokens_id: int = suppress_tokens_id
self._assistant_to_target_input_ids = self._get_assistant_to_target_input_ids()
self.logits_processors: LogitsProcessorList = LogitsProcessorList(
[SuppressTokensLogitsProcessor(self._get_suppress_input_ids(), self._assistant_model_device)]
)
self._suppress_input_ids: list[int] = self._get_suppress_input_ids()
self.logits_processors: Optional[LogitsProcessorList] = None
if len(self._suppress_input_ids) > 0:
# len(self._suppress_input_ids) = 0 if the assistant vocab is a subset of the target vocab
self.logits_processors = LogitsProcessorList(
[SuppressTokensLogitsProcessor(self._get_suppress_input_ids(), self._assistant_model_device)]
)

def _get_assistant_to_target_input_ids(self):
target_vocab = self._target_tokenizer.get_vocab()
Expand Down Expand Up @@ -707,8 +714,8 @@ def get_translator(
cls,
target_tokenizer: "PreTrainedTokenizerBase",
assistant_tokenizer: "PreTrainedTokenizerBase",
assistant_model_device,
target_vocab_size: int,
assistant_model_device: str = "cpu",
target_vocab_size: Optional[int] = None,
) -> AssistantToTargetTranslator:
with cls._lock:
assistant_dict = cls._cache.get(target_tokenizer)
Expand Down Expand Up @@ -801,7 +808,8 @@ def get_candidates(self, input_ids: torch.LongTensor) -> Tuple[torch.LongTensor,
generation_args["generation_config"].return_dict_in_generate = True

# Generate and process outputs using translator
generation_args["logits_processor"] = self._atm_translator.logits_processors
if self._atm_translator.logits_processors is not None:
generation_args["logits_processor"] = self._atm_translator.logits_processors
self._prev_assistant_ids, assistant_candidate_logits = self._generate_candidates(generation_args)

# Use translator to convert tokens and logits
Expand Down