From 14cf685678ec3fbbc19314b1049697eb8a24e4ee Mon Sep 17 00:00:00 2001 From: Peter Chapman Date: Thu, 27 Aug 2026 09:13:34 +1200 Subject: [PATCH 01/14] Migrate to transformers 5 --- local_check.sh | 2 +- .../hugging_face_nmt_model_factory.py | 5 +- machine/jobs/nmt_build_options.py | 2 +- machine/jobs/nmt_engine_build_job.py | 3 +- machine/jobs/settings.yaml | 2 +- .../custom_normalizer/tokenizer_config.json | 2 +- .../huggingface/hugging_face_nmt_engine.py | 175 +++++++--- .../huggingface/hugging_face_nmt_model.py | 7 +- .../hugging_face_nmt_model_trainer.py | 98 +++--- poetry.lock | 324 +++++++++++------- pyproject.toml | 4 +- .../test_hugging_face_nmt_engine.py | 36 +- .../test_hugging_face_nmt_model_trainer.py | 51 +-- 13 files changed, 406 insertions(+), 305 deletions(-) diff --git a/local_check.sh b/local_check.sh index 69aa3825..c606be14 100755 --- a/local_check.sh +++ b/local_check.sh @@ -15,7 +15,7 @@ if [ "$#" -ne 0 ]; then exit 2 fi -poetry install +poetry install --all-extras if [ "$agent_strict" = true ]; then echo "=================== comment hygiene =================" diff --git a/machine/jobs/huggingface/hugging_face_nmt_model_factory.py b/machine/jobs/huggingface/hugging_face_nmt_model_factory.py index 9b0e3356..5a0a23ef 100644 --- a/machine/jobs/huggingface/hugging_face_nmt_model_factory.py +++ b/machine/jobs/huggingface/hugging_face_nmt_model_factory.py @@ -6,8 +6,8 @@ import datasets.utils.logging as datasets_logging import transformers.utils.logging as transformers_logging from transformers import AutoConfig, AutoModelForSeq2SeqLM, HfArgumentParser, PreTrainedModel, Seq2SeqTrainingArguments -from transformers.integrations import ClearMLCallback -from transformers.tokenization_utils import TruncationStrategy +from transformers.integrations.integration_utils import ClearMLCallback +from transformers.tokenization_utils_base import TruncationStrategy from ...corpora.parallel_text_corpus import ParallelTextCorpus from ...corpora.text_corpus import TextCorpus @@ -26,7 +26,6 @@ def __init__(self, config: Any) -> None: self._config = config args = config.huggingface.train_params.to_dict() args["output_dir"] = str(self._model_dir) - args["overwrite_output_dir"] = True # Use "max_steps" from root for backward compatibility if "max_steps" in self._config.huggingface: args["max_steps"] = self._config.huggingface.max_steps diff --git a/machine/jobs/nmt_build_options.py b/machine/jobs/nmt_build_options.py index 03a777eb..d0954de6 100644 --- a/machine/jobs/nmt_build_options.py +++ b/machine/jobs/nmt_build_options.py @@ -15,7 +15,7 @@ class TrainParams(BaseModel): per_device_train_batch_size: int | None = None gradient_accumulation_steps: int | None = None label_smoothing_factor: float | None = None - group_by_length: bool | None = None + train_sampling_strategy: str | None = None gradient_checkpointing: bool | None = None lr_scheduler_type: str | None = None learning_rate: float | None = None diff --git a/machine/jobs/nmt_engine_build_job.py b/machine/jobs/nmt_engine_build_job.py index 69304f32..cf25d4b7 100644 --- a/machine/jobs/nmt_engine_build_job.py +++ b/machine/jobs/nmt_engine_build_job.py @@ -149,7 +149,8 @@ def _translate( if check_canceled is not None: check_canceled() source_segments = [pt_info["translation"] for pt_info in pt_batch] - for pt_info, result in zip(pt_batch, engine.translate_batch(source_segments), strict=True): + t_batch = engine.translate_batch(source_segments) + for pt_info, result in zip(pt_batch, t_batch, strict=True): pt_info["translation"] = result.translation pt_info["sequenceConfidence"] = result.sequence_confidence current_inference_step += len(pt_batch) diff --git a/machine/jobs/settings.yaml b/machine/jobs/settings.yaml index cdddcf3d..18174033 100644 --- a/machine/jobs/settings.yaml +++ b/machine/jobs/settings.yaml @@ -13,7 +13,7 @@ default: per_device_train_batch_size: 64 gradient_accumulation_steps: 1 label_smoothing_factor: 0.2 - group_by_length: true + train_sampling_strategy : group_by_length gradient_checkpointing: true lr_scheduler_type: cosine learning_rate: 0.0002 diff --git a/machine/translation/huggingface/custom_normalizer/tokenizer_config.json b/machine/translation/huggingface/custom_normalizer/tokenizer_config.json index 742aa330..7d0fc87c 100644 --- a/machine/translation/huggingface/custom_normalizer/tokenizer_config.json +++ b/machine/translation/huggingface/custom_normalizer/tokenizer_config.json @@ -1,5 +1,5 @@ { - "additional_special_tokens": null, + "extra_special_tokens": null, "bos_token": "", "cls_token": "", "eos_token": "", diff --git a/machine/translation/huggingface/hugging_face_nmt_engine.py b/machine/translation/huggingface/hugging_face_nmt_engine.py index 0ed23519..e087f692 100644 --- a/machine/translation/huggingface/hugging_face_nmt_engine.py +++ b/machine/translation/huggingface/hugging_face_nmt_engine.py @@ -4,7 +4,8 @@ import logging import re from math import exp, prod -from typing import Collection, Iterable, List, Optional, Sequence, Tuple, Union, cast +from pathlib import Path +from typing import Any, Collection, Iterable, List, Optional, Sequence, Tuple, Union, cast import torch # pyright: ignore[reportMissingImports] from sacremoses import MosesPunctNormalizer @@ -14,18 +15,16 @@ AutoTokenizer, M2M100Tokenizer, NllbTokenizer, - NllbTokenizerFast, + Pipeline, PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast, - TranslationPipeline, ) -from transformers.generation import BeamSearchEncoderDecoderOutput, GreedySearchEncoderDecoderOutput -from transformers.tokenization_utils import BatchEncoding, TruncationStrategy +from transformers.generation.utils import GenerateBeamEncoderDecoderOutput, GenerateEncoderDecoderOutput +from transformers.tokenization_utils_base import TruncationStrategy from ...annotations.range import Range from ...corpora.aligned_word_pair import AlignedWordPair -from ...utils.typeshed import StrPath from ..translation_engine import TranslationEngine from ..translation_result import TranslationResult from ..translation_result_builder import TranslationResultBuilder @@ -38,23 +37,30 @@ class HuggingFaceNmtEngine(TranslationEngine): def __init__( self, - model: Union[PreTrainedModel, StrPath, str], + model: Union[PreTrainedModel, Path, str], oom_batch_size_backoff_mult: float = 1.0, **pipeline_kwargs, ) -> None: - self._model = model self._pipeline_kwargs = pipeline_kwargs - if isinstance(self._model, PreTrainedModel): + if isinstance(model, PreTrainedModel): + self._model = model self._model.eval() self._is_model_owned = False else: - model_config = AutoConfig.from_pretrained(str(self._model), label2id={}, id2label={}, num_labels=0) + model_config = AutoConfig.from_pretrained(str(model), label2id={}, id2label={}, num_labels=0) + + # If output_attentions is True or None, we need to set the attn_implementation to eager to get the attentions + attn_implementation = "eager" if self._pipeline_kwargs.get("output_attentions", True) else "sdpa" + self._model = cast( - PreTrainedModel, AutoModelForSeq2SeqLM.from_pretrained(str(self._model), config=model_config) + PreTrainedModel, + AutoModelForSeq2SeqLM.from_pretrained( + str(model), config=model_config, attn_implementation=attn_implementation + ), ) self._is_model_owned = True - self._tokenizer = AutoTokenizer.from_pretrained(self._model.name_or_path, use_fast=True) - if isinstance(self._tokenizer, (NllbTokenizer, NllbTokenizerFast)): + self._tokenizer = AutoTokenizer.from_pretrained(self._model.name_or_path) + if isinstance(self._tokenizer, NllbTokenizer): self._mpn = MosesPunctNormalizer() self._mpn.substitutions = [ # type: ignore (re.compile(r), sub) @@ -70,11 +76,12 @@ def __init__( src_lang is not None and tgt_lang is not None and "prefix" not in self._pipeline_kwargs + and self._model.name_or_path is not None and (self._model.name_or_path.startswith("t5-") or self._model.name_or_path.startswith("google/mt5-")) ): self._pipeline_kwargs["prefix"] = f"translate {src_lang} to {tgt_lang}: " else: - additional_special_tokens = cast(list[str], self._tokenizer.additional_special_tokens or []) + extra_special_tokens = cast(list[str], self._tokenizer.extra_special_tokens or []) if isinstance(self._tokenizer, M2M100Tokenizer): src_lang_token = self._tokenizer.lang_code_to_token.get(src_lang) if src_lang is not None else None tgt_lang_token = self._tokenizer.lang_code_to_token.get(tgt_lang) if tgt_lang is not None else None @@ -84,14 +91,14 @@ def __init__( if ( src_lang is not None and src_lang_token not in self._tokenizer.added_tokens_encoder - and src_lang_token not in additional_special_tokens + and src_lang_token not in extra_special_tokens ): raise ValueError(f"The specified model does not support the language code '{src_lang}'") if ( tgt_lang is not None and tgt_lang_token not in self._tokenizer.added_tokens_encoder - and tgt_lang_token not in additional_special_tokens + and tgt_lang_token not in extra_special_tokens ): raise ValueError(f"The specified model does not support the language code '{tgt_lang}'") @@ -101,7 +108,7 @@ def __init__( self._pipeline = _TranslationPipeline( model=self._model, - tokenizer=self._tokenizer, + tokenizer=cast(PreTrainedTokenizer, self.tokenizer), mpn=self._mpn, batch_size=self._batch_size, **self._pipeline_kwargs, @@ -189,17 +196,44 @@ def close(self) -> None: torch.cuda.empty_cache() -class _TranslationPipeline(TranslationPipeline): +class _TranslationPipeline(Pipeline): def __init__( self, - model: Union[PreTrainedModel, StrPath, str], - tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast], - batch_size: int, + model: PreTrainedModel, + tokenizer: PreTrainedTokenizer, mpn: Optional[MosesPunctNormalizer] = None, **kwargs, ) -> None: - super().__init__(model=model, tokenizer=tokenizer, batch_size=batch_size, **kwargs) + if isinstance(model, (str, Path)): + model = cast(PreTrainedModel, AutoModelForSeq2SeqLM.from_pretrained(model)) + self.device: torch.device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu") + self.model: Any = cast(PreTrainedModel, cast(torch.nn.Module, model).to(self.device)) + self.tokenizer = tokenizer self._mpn = mpn + self._kwargs = kwargs + + def __call__(self, sentences: Iterable[Any], **generate_kwargs) -> List[List[dict]]: + model_inputs = self.preprocess(*sentences).to(self.device) + with torch.no_grad(): + model_outputs = self._forward(model_inputs, **generate_kwargs) + return self.postprocess(model_outputs) + + def _sanitize_parameters(self, **pipeline_parameters): + return pipeline_parameters + + def check_inputs(self, input_length: int, min_length: Optional[int], max_length: Optional[int]) -> None: + """ + Checks whether there might be something wrong with given input with regard to the model. + """ + if min_length is not None and max_length is not None and max_length < min_length: + logger.warning(f"Your min_length={min_length} must be inferior than your max_length={max_length}.") + + if max_length is not None and input_length < max_length: + logger.warning( + f"Your max_length is set to {max_length}, but your input_length is only {input_length}. Since this is " + "a summarization task, where outputs shorter than the input are typically wanted, you might " + f"consider decreasing max_length manually, e.g. summarizer('...', max_length={input_length // 2})" + ) def preprocess(self, *args, truncation=TruncationStrategy.DO_NOT_TRUNCATE, src_lang=None, tgt_lang=None): if self.tokenizer is None: @@ -222,31 +256,38 @@ def preprocess(self, *args, truncation=TruncationStrategy.DO_NOT_TRUNCATE, src_l ) for s in args ] - inputs = cast( - BatchEncoding, - super().preprocess(*sentences, truncation=truncation, src_lang=src_lang, tgt_lang=tgt_lang), + + # The source language prefix and the forced target language token are configured on the tokenizer and the + # generation config in _configure_model, so a plain tokenizer call is all that is needed here. + return self.tokenizer( + cast(list[str] | list[list[str]], sentences), + return_tensors="pt", + truncation=truncation, + padding=True, ) - if inputs.encodings is not None: - inputs["input_tokens"] = [ - _get_encoding_fast_tokens(inputs.encodings[i]) if isinstance(args[i], str) else args[i] - for i in range(len(args)) - ] - else: - inputs["input_tokens"] = [self.tokenizer.tokenize(s) if isinstance(s, str) else s for s in args] - return inputs def _forward(self, model_inputs, **generate_kwargs): in_b, input_length = model_inputs["input_ids"].shape - input_tokens = model_inputs["input_tokens"] - del model_inputs["input_tokens"] + if "input_tokens" in model_inputs: + input_tokens = model_inputs.pop("input_tokens") + else: + input_tokens = [self.tokenizer.convert_ids_to_tokens(seq) for seq in model_inputs["input_ids"]] + if hasattr(self.model, "generation_config") and self.model.generation_config is not None: config = self.model.generation_config else: config = self.model.config - generate_kwargs["min_length"] = generate_kwargs.get("min_length", config.min_length) - generate_kwargs["max_length"] = generate_kwargs.get("max_length", config.max_length) - generate_kwargs["output_attentions"] = generate_kwargs.get("output_attentions", True) + generate_kwargs["min_length"] = generate_kwargs.get( + "min_length", self._kwargs.get("min_length", config.min_length or 0) + ) + generate_kwargs["max_length"] = generate_kwargs.get( + "max_length", self._kwargs.get("max_length", config.max_length or 200) + ) + generate_kwargs["output_attentions"] = generate_kwargs.get( + "output_attentions", self._kwargs.get("output_attentions", True) + ) + generate_kwargs["num_beams"] = self._kwargs.get("num_beams", config.num_beams or 1) self.check_inputs(input_length, generate_kwargs["min_length"], generate_kwargs["max_length"]) output = self.model.generate( **model_inputs, @@ -255,14 +296,14 @@ def _forward(self, model_inputs, **generate_kwargs): return_dict_in_generate=True, ) - if isinstance(output, BeamSearchEncoderDecoderOutput): + if isinstance(output, GenerateBeamEncoderDecoderOutput): output_ids = output.sequences beam_indices = output.beam_indices scores = output.scores assert scores is not None and beam_indices is not None sequences_scores = output.sequences_scores attentions = output.cross_attentions - elif isinstance(output, GreedySearchEncoderDecoderOutput): + elif isinstance(output, GenerateEncoderDecoderOutput): output_ids = output.sequences beam_indices = None assert output.scores is not None @@ -272,15 +313,30 @@ def _forward(self, model_inputs, **generate_kwargs): else: raise RuntimeError("Cannot postprocess the output of the model.") - transition_scores = cast( - torch.Tensor, - self.model.compute_transition_scores( - output_ids, # type: ignore - scores, # type: ignore - beam_indices, # type: ignore - normalize_logits=True, - ), - ) + try: + transition_scores = cast( + torch.Tensor, + self.model.compute_transition_scores( + output_ids, + scores, + beam_indices, + normalize_logits=True, + ), + ) + except Exception: + output_ids = output_ids.to("cpu") + scores = tuple(score.to("cpu") for score in scores) + beam_indices = beam_indices.to("cpu") if beam_indices is not None else None + transition_scores = cast( + torch.Tensor, + self.model.compute_transition_scores( + output_ids, + scores, + beam_indices, + normalize_logits=True, + ), + ) + sequences_scores = getattr(output, "sequences_scores", None) if beam_indices is None: beam_indices = torch.zeros_like(output_ids) @@ -312,12 +368,23 @@ def _forward(self, model_inputs, **generate_kwargs): if generate_kwargs["output_attentions"] is True: assert attentions is not None num_heads = attentions[0][0].shape[1] + + # Truncate/Pad beam_indices to match output_ids length exact slice + target_seq_len = output_ids.shape[1] - start_index + sliced_beam_indices = beam_indices[:, start_index:] + if sliced_beam_indices.shape[1] > target_seq_len: + sliced_beam_indices = sliced_beam_indices[:, :target_seq_len] + elif sliced_beam_indices.shape[1] < target_seq_len: + sliced_beam_indices = torch.nn.functional.pad( + sliced_beam_indices, (0, target_seq_len - sliced_beam_indices.shape[1]) + ) + indices = torch.stack( ( torch.arange(output_ids.shape[1] - start_index, device=output_ids.device).expand( in_b, n_sequences, -1 ), - torch.reshape(beam_indices[:, start_index:] % num_beams, (in_b, n_sequences, -1)), + torch.reshape(sliced_beam_indices % num_beams, (in_b, n_sequences, -1)), ), dim=3, ) @@ -354,7 +421,7 @@ def _forward(self, model_inputs, **generate_kwargs): "attentions": attentions, } - def postprocess(self, model_outputs, clean_up_tokenization_spaces=False): + def postprocess(self, model_outputs, clean_up_tokenization_spaces=False) -> List[List[dict]]: if self.tokenizer is None: raise RuntimeError("No tokenizer is specified.") all_special_ids = set(self.tokenizer.all_special_ids) @@ -411,7 +478,7 @@ def postprocess(self, model_outputs, clean_up_tokenization_spaces=False): records.append(record) - return records + return [records] def torch_gather_nd(params: torch.Tensor, indices: torch.Tensor, batch_dim: int = 0) -> torch.Tensor: @@ -444,7 +511,3 @@ def torch_gather_nd(params: torch.Tensor, indices: torch.Tensor, batch_dim: int out = torch.gather(params, dim=batch_dim, index=indices) return out.reshape(*index_shape, *tail_sizes) - - -def _get_encoding_fast_tokens(encoding) -> List[str]: - return [token for (token, mask) in zip(encoding.tokens, encoding.special_tokens_mask) if not mask] diff --git a/machine/translation/huggingface/hugging_face_nmt_model.py b/machine/translation/huggingface/hugging_face_nmt_model.py index 253eb9be..26d65897 100644 --- a/machine/translation/huggingface/hugging_face_nmt_model.py +++ b/machine/translation/huggingface/hugging_face_nmt_model.py @@ -8,7 +8,6 @@ from transformers import PreTrainedModel, Seq2SeqTrainingArguments from ...corpora.parallel_text_corpus import ParallelTextCorpus -from ...utils.typeshed import StrPath from ..translation_model import TranslationModel from ..translation_result import TranslationResult from .hugging_face_nmt_engine import HuggingFaceNmtEngine @@ -18,14 +17,14 @@ class HuggingFaceNmtModel(TranslationModel): def __init__( self, - model: Union[PreTrainedModel, StrPath], + model: Union[PreTrainedModel, Path], parent_model_name: str, training_args: Optional[Seq2SeqTrainingArguments] = None, **pipeline_kwargs, ) -> None: self._model = model if isinstance(model, PreTrainedModel): - self._model_path = Path(model.name_or_path) + self._model_path = Path(str(model.name_or_path)) else: self._model_path = Path(model) self._parent_model_name = parent_model_name @@ -89,7 +88,7 @@ def __init__(self, model: HuggingFaceNmtModel, corpus: Union[ParallelTextCorpus, def save(self) -> None: super().save() - output_dir = Path(self._model.training_args.output_dir) + output_dir = Path(str(self._model.training_args.output_dir)) if output_dir != self._model._model_path: shutil.copytree(output_dir, self._model._model_path) self._model.reset_engine() diff --git a/machine/translation/huggingface/hugging_face_nmt_model_trainer.py b/machine/translation/huggingface/hugging_face_nmt_model_trainer.py index 26e9bf94..0f69e55b 100644 --- a/machine/translation/huggingface/hugging_face_nmt_model_trainer.py +++ b/machine/translation/huggingface/hugging_face_nmt_model_trainer.py @@ -25,11 +25,8 @@ M2M100ForConditionalGeneration, M2M100Tokenizer, MBart50Tokenizer, - MBart50TokenizerFast, MBartTokenizer, - MBartTokenizerFast, NllbTokenizer, - NllbTokenizerFast, PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerBase, @@ -40,7 +37,7 @@ TrainerCallback, set_seed, ) -from transformers.tokenization_utils import BatchEncoding +from transformers.tokenization_utils_base import BatchEncoding from transformers.trainer_callback import TrainerControl, TrainerState from transformers.trainer_utils import get_last_checkpoint from transformers.training_args import TrainingArguments @@ -76,12 +73,9 @@ def prepare_decoder_input_ids_from_labels(self: M2M100ForConditionalGeneration, MULTILINGUAL_TOKENIZERS = ( MBartTokenizer, - MBartTokenizerFast, MBart50Tokenizer, - MBart50TokenizerFast, M2M100Tokenizer, NllbTokenizer, - NllbTokenizerFast, ) @@ -125,17 +119,21 @@ def train( check_canceled: Optional[Callable[[], None]] = None, ) -> None: last_checkpoint = None - if os.path.isdir(self._training_args.output_dir) and not self._training_args.overwrite_output_dir: + if ( + self._training_args.output_dir is not None + and os.path.isdir(self._training_args.output_dir) + and self._training_args.resume_from_checkpoint is not None + ): last_checkpoint = get_last_checkpoint(self._training_args.output_dir) if last_checkpoint is None and any(os.path.isfile(p) for p in os.listdir(self._training_args.output_dir)): raise ValueError( f"Output directory ({self._training_args.output_dir}) already exists and is not empty. " - "Use --overwrite_output_dir to overcome." + "Remove --resume_from_checkpoint to overcome." ) elif last_checkpoint is not None and self._training_args.resume_from_checkpoint is None: logger.info( f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change " - "the `--output_dir` or add `--overwrite_output_dir` to train from scratch." + "the `--output_dir` or remove `--resume_from_checkpoint` to train from scratch." ) # Set seed before initializing model. @@ -156,7 +154,7 @@ def train( model = cast(PreTrainedModel, AutoModelForSeq2SeqLM.from_pretrained(self._model, config=config)) logger.info("Initializing tokenizer") - tokenizer = AutoTokenizer.from_pretrained(model.name_or_path, use_fast=True) + tokenizer = AutoTokenizer.from_pretrained(model.name_or_path) src_lang = self._src_lang if src_lang is None: @@ -176,19 +174,22 @@ def train( def find_missing_characters(tokenizer: Any, train_dataset: Dataset, lang_codes: List[str]) -> List[str]: vocab = tokenizer.get_vocab().keys() charset = set() - mpn_normalize = True if isinstance(tokenizer, (NllbTokenizerFast)) else False + mpn_normalize = True if isinstance(tokenizer, NllbTokenizer) else False for ex in train_dataset["translation"]: for lang_code in lang_codes: ex_text = ex[lang_code] if mpn_normalize: ex_text = self._mpn.normalize(ex_text) - ex_text = tokenizer.backend_tokenizer.normalizer.normalize_str(ex_text) + if tokenizer.backend_tokenizer.normalizer is not None: + ex_text = tokenizer.backend_tokenizer.normalizer.normalize_str(ex_text) charset = charset | set(ex_text) charset = set(filter(None, {char.strip() for char in charset})) missing_characters = sorted(list(charset - vocab)) return missing_characters def add_tokens(tokenizer: Any, missing_tokens: List[str]) -> Any: + if self._training_args.output_dir is None: + raise ValueError("Missing output_dir from training arguments") tokenizer_dir = Path(self._training_args.output_dir) tokenizer.save_pretrained(str(tokenizer_dir)) with open(tokenizer_dir / "tokenizer.json", "r+", encoding="utf-8") as file: @@ -204,7 +205,7 @@ def add_tokens(tokenizer: Any, missing_tokens: List[str]) -> Any: json.dump(data, file, ensure_ascii=False, indent=4) file.truncate() logger.info(f"Added {len(missing_tokens)} tokens to the tokenizer: {missing_tokens}") - return AutoTokenizer.from_pretrained(str(tokenizer_dir), use_fast=True) + return AutoTokenizer.from_pretrained(str(tokenizer_dir)) if self._add_unk_src_tokens or self._add_unk_tgt_tokens: logger.info("Checking for missing tokens") @@ -215,8 +216,7 @@ def add_tokens(tokenizer: Any, missing_tokens: List[str]) -> Any: ) else: norm_tok = PreTrainedTokenizerFast.from_pretrained( - str(Path(os.path.dirname(os.path.abspath(__file__))) / "custom_normalizer"), - use_fast=True, + str(Path(os.path.dirname(os.path.abspath(__file__))) / "custom_normalizer") ) # using unofficially supported behavior to set the normalizer lang_codes = [] @@ -249,12 +249,9 @@ def add_tokens(tokenizer: Any, missing_tokens: List[str]) -> Any: if ( self._tgt_lang is not None and model.config.decoder_start_token_id is None - and isinstance(tokenizer, (MBartTokenizer, MBartTokenizerFast)) + and isinstance(tokenizer, MBartTokenizer) ): - if isinstance(tokenizer, MBartTokenizer): - model.config.decoder_start_token_id = tokenizer.lang_code_to_id[self._tgt_lang] - else: - model.config.decoder_start_token_id = tokenizer.convert_tokens_to_ids(self._tgt_lang) + model.config.decoder_start_token_id = tokenizer.convert_tokens_to_ids(self._tgt_lang) if model.config.decoder_start_token_id is None: raise ValueError("Make sure that `config.decoder_start_token_id` is correctly defined") @@ -280,15 +277,17 @@ def add_tokens(tokenizer: Any, missing_tokens: List[str]) -> Any: model.config.forced_bos_token_id = forced_bos_token_id prefix = "" - if model.name_or_path.startswith("t5-") or model.name_or_path.startswith("google/mt5-"): + if model.name_or_path is not None and ( + model.name_or_path.startswith("t5-") or model.name_or_path.startswith("google/mt5-") + ): prefix = f"translate {self._src_lang} to {self._tgt_lang}: " max_src_length = self.max_src_length if max_src_length is None: - max_src_length = model.config.max_length + max_src_length = model.generation_config.max_length or 200 max_tgt_length = self.max_tgt_length if max_tgt_length is None: - max_tgt_length = model.config.max_length + max_tgt_length = model.generation_config.max_length or 200 if self._training_args.label_smoothing_factor > 0 and not hasattr( model, "prepare_decoder_input_ids_from_labels" @@ -304,37 +303,23 @@ def batch_prepare_for_model( batch_tokens: List[List[str]], return_tensors: Optional[Union[str, TensorType]] = None, ) -> BatchEncoding: - batch_outputs: Dict[str, Any] = {} - for tokens in batch_tokens: - ids = cast(List[int], tokenizer.convert_tokens_to_ids(tokens)) - outputs = tokenizer.prepare_for_model(ids, add_special_tokens=False) - - for key, value in outputs.items(): - if key not in batch_outputs: - batch_outputs[key] = [] - batch_outputs[key].append(value) - return BatchEncoding(batch_outputs, tensor_type=return_tensors) + return tokenizer( + batch_tokens, + is_split_into_words=True, + add_special_tokens=False, + return_tensors=return_tensors, + ) def preprocess_function(examples): # Add one to the content_type in order to convert back from ClassLabels which are enumerated from 0, not 1 - if isinstance(tokenizer, (NllbTokenizer, NllbTokenizerFast)): - inputs = [ - (self._mpn.normalize(ex[src_lang]), TextRowContentType(d + 1)) - for ex, d in zip(examples["translation"], examples["content_type"]) - ] - targets = [ - (self._mpn.normalize(ex[tgt_lang]), TextRowContentType(d + 1)) - for ex, d in zip(examples["translation"], examples["content_type"]) - ] - else: - inputs = [ - (self._mpn.normalize(ex[src_lang]), TextRowContentType(d + 1)) - for ex, d in zip(examples["translation"], examples["content_type"]) - ] - targets = [ - (self._mpn.normalize(ex[tgt_lang]), TextRowContentType(d + 1)) - for ex, d in zip(examples["translation"], examples["content_type"]) - ] + inputs = [ + (self._mpn.normalize(ex[src_lang]), TextRowContentType(d + 1)) + for ex, d in zip(examples["translation"], examples["content_type"]) + ] + targets = [ + (self._mpn.normalize(ex[tgt_lang]), TextRowContentType(d + 1)) + for ex, d in zip(examples["translation"], examples["content_type"]) + ] num_glosses = len([1 for _, d in inputs if d == TextRowContentType.WORD]) if not isinstance(tokenizer, PreTrainedTokenizerFast) or num_glosses == 0: @@ -503,6 +488,7 @@ def __init__( eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None, tokenizer: Optional[PreTrainedTokenizerBase] = None, model_init: Optional[Callable[[], PreTrainedModel]] = None, + compute_loss_func: Callable | None = None, compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None, callbacks: Optional[List[TrainerCallback]] = None, optimizers: Tuple[Optional[Optimizer], Optional[LambdaLR]] = (None, None), @@ -512,13 +498,14 @@ def __init__( model, args, data_collator, - train_dataset, # type: ignore + train_dataset, eval_dataset, # type: ignore tokenizer, model_init, + compute_loss_func, compute_metrics, callbacks, - optimizers, # type: ignore + optimizers, preprocess_logits_for_metrics, ) @@ -570,13 +557,12 @@ def add_lang_code_to_tokenizer(tokenizer: Union[PreTrainedTokenizer, PreTrainedT return tokenizer.add_special_tokens( - {"additional_special_tokens": tokenizer.additional_special_tokens + [lang_token]} # type: ignore + {"extra_special_tokens": tokenizer.extra_special_tokens + [lang_token]} # type: ignore ) lang_id = cast(int, tokenizer.convert_tokens_to_ids(lang_token)) if isinstance(tokenizer, (MBart50Tokenizer, MBartTokenizer)): tokenizer.lang_code_to_id[lang_code] = lang_id - tokenizer.id_to_lang_code[lang_id] = lang_code tokenizer.fairseq_tokens_to_ids[lang_code] = lang_id tokenizer.fairseq_ids_to_tokens[lang_id] = lang_code elif isinstance(tokenizer, M2M100Tokenizer): diff --git a/poetry.lock b/poetry.lock index 0cc34e54..a119e1a2 100644 --- a/poetry.lock +++ b/poetry.lock @@ -208,6 +208,19 @@ files = [ frozenlist = ">=1.1.0" typing-extensions = {version = ">=4.2", markers = "python_version < \"3.13\""} +[[package]] +name = "annotated-doc" +version = "0.0.5" +description = "Document parameters, class attributes, return types, and variables inline, with Annotated." +optional = true +python-versions = ">=3.9" +groups = ["main"] +markers = "extra == \"huggingface\"" +files = [ + {file = "annotated_doc-0.0.5-py3-none-any.whl", hash = "sha256:117bac03a25ede5df5440e855b32d556049ca169ead221505badf432fed4b101"}, + {file = "annotated_doc-0.0.5.tar.gz", hash = "sha256:c7e58ce09192557605d8bbd92836d7e1d520ac9580096042c0bfd197efacf1bb"}, +] + [[package]] name = "annotated-types" version = "0.8.0" @@ -227,12 +240,12 @@ version = "4.14.2" description = "High-level concurrency and networking framework on top of asyncio or Trio" optional = false python-versions = ">=3.10" -groups = ["main", "dev"] +groups = ["main", "dev", "gpu"] files = [ {file = "anyio-4.14.2-py3-none-any.whl", hash = "sha256:9f505dda5ac9f0c8309b5e8bd445a8c2bf7246f3ce950121e45ea15bc41d1494"}, {file = "anyio-4.14.2.tar.gz", hash = "sha256:cfa139f3ed1a23ee8f88a145ddb5ac7605b8bbfd8592baacd7ce3d8bb4313c7f"}, ] -markers = {main = "extra == \"huggingface\""} +markers = {main = "extra == \"huggingface\"", gpu = "sys_platform == \"win32\" or sys_platform == \"linux\""} [package.dependencies] exceptiongroup = {version = ">=1.0.2", markers = "python_version < \"3.11\""} @@ -680,12 +693,11 @@ version = "2.1.1" description = "The Real First Universal Charset Detector. Open, modern and actively maintained alternative to Chardet." optional = false python-versions = ">=3.6.0" -groups = ["main", "dev", "gpu"] +groups = ["main", "dev"] files = [ {file = "charset-normalizer-2.1.1.tar.gz", hash = "sha256:5a3d016c7c547f69d6f81fb0db9449ce888b418b5b9952cc5e6e66843e9dd845"}, {file = "charset_normalizer-2.1.1-py3-none-any.whl", hash = "sha256:83e9a75d1911279afd89352c68b45348559d1fc0506b054b346651b5e7fee29f"}, ] -markers = {gpu = "sys_platform == \"win32\" or sys_platform == \"linux\""} [package.extras] unicode-backport = ["unicodedata2"] @@ -732,12 +744,12 @@ version = "8.5.0" description = "Composable command line interface toolkit" optional = false python-versions = ">=3.10" -groups = ["main", "dev"] +groups = ["main", "dev", "gpu"] files = [ {file = "click-8.5.0-py3-none-any.whl", hash = "sha256:255bc9599cf7748b4b1a446ccc735421bd08a2ae529a8b88597d3de5664ee360"}, {file = "click-8.5.0.tar.gz", hash = "sha256:ba0d2089de75ea0310e2dde03160e6ca10009947fb95a182f9b54021bb272e34"}, ] -markers = {main = "extra == \"huggingface\""} +markers = {main = "extra == \"huggingface\"", gpu = "sys_platform == \"win32\" or sys_platform == \"linux\""} [[package]] name = "colorama" @@ -956,21 +968,21 @@ files = [ [[package]] name = "datasets" -version = "4.8.5" +version = "5.0.1" description = "HuggingFace community-driven open-source library of datasets" optional = true python-versions = ">=3.10.0" groups = ["main"] markers = "extra == \"huggingface\"" files = [ - {file = "datasets-4.8.5-py3-none-any.whl", hash = "sha256:5079900781719c0e063a8efdd2cd95a31ad0c63209178669cd23cf1b926149ff"}, - {file = "datasets-4.8.5.tar.gz", hash = "sha256:0f0c1c3d56ffff2c93b2f4c63c95bac94f3d7e8621aea2a2a576275233bba772"}, + {file = "datasets-5.0.1-py3-none-any.whl", hash = "sha256:9fbf73688f8c18f7529b4fe592abd04015f81d1e58001e4bac73ffb2b39d7cc4"}, + {file = "datasets-5.0.1.tar.gz", hash = "sha256:ce22bb851efd7494f08aad33b940803784434f6e77763d00679a0dc45fcf686a"}, ] [package.dependencies] dill = ">=0.3.0,<0.4.2" filelock = "*" -fsspec = {version = ">=2023.1.0,<=2026.2.0", extras = ["http"]} +fsspec = {version = ">=2023.1.0,<=2026.6.0", extras = ["http"]} httpx = "<1.0.0" huggingface-hub = ">=0.25.0,<2.0" multiprocess = "<0.70.20" @@ -986,16 +998,18 @@ xxhash = "*" [package.extras] audio = ["torch (>=2.8.0)", "torchcodec (>=0.6.0)"] benchmarks = ["tensorflow (==2.12.0)", "torch (==2.0.1)", "transformers (==4.30.1)"] -dev = ["Pillow (>=9.4.0)", "absl-py", "aiohttp", "decorator", "elasticsearch (>=7.17.12,<8.0.0)", "faiss-cpu (>=1.8.0.post1)", "h5py", "jax (>=0.3.14) ; sys_platform != \"win32\"", "jaxlib (>=0.3.14) ; sys_platform != \"win32\"", "joblib (<1.3.0)", "joblibspark ; python_version < \"3.14\"", "lz4 ; python_version < \"3.14\"", "moto[server]", "nibabel (>=5.3.1)", "numba (>=0.56.4) ; python_version < \"3.14\"", "polars[timezone] (>=0.20.0)", "protobuf (<4.0.0)", "py7zr", "pylance", "pyspark (>=3.4)", "pytest", "pytest-datadir", "pytest-xdist", "rarfile (>=4.0)", "ruff (>=0.3.0)", "sqlalchemy", "tensorflow (>=2.16.0) ; python_version >= \"3.10\" and sys_platform != \"win32\" and python_version < \"3.14\"", "tensorflow (>=2.6.0)", "tensorflow (>=2.6.0) ; python_version < \"3.10\" and sys_platform != \"win32\"", "tiktoken", "torch", "torch (>=2.8.0)", "torchcodec (>=0.7.0) ; python_version < \"3.14\"", "torchdata", "transformers", "transformers (>=4.42.0)", "zstandard"] +dev = ["Pillow (>=9.4.0)", "absl-py", "aiohttp", "decorator", "elasticsearch (>=7.17.12,<8.0.0)", "faiss-cpu (>=1.8.0.post1)", "h5py", "jax (>=0.3.14) ; sys_platform != \"win32\"", "jaxlib (>=0.3.14) ; sys_platform != \"win32\"", "joblib (<1.3.0)", "joblibspark ; python_version < \"3.14\"", "lz4 ; python_version < \"3.14\"", "moto[server]", "nibabel (>=5.3.1)", "numba (>=0.56.4) ; python_version < \"3.14\"", "polars[timezone] (>=0.20.0)", "protobuf (<4.0.0)", "py7zr", "pyiceberg[pyarrow,sql-sqlite]", "pylance", "pyspark (>=3.4)", "pytest", "pytest-datadir", "pytest-xdist", "rarfile (>=4.0)", "ruff (>=0.3.0)", "sqlalchemy", "teich (==0.1.5)", "tensorflow (>=2.16.0) ; python_version >= \"3.10\" and sys_platform != \"win32\" and python_version < \"3.14\"", "tensorflow (>=2.6.0)", "tensorflow (>=2.6.0) ; python_version < \"3.10\" and sys_platform != \"win32\"", "tiktoken", "torch", "torch (>=2.8.0)", "torchcodec (>=0.7.0) ; python_version < \"3.14\"", "torchdata", "transformers", "transformers (>=4.42.0)", "trimesh (>=4.10.0)", "zstandard"] docs = ["tensorflow (>=2.6.0)", "torch", "transformers"] +iceberg = ["pyiceberg (>=0.7.0)"] jax = ["jax (>=0.3.14)", "jaxlib (>=0.3.14)"] +mesh = ["trimesh (>=4.10.0)"] nibabel = ["ipyniivue (==2.4.2)", "nibabel (>=5.3.2)"] pdfs = ["pdfplumber (>=0.11.4)"] quality = ["ruff (>=0.3.0)"] tensorflow = ["tensorflow (>=2.6.0)"] tensorflow-gpu = ["tensorflow (>=2.6.0)"] -tests = ["Pillow (>=9.4.0)", "absl-py", "aiohttp", "decorator", "elasticsearch (>=7.17.12,<8.0.0)", "faiss-cpu (>=1.8.0.post1)", "h5py", "jax (>=0.3.14) ; sys_platform != \"win32\"", "jaxlib (>=0.3.14) ; sys_platform != \"win32\"", "joblib (<1.3.0)", "joblibspark ; python_version < \"3.14\"", "lz4 ; python_version < \"3.14\"", "moto[server]", "nibabel (>=5.3.1)", "numba (>=0.56.4) ; python_version < \"3.14\"", "polars[timezone] (>=0.20.0)", "protobuf (<4.0.0)", "py7zr", "pylance", "pyspark (>=3.4)", "pytest", "pytest-datadir", "pytest-xdist", "rarfile (>=4.0)", "sqlalchemy", "tensorflow (>=2.16.0) ; python_version >= \"3.10\" and sys_platform != \"win32\" and python_version < \"3.14\"", "tensorflow (>=2.6.0) ; python_version < \"3.10\" and sys_platform != \"win32\"", "tiktoken", "torch (>=2.8.0)", "torchcodec (>=0.7.0) ; python_version < \"3.14\"", "torchdata", "transformers (>=4.42.0)", "zstandard"] -tests-numpy2 = ["Pillow (>=9.4.0)", "absl-py", "aiohttp", "decorator", "elasticsearch (>=7.17.12,<8.0.0)", "h5py", "jax (>=0.3.14) ; sys_platform != \"win32\"", "jaxlib (>=0.3.14) ; sys_platform != \"win32\"", "joblib (<1.3.0)", "joblibspark ; python_version < \"3.14\"", "lz4 ; python_version < \"3.14\"", "moto[server]", "nibabel (>=5.3.1)", "numba (>=0.56.4) ; python_version < \"3.14\"", "polars[timezone] (>=0.20.0)", "protobuf (<4.0.0)", "py7zr", "pylance", "pyspark (>=3.4)", "pytest", "pytest-datadir", "pytest-xdist", "rarfile (>=4.0)", "sqlalchemy", "tiktoken", "torch (>=2.8.0)", "torchcodec (>=0.7.0) ; python_version < \"3.14\"", "torchdata", "transformers (>=4.42.0)", "zstandard"] +tests = ["Pillow (>=9.4.0)", "absl-py", "aiohttp", "decorator", "elasticsearch (>=7.17.12,<8.0.0)", "faiss-cpu (>=1.8.0.post1)", "h5py", "jax (>=0.3.14) ; sys_platform != \"win32\"", "jaxlib (>=0.3.14) ; sys_platform != \"win32\"", "joblib (<1.3.0)", "joblibspark ; python_version < \"3.14\"", "lz4 ; python_version < \"3.14\"", "moto[server]", "nibabel (>=5.3.1)", "numba (>=0.56.4) ; python_version < \"3.14\"", "polars[timezone] (>=0.20.0)", "protobuf (<4.0.0)", "py7zr", "pyiceberg[pyarrow,sql-sqlite]", "pylance", "pyspark (>=3.4)", "pytest", "pytest-datadir", "pytest-xdist", "rarfile (>=4.0)", "sqlalchemy", "teich (==0.1.5)", "tensorflow (>=2.16.0) ; python_version >= \"3.10\" and sys_platform != \"win32\" and python_version < \"3.14\"", "tensorflow (>=2.6.0) ; python_version < \"3.10\" and sys_platform != \"win32\"", "tiktoken", "torch (>=2.8.0)", "torchcodec (>=0.7.0) ; python_version < \"3.14\"", "torchdata", "transformers (>=4.42.0)", "trimesh (>=4.10.0)", "zstandard"] +tests-numpy2 = ["Pillow (>=9.4.0)", "absl-py", "aiohttp", "decorator", "elasticsearch (>=7.17.12,<8.0.0)", "h5py", "jax (>=0.3.14) ; sys_platform != \"win32\"", "jaxlib (>=0.3.14) ; sys_platform != \"win32\"", "joblib (<1.3.0)", "joblibspark ; python_version < \"3.14\"", "lz4 ; python_version < \"3.14\"", "moto[server]", "nibabel (>=5.3.1)", "numba (>=0.56.4) ; python_version < \"3.14\"", "polars[timezone] (>=0.20.0)", "protobuf (<4.0.0)", "py7zr", "pyiceberg[pyarrow,sql-sqlite]", "pylance", "pyspark (>=3.4)", "pytest", "pytest-datadir", "pytest-xdist", "rarfile (>=4.0)", "sqlalchemy", "teich (==0.1.5)", "tiktoken", "torch (>=2.8.0)", "torchcodec (>=0.7.0) ; python_version < \"3.14\"", "torchdata", "transformers (>=4.42.0)", "trimesh (>=4.10.0)", "zstandard"] torch = ["torch"] vision = ["Pillow (>=9.4.0)"] @@ -1143,12 +1157,12 @@ version = "1.3.1" description = "Backport of PEP 654 (exception groups)" optional = false python-versions = ">=3.7" -groups = ["main", "dev"] +groups = ["main", "dev", "gpu"] files = [ {file = "exceptiongroup-1.3.1-py3-none-any.whl", hash = "sha256:a7a39a3bd276781e98394987d3a5701d0c4edffb633bb7a5144577f82c773598"}, {file = "exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219"}, ] -markers = {main = "extra == \"huggingface\" and python_version == \"3.10\"", dev = "python_version == \"3.10\""} +markers = {main = "extra == \"huggingface\" and python_version == \"3.10\"", dev = "python_version == \"3.10\"", gpu = "(sys_platform == \"win32\" or sys_platform == \"linux\") and python_version == \"3.10\""} [package.dependencies] typing-extensions = {version = ">=4.6.0", markers = "python_version < \"3.13\""} @@ -1436,12 +1450,12 @@ version = "0.16.0" description = "A pure-Python, bring-your-own-I/O implementation of HTTP/1.1" optional = false python-versions = ">=3.8" -groups = ["main", "dev"] +groups = ["main", "dev", "gpu"] files = [ {file = "h11-0.16.0-py3-none-any.whl", hash = "sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86"}, {file = "h11-0.16.0.tar.gz", hash = "sha256:4e35b956cf45792e4caa5885e69fba00bdbc6ffafbfa020300e549b208ee5ff1"}, ] -markers = {main = "extra == \"huggingface\""} +markers = {main = "extra == \"huggingface\"", gpu = "sys_platform == \"win32\" or sys_platform == \"linux\""} [[package]] name = "hf-xet" @@ -1469,7 +1483,7 @@ files = [ {file = "hf_xet-1.6.0-cp38-abi3-win_arm64.whl", hash = "sha256:3dc3e35441ba395006af5aaacc40ef2e603c51ef46c3530b9156185f00935ea3"}, {file = "hf_xet-1.6.0.tar.gz", hash = "sha256:2e58454a340b3556dfa4972d5451aff4fba8dd42a236600ba1a1d2b1514f0fef"}, ] -markers = {main = "extra == \"huggingface\" and (platform_machine == \"x86_64\" or platform_machine == \"amd64\" or platform_machine == \"arm64\" or platform_machine == \"aarch64\")", gpu = "(sys_platform == \"win32\" or sys_platform == \"linux\") and (platform_machine == \"x86_64\" or platform_machine == \"amd64\" or platform_machine == \"arm64\" or platform_machine == \"aarch64\")"} +markers = {main = "extra == \"huggingface\" and (platform_machine == \"x86_64\" or platform_machine == \"amd64\" or platform_machine == \"AMD64\" or platform_machine == \"arm64\" or platform_machine == \"aarch64\")", gpu = "(sys_platform == \"win32\" or sys_platform == \"linux\") and (platform_machine == \"x86_64\" or platform_machine == \"amd64\" or platform_machine == \"AMD64\" or platform_machine == \"arm64\" or platform_machine == \"aarch64\")"} [package.extras] tests = ["pytest"] @@ -1480,12 +1494,12 @@ version = "1.0.9" description = "A minimal low-level HTTP client." optional = false python-versions = ">=3.8" -groups = ["main", "dev"] +groups = ["main", "dev", "gpu"] files = [ {file = "httpcore-1.0.9-py3-none-any.whl", hash = "sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55"}, {file = "httpcore-1.0.9.tar.gz", hash = "sha256:6e34463af53fd2ab5d807f399a9b45ea31c3dfa2276f15a2c3f00afff6e176e8"}, ] -markers = {main = "extra == \"huggingface\""} +markers = {main = "extra == \"huggingface\"", gpu = "sys_platform == \"win32\" or sys_platform == \"linux\""} [package.dependencies] certifi = "*" @@ -1503,12 +1517,12 @@ version = "0.28.1" description = "The next generation HTTP client." optional = false python-versions = ">=3.8" -groups = ["main", "dev"] +groups = ["main", "dev", "gpu"] files = [ {file = "httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad"}, {file = "httpx-0.28.1.tar.gz", hash = "sha256:75e98c5f16b0f35b567856f597f06ff2270a374470a5c2392242528e3e3e42fc"}, ] -markers = {main = "extra == \"huggingface\""} +markers = {main = "extra == \"huggingface\"", gpu = "sys_platform == \"win32\" or sys_platform == \"linux\""} [package.dependencies] anyio = "*" @@ -1525,43 +1539,40 @@ zstd = ["zstandard (>=0.18.0)"] [[package]] name = "huggingface-hub" -version = "0.36.2" +version = "1.28.0" description = "Client library to download and publish models, datasets and other repos on the huggingface.co hub" optional = false -python-versions = ">=3.8.0" +python-versions = ">=3.10.0" groups = ["main", "gpu"] files = [ - {file = "huggingface_hub-0.36.2-py3-none-any.whl", hash = "sha256:48f0c8eac16145dfce371e9d2d7772854a4f591bcb56c9cf548accf531d54270"}, - {file = "huggingface_hub-0.36.2.tar.gz", hash = "sha256:1934304d2fb224f8afa3b87007d58501acfda9215b334eed53072dd5e815ff7a"}, + {file = "huggingface_hub-1.28.0-py3-none-any.whl", hash = "sha256:58a8bacb03072edfc38067065e9dc24bbb34805410fcd36a1632de0b329660bb"}, + {file = "huggingface_hub-1.28.0.tar.gz", hash = "sha256:46a2e950c09234de54093d587d1675382f0d08dbd600d9fb599b5932f5b2c6cb"}, ] markers = {main = "extra == \"huggingface\"", gpu = "sys_platform == \"win32\" or sys_platform == \"linux\""} [package.dependencies] -filelock = "*" +click = ">=8.4.2,<9.0.0" +filelock = ">=3.10.0" fsspec = ">=2023.5.0" -hf-xet = {version = ">=1.1.3,<2.0.0", markers = "platform_machine == \"x86_64\" or platform_machine == \"amd64\" or platform_machine == \"arm64\" or platform_machine == \"aarch64\""} +hf-xet = {version = ">=1.5.2,<2.0.0", markers = "platform_machine == \"x86_64\" or platform_machine == \"amd64\" or platform_machine == \"AMD64\" or platform_machine == \"arm64\" or platform_machine == \"aarch64\""} +httpx = ">=0.23.0,<1" packaging = ">=20.9" pyyaml = ">=5.1" -requests = "*" tqdm = ">=4.42.1" -typing-extensions = ">=3.7.4.3" +typing-extensions = ">=4.1.0" [package.extras] -all = ["InquirerPy (==0.3.4)", "Jinja2", "Pillow", "aiohttp", "authlib (>=1.3.2)", "fastapi", "fastapi", "gradio (>=4.0.0)", "httpx", "itsdangerous", "jedi", "libcst (>=1.4.0)", "mypy (==1.15.0) ; python_version >= \"3.9\"", "mypy (>=1.14.1,<1.15.0) ; python_version == \"3.8\"", "numpy", "pytest (>=8.1.1,<8.2.2)", "pytest-asyncio", "pytest-cov", "pytest-env", "pytest-mock", "pytest-rerunfailures (<16.0)", "pytest-vcr", "pytest-xdist", "ruff (>=0.9.0)", "soundfile", "ty", "types-PyYAML", "types-requests", "types-simplejson", "types-toml", "types-tqdm", "types-urllib3", "typing-extensions (>=4.8.0)", "urllib3 (<2.0)"] -cli = ["InquirerPy (==0.3.4)"] -dev = ["InquirerPy (==0.3.4)", "Jinja2", "Pillow", "aiohttp", "authlib (>=1.3.2)", "fastapi", "fastapi", "gradio (>=4.0.0)", "httpx", "itsdangerous", "jedi", "libcst (>=1.4.0)", "mypy (==1.15.0) ; python_version >= \"3.9\"", "mypy (>=1.14.1,<1.15.0) ; python_version == \"3.8\"", "numpy", "pytest (>=8.1.1,<8.2.2)", "pytest-asyncio", "pytest-cov", "pytest-env", "pytest-mock", "pytest-rerunfailures (<16.0)", "pytest-vcr", "pytest-xdist", "ruff (>=0.9.0)", "soundfile", "ty", "types-PyYAML", "types-requests", "types-simplejson", "types-toml", "types-tqdm", "types-urllib3", "typing-extensions (>=4.8.0)", "urllib3 (<2.0)"] +all = ["Jinja2", "Pillow", "authlib (>=1.3.2)", "duckdb", "fastapi", "fastapi", "httpx", "itsdangerous", "jedi", "libcst (>=1.4.0)", "mypy (==1.15.0)", "numpy", "pytest (>=8.4.2)", "pytest-asyncio", "pytest-cov", "pytest-env", "pytest-mock", "pytest-rerunfailures (>=16.2)", "pytest-vcr", "pytest-xdist", "ruff (>=0.9.0)", "soundfile", "ty", "types-PyYAML", "types-simplejson", "types-toml", "types-tqdm", "types-urllib3", "typing-extensions (>=4.8.0)", "urllib3 (<2.0)"] +dev = ["Jinja2", "Pillow", "authlib (>=1.3.2)", "duckdb", "fastapi", "fastapi", "httpx", "itsdangerous", "jedi", "libcst (>=1.4.0)", "mypy (==1.15.0)", "numpy", "pytest (>=8.4.2)", "pytest-asyncio", "pytest-cov", "pytest-env", "pytest-mock", "pytest-rerunfailures (>=16.2)", "pytest-vcr", "pytest-xdist", "ruff (>=0.9.0)", "soundfile", "ty", "types-PyYAML", "types-simplejson", "types-toml", "types-tqdm", "types-urllib3", "typing-extensions (>=4.8.0)", "urllib3 (<2.0)"] fastai = ["fastai (>=2.4)", "fastcore (>=1.3.27)", "toml"] -hf-transfer = ["hf_transfer (>=0.1.4)"] -hf-xet = ["hf-xet (>=1.1.2,<2.0.0)"] -inference = ["aiohttp"] -mcp = ["aiohttp", "mcp (>=1.8.0)", "typer"] +gradio = ["gradio (>=5.0.0)", "requests"] +hf-xet = ["hf-xet (>=1.5.2,<2.0.0)"] +mcp = ["mcp (>=1.8.0)"] oauth = ["authlib (>=1.3.2)", "fastapi", "httpx", "itsdangerous"] -quality = ["libcst (>=1.4.0)", "mypy (==1.15.0) ; python_version >= \"3.9\"", "mypy (>=1.14.1,<1.15.0) ; python_version == \"3.8\"", "ruff (>=0.9.0)", "ty"] -tensorflow = ["graphviz", "pydot", "tensorflow"] -tensorflow-testing = ["keras (<3.0)", "tensorflow"] -testing = ["InquirerPy (==0.3.4)", "Jinja2", "Pillow", "aiohttp", "authlib (>=1.3.2)", "fastapi", "fastapi", "gradio (>=4.0.0)", "httpx", "itsdangerous", "jedi", "numpy", "pytest (>=8.1.1,<8.2.2)", "pytest-asyncio", "pytest-cov", "pytest-env", "pytest-mock", "pytest-rerunfailures (<16.0)", "pytest-vcr", "pytest-xdist", "soundfile", "urllib3 (<2.0)"] +quality = ["libcst (>=1.4.0)", "mypy (==1.15.0)", "ruff (>=0.9.0)", "ty"] +testing = ["Jinja2", "Pillow", "authlib (>=1.3.2)", "duckdb", "fastapi", "fastapi", "httpx", "itsdangerous", "jedi", "numpy", "pytest (>=8.4.2)", "pytest-asyncio", "pytest-cov", "pytest-env", "pytest-mock", "pytest-rerunfailures (>=16.2)", "pytest-vcr", "pytest-xdist", "soundfile", "urllib3 (<2.0)"] torch = ["safetensors[torch]", "torch"] -typing = ["types-PyYAML", "types-requests", "types-simplejson", "types-toml", "types-tqdm", "types-urllib3", "typing-extensions (>=4.8.0)"] +typing = ["types-PyYAML", "types-simplejson", "types-toml", "types-tqdm", "types-urllib3", "typing-extensions (>=4.8.0)"] [[package]] name = "idna" @@ -2231,6 +2242,31 @@ interegular = ["interegular (>=0.3.1,<0.4.0)"] nearley = ["js2py"] regex = ["regex"] +[[package]] +name = "markdown-it-py" +version = "4.2.0" +description = "Python port of markdown-it. Markdown parsing, done right!" +optional = true +python-versions = ">=3.10" +groups = ["main"] +markers = "extra == \"huggingface\"" +files = [ + {file = "markdown_it_py-4.2.0-py3-none-any.whl", hash = "sha256:9f7ebbcd14fe59494226453aed97c1070d83f8d24b6fc3a3bcf9a38092641c4a"}, + {file = "markdown_it_py-4.2.0.tar.gz", hash = "sha256:04a21681d6fbb623de53f6f364d352309d4094dd4194040a10fd51833e418d49"}, +] + +[package.dependencies] +mdurl = ">=0.1,<1.0" + +[package.extras] +benchmarking = ["psutil", "pytest", "pytest-benchmark"] +compare = ["commonmark (>=0.9,<1.0)", "markdown (>=3.4,<4.0)", "markdown-it-pyrs", "mistletoe (>=1.0,<2.0)", "mistune (>=3.0,<4.0)", "panflute (>=2.3,<3.0)"] +linkify = ["linkify-it-py (>=1,<3)"] +plugins = ["mdit-py-plugins (>=0.5.0)"] +profiling = ["gprof2dot"] +rtd = ["ipykernel", "jupyter_sphinx", "mdit-py-plugins (>=0.5.0)", "myst-parser", "pyyaml", "sphinx", "sphinx-book-theme (>=1.0,<2.0)", "sphinx-copybutton", "sphinx-design"] +testing = ["coverage", "pytest", "pytest-cov", "pytest-regressions", "pytest-timeout", "requests"] + [[package]] name = "markupsafe" version = "3.0.3" @@ -2361,6 +2397,19 @@ files = [ {file = "mccabe-0.7.0.tar.gz", hash = "sha256:348e0240c33b60bbdf4e523192ef919f28cb2c3d7d5c7794f74009290f236325"}, ] +[[package]] +name = "mdurl" +version = "0.1.2" +description = "Markdown URL utilities" +optional = true +python-versions = ">=3.7" +groups = ["main"] +markers = "extra == \"huggingface\"" +files = [ + {file = "mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8"}, + {file = "mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba"}, +] + [[package]] name = "mistune" version = "3.3.4" @@ -4132,11 +4181,12 @@ version = "2.21.0" description = "Pygments is a syntax highlighting package written in Python." optional = false python-versions = ">=3.9" -groups = ["dev"] +groups = ["main", "dev"] files = [ {file = "pygments-2.21.0-py3-none-any.whl", hash = "sha256:2363c69b61c4a97c838da3b130dcd6468f4848992b21a82f2a63ec34377137d9"}, {file = "pygments-2.21.0.tar.gz", hash = "sha256:610ca751c9bc2492b38eb9a38a7fbc93edbbb2d7182edaf34e66ae493dee5c8c"}, ] +markers = {main = "extra == \"huggingface\""} [package.extras] windows-terminal = ["colorama (>=0.4.6)"] @@ -4638,12 +4688,12 @@ version = "2.34.2" description = "Python HTTP for Humans." optional = false python-versions = ">=3.10" -groups = ["main", "dev", "gpu"] +groups = ["main", "dev"] files = [ {file = "requests-2.34.2-py3-none-any.whl", hash = "sha256:2a0d60c172f83ac6ab31e4554906c0f3b3588d37b5cb939b1c061f4907e278e0"}, {file = "requests-2.34.2.tar.gz", hash = "sha256:f288924cae4e29463698d6d60bc6a4da69c89185ad1e0bcc4104f584e960b9ed"}, ] -markers = {main = "extra == \"huggingface\" or extra == \"jobs\"", gpu = "sys_platform == \"win32\" or sys_platform == \"linux\""} +markers = {main = "extra == \"huggingface\" or extra == \"jobs\""} [package.dependencies] certifi = ">=2023.5.7" @@ -4700,6 +4750,26 @@ lark = ">=1.2.2" [package.extras] testing = ["pytest (>=8.3.5)"] +[[package]] +name = "rich" +version = "15.0.0" +description = "Render rich text, tables, progress bars, syntax highlighting, markdown and more to the terminal" +optional = true +python-versions = ">=3.9.0" +groups = ["main"] +markers = "extra == \"huggingface\"" +files = [ + {file = "rich-15.0.0-py3-none-any.whl", hash = "sha256:33bd4ef74232fb73fe9279a257718407f169c09b78a87ad3d296f548e27de0bb"}, + {file = "rich-15.0.0.tar.gz", hash = "sha256:edd07a4824c6b40189fb7ac9bc4c52536e9780fbbfbddf6f1e2502c31b068c36"}, +] + +[package.dependencies] +markdown-it-py = ">=2.2.0" +pygments = ">=2.13.0,<3.0.0" + +[package.extras] +jupyter = ["ipywidgets (>=7.5.1,<9)"] + [[package]] name = "rpds-py" version = "0.30.0" @@ -5144,6 +5214,19 @@ enabler = ["pytest-enabler (>=3.4)"] test = ["build[virtualenv] (>=1.0.3)", "filelock (>=3.4.0)", "ini2toml[lite] (>=0.14)", "jaraco.develop (>=7.21) ; python_version >= \"3.9\" and sys_platform != \"cygwin\"", "jaraco.envs (>=2.2)", "jaraco.path (>=3.7.2)", "jaraco.test (>=5.5)", "packaging (>=24.2)", "pip (>=19.1)", "pyproject-hooks (!=1.1)", "pytest (>=6,!=8.1.*)", "pytest-home (>=0.5)", "pytest-perf ; sys_platform != \"cygwin\"", "pytest-subprocess", "pytest-timeout", "pytest-xdist (>=3)", "tomli-w (>=1.0.0)", "virtualenv (>=13.0.0)", "wheel (>=0.44.0)"] type = ["importlib_metadata (>=7.0.2) ; python_version < \"3.10\"", "jaraco.develop (>=7.21) ; sys_platform != \"cygwin\"", "mypy (==1.18.*)", "pytest-mypy (>=1.0.1) ; platform_python_implementation != \"PyPy\""] +[[package]] +name = "shellingham" +version = "1.5.4" +description = "Tool to Detect Surrounding Shell" +optional = true +python-versions = ">=3.7" +groups = ["main"] +markers = "extra == \"huggingface\"" +files = [ + {file = "shellingham-1.5.4-py2.py3-none-any.whl", hash = "sha256:7ecfff8f2fd72616f7481040475a65b2bf8af90a56c89140852d1120324e8686"}, + {file = "shellingham-1.5.4.tar.gz", hash = "sha256:8dbca0739d487e5bd35ab3ca4b36e11c4078f3a234bfce294b0a0291363404de"}, +] + [[package]] name = "sil-thot" version = "3.5.2" @@ -5307,37 +5390,46 @@ test = ["pytest", "ruff"] [[package]] name = "tokenizers" -version = "0.21.4" +version = "0.22.2" description = "" optional = true python-versions = ">=3.9" groups = ["main"] markers = "extra == \"huggingface\"" files = [ - {file = "tokenizers-0.21.4-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:2ccc10a7c3bcefe0f242867dc914fc1226ee44321eb618cfe3019b5df3400133"}, - {file = "tokenizers-0.21.4-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:5e2f601a8e0cd5be5cc7506b20a79112370b9b3e9cb5f13f68ab11acd6ca7d60"}, - {file = "tokenizers-0.21.4-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:39b376f5a1aee67b4d29032ee85511bbd1b99007ec735f7f35c8a2eb104eade5"}, - {file = "tokenizers-0.21.4-cp39-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:2107ad649e2cda4488d41dfd031469e9da3fcbfd6183e74e4958fa729ffbf9c6"}, - {file = "tokenizers-0.21.4-cp39-abi3-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:3c73012da95afafdf235ba80047699df4384fdc481527448a078ffd00e45a7d9"}, - {file = "tokenizers-0.21.4-cp39-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:f23186c40395fc390d27f519679a58023f368a0aad234af145e0f39ad1212732"}, - {file = "tokenizers-0.21.4-cp39-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:cc88bb34e23a54cc42713d6d98af5f1bf79c07653d24fe984d2d695ba2c922a2"}, - {file = "tokenizers-0.21.4-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:51b7eabb104f46c1c50b486520555715457ae833d5aee9ff6ae853d1130506ff"}, - {file = "tokenizers-0.21.4-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:714b05b2e1af1288bd1bc56ce496c4cebb64a20d158ee802887757791191e6e2"}, - {file = "tokenizers-0.21.4-cp39-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:1340ff877ceedfa937544b7d79f5b7becf33a4cfb58f89b3b49927004ef66f78"}, - {file = "tokenizers-0.21.4-cp39-abi3-musllinux_1_2_i686.whl", hash = "sha256:3c1f4317576e465ac9ef0d165b247825a2a4078bcd01cba6b54b867bdf9fdd8b"}, - {file = "tokenizers-0.21.4-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:c212aa4e45ec0bb5274b16b6f31dd3f1c41944025c2358faaa5782c754e84c24"}, - {file = "tokenizers-0.21.4-cp39-abi3-win32.whl", hash = "sha256:6c42a930bc5f4c47f4ea775c91de47d27910881902b0f20e4990ebe045a415d0"}, - {file = "tokenizers-0.21.4-cp39-abi3-win_amd64.whl", hash = "sha256:475d807a5c3eb72c59ad9b5fcdb254f6e17f53dfcbb9903233b0dfa9c943b597"}, - {file = "tokenizers-0.21.4.tar.gz", hash = "sha256:fa23f85fbc9a02ec5c6978da172cdcbac23498c3ca9f3645c5c68740ac007880"}, + {file = "tokenizers-0.22.2-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:544dd704ae7238755d790de45ba8da072e9af3eea688f698b137915ae959281c"}, + {file = "tokenizers-0.22.2-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:1e418a55456beedca4621dbab65a318981467a2b188e982a23e117f115ce5001"}, + {file = "tokenizers-0.22.2-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2249487018adec45d6e3554c71d46eb39fa8ea67156c640f7513eb26f318cec7"}, + {file = "tokenizers-0.22.2-cp39-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:25b85325d0815e86e0bac263506dd114578953b7b53d7de09a6485e4a160a7dd"}, + {file = "tokenizers-0.22.2-cp39-abi3-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:bfb88f22a209ff7b40a576d5324bf8286b519d7358663db21d6246fb17eea2d5"}, + {file = "tokenizers-0.22.2-cp39-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:1c774b1276f71e1ef716e5486f21e76333464f47bece56bbd554485982a9e03e"}, + {file = "tokenizers-0.22.2-cp39-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:df6c4265b289083bf710dff49bc51ef252f9d5be33a45ee2bed151114a56207b"}, + {file = "tokenizers-0.22.2-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:369cc9fc8cc10cb24143873a0d95438bb8ee257bb80c71989e3ee290e8d72c67"}, + {file = "tokenizers-0.22.2-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:29c30b83d8dcd061078b05ae0cb94d3c710555fbb44861139f9f83dcca3dc3e4"}, + {file = "tokenizers-0.22.2-cp39-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:37ae80a28c1d3265bb1f22464c856bd23c02a05bb211e56d0c5301a435be6c1a"}, + {file = "tokenizers-0.22.2-cp39-abi3-musllinux_1_2_i686.whl", hash = "sha256:791135ee325f2336f498590eb2f11dc5c295232f288e75c99a36c5dbce63088a"}, + {file = "tokenizers-0.22.2-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:38337540fbbddff8e999d59970f3c6f35a82de10053206a7562f1ea02d046fa5"}, + {file = "tokenizers-0.22.2-cp39-abi3-win32.whl", hash = "sha256:a6bf3f88c554a2b653af81f3204491c818ae2ac6fbc09e76ef4773351292bc92"}, + {file = "tokenizers-0.22.2-cp39-abi3-win_amd64.whl", hash = "sha256:c9ea31edff2968b44a88f97d784c2f16dc0729b8b143ed004699ebca91f05c48"}, + {file = "tokenizers-0.22.2-cp39-abi3-win_arm64.whl", hash = "sha256:9ce725d22864a1e965217204946f830c37876eee3b2ba6fc6255e8e903d5fcbc"}, + {file = "tokenizers-0.22.2-pp310-pypy310_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:753d47ebd4542742ef9261d9da92cd545b2cacbb48349a1225466745bb866ec4"}, + {file = "tokenizers-0.22.2-pp310-pypy310_pp73-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:e10bf9113d209be7cd046d40fbabbaf3278ff6d18eb4da4c500443185dc1896c"}, + {file = "tokenizers-0.22.2-pp310-pypy310_pp73-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:64d94e84f6660764e64e7e0b22baa72f6cd942279fdbb21d46abd70d179f0195"}, + {file = "tokenizers-0.22.2-pp310-pypy310_pp73-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:f01a9c019878532f98927d2bacb79bbb404b43d3437455522a00a30718cdedb5"}, + {file = "tokenizers-0.22.2-pp39-pypy39_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:319f659ee992222f04e58f84cbf407cfa66a65fe3a8de44e8ad2bc53e7d99012"}, + {file = "tokenizers-0.22.2-pp39-pypy39_pp73-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:1e50f8554d504f617d9e9d6e4c2c2884a12b388a97c5c77f0bc6cf4cd032feee"}, + {file = "tokenizers-0.22.2-pp39-pypy39_pp73-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:1a62ba2c5faa2dd175aaeed7b15abf18d20266189fb3406c5d0550dd34dd5f37"}, + {file = "tokenizers-0.22.2-pp39-pypy39_pp73-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:143b999bdc46d10febb15cbffb4207ddd1f410e2c755857b5a0797961bbdc113"}, + {file = "tokenizers-0.22.2.tar.gz", hash = "sha256:473b83b915e547aa366d1eee11806deaf419e17be16310ac0a14077f1e28f917"}, ] [package.dependencies] -huggingface-hub = ">=0.16.4,<1.0" +huggingface-hub = ">=0.16.4,<2.0" [package.extras] dev = ["tokenizers[testing]"] docs = ["setuptools-rust", "sphinx", "sphinx-rtd-theme"] -testing = ["black (==22.3)", "datasets", "numpy", "pytest", "requests", "ruff"] +testing = ["datasets", "numpy", "pytest", "pytest-asyncio", "requests", "ruff", "ty"] [[package]] name = "tomli" @@ -5526,74 +5618,58 @@ test = ["argcomplete (>=3.0.3) ; python_version < \"3.12\"", "argcomplete (>=3.5 [[package]] name = "transformers" -version = "4.47.1" -description = "State-of-the-art Machine Learning for JAX, PyTorch and TensorFlow" +version = "5.14.1" +description = "Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training." optional = true -python-versions = ">=3.9.0" +python-versions = ">=3.10.0" groups = ["main"] markers = "extra == \"huggingface\"" files = [ - {file = "transformers-4.47.1-py3-none-any.whl", hash = "sha256:d2f5d19bb6283cd66c893ec7e6d931d6370bbf1cc93633326ff1f41a40046c9c"}, - {file = "transformers-4.47.1.tar.gz", hash = "sha256:6c29c05a5f595e278481166539202bf8641281536df1c42357ee58a45d0a564a"}, + {file = "transformers-5.14.1-py3-none-any.whl", hash = "sha256:9db974c4079ede2d1a3ea7ca5a240df33f2cc26fc2b36ba64c5f2a4f43b6e725"}, + {file = "transformers-5.14.1.tar.gz", hash = "sha256:60d196c27781eacf8637e2b533f517582907ad6f9ae142046d6b69431a5b2173"}, ] [package.dependencies] -filelock = "*" -huggingface-hub = ">=0.24.0,<1.0" +huggingface-hub = ">=1.5.0,<2.0" numpy = ">=1.17" packaging = ">=20.0" pyyaml = ">=5.1" -regex = "!=2019.12.17" -requests = "*" -safetensors = ">=0.4.1" -tokenizers = ">=0.21,<0.22" -tqdm = ">=4.27" +regex = ">=2025.10.22" +safetensors = ">=0.8.0" +tokenizers = ">=0.22.0,<=0.23.0" +tqdm = ">=4.60" +typer = "*" [package.extras] -accelerate = ["accelerate (>=0.26.0)"] -agents = ["Pillow (>=10.0.1,<=15.0)", "accelerate (>=0.26.0)", "datasets (!=2.5.0)", "diffusers", "opencv-python", "sentencepiece (>=0.1.91,!=0.1.92)", "torch"] -all = ["Pillow (>=10.0.1,<=15.0)", "accelerate (>=0.26.0)", "av (==9.2.0)", "codecarbon (==1.2.0)", "flax (>=0.4.1,<=0.7.0)", "jax (>=0.4.1,<=0.4.13)", "jaxlib (>=0.4.1,<=0.4.13)", "kenlm", "keras-nlp (>=0.3.1,<0.14.0)", "librosa", "onnxconverter-common", "optax (>=0.0.8,<=0.1.4)", "optuna", "phonemizer", "protobuf", "pyctcdecode (>=0.4.0)", "ray[tune] (>=2.7.0)", "scipy (<1.13.0)", "sentencepiece (>=0.1.91,!=0.1.92)", "sigopt", "tensorflow (>2.9,<2.16)", "tensorflow-text (<2.16)", "tf2onnx", "timm (<=1.0.11)", "tokenizers (>=0.21,<0.22)", "torch", "torchaudio", "torchvision"] -audio = ["kenlm", "librosa", "phonemizer", "pyctcdecode (>=0.4.0)"] +accelerate = ["accelerate (>=1.1.0)"] +all = ["Pillow (>=10.0.1,<=15.0)", "accelerate (>=1.1.0)", "av", "blobfile", "jinja2 (>=3.1.0)", "kernels (>=0.15.2,<0.16)", "librosa", "mistral-common[image] (>=1.11.5)", "num2words", "phonemizer", "protobuf", "pyctcdecode (>=0.4.0)", "sentencepiece (>=0.1.91,!=0.1.92)", "tiktoken", "timm (>=1.0.23)", "torch (>=2.4)", "torchaudio", "torchvision"] +audio = ["librosa", "phonemizer", "pyctcdecode (>=0.4.0)", "torchaudio"] benchmark = ["optimum-benchmark (>=0.3.0)"] -codecarbon = ["codecarbon (==1.2.0)"] -deepspeed = ["accelerate (>=0.26.0)", "deepspeed (>=0.9.3)"] -deepspeed-testing = ["GitPython (<3.1.19)", "accelerate (>=0.26.0)", "beautifulsoup4", "cookiecutter (==1.7.3)", "datasets (!=2.5.0)", "deepspeed (>=0.9.3)", "dill (<0.3.5)", "evaluate (>=0.2.0)", "faiss-cpu", "nltk (<=3.8.1)", "optuna", "parameterized", "protobuf", "psutil", "pydantic", "pytest (>=7.2.0,<8.0.0)", "pytest-rich", "pytest-timeout", "pytest-xdist", "rjieba", "rouge-score (!=0.0.7,!=0.0.8,!=0.1,!=0.1.1)", "ruff (==0.5.1)", "sacrebleu (>=1.4.12,<2.0.0)", "sacremoses", "sentencepiece (>=0.1.91,!=0.1.92)", "tensorboard", "timeout-decorator"] -dev = ["GitPython (<3.1.19)", "Pillow (>=10.0.1,<=15.0)", "accelerate (>=0.26.0)", "av (==9.2.0)", "beautifulsoup4", "codecarbon (==1.2.0)", "cookiecutter (==1.7.3)", "datasets (!=2.5.0)", "dill (<0.3.5)", "evaluate (>=0.2.0)", "faiss-cpu", "flax (>=0.4.1,<=0.7.0)", "fugashi (>=1.0)", "ipadic (>=1.0.0,<2.0)", "isort (>=5.5.4)", "jax (>=0.4.1,<=0.4.13)", "jaxlib (>=0.4.1,<=0.4.13)", "kenlm", "keras-nlp (>=0.3.1,<0.14.0)", "libcst", "librosa", "nltk (<=3.8.1)", "onnxconverter-common", "optax (>=0.0.8,<=0.1.4)", "optuna", "parameterized", "phonemizer", "protobuf", "psutil", "pyctcdecode (>=0.4.0)", "pydantic", "pytest (>=7.2.0,<8.0.0)", "pytest-rich", "pytest-timeout", "pytest-xdist", "ray[tune] (>=2.7.0)", "rhoknp (>=1.1.0,<1.3.1)", "rich", "rjieba", "rouge-score (!=0.0.7,!=0.0.8,!=0.1,!=0.1.1)", "ruff (==0.5.1)", "sacrebleu (>=1.4.12,<2.0.0)", "sacremoses", "scikit-learn", "scipy (<1.13.0)", "sentencepiece (>=0.1.91,!=0.1.92)", "sigopt", "sudachidict-core (>=20220729)", "sudachipy (>=0.6.6)", "tensorboard", "tensorflow (>2.9,<2.16)", "tensorflow-text (<2.16)", "tf2onnx", "timeout-decorator", "timm (<=1.0.11)", "tokenizers (>=0.21,<0.22)", "torch", "torchaudio", "torchvision", "unidic (>=1.0.2)", "unidic-lite (>=1.0.7)", "urllib3 (<2.0.0)"] -dev-tensorflow = ["GitPython (<3.1.19)", "Pillow (>=10.0.1,<=15.0)", "beautifulsoup4", "cookiecutter (==1.7.3)", "datasets (!=2.5.0)", "dill (<0.3.5)", "evaluate (>=0.2.0)", "faiss-cpu", "isort (>=5.5.4)", "kenlm", "keras-nlp (>=0.3.1,<0.14.0)", "libcst", "librosa", "nltk (<=3.8.1)", "onnxconverter-common", "onnxruntime (>=1.4.0)", "onnxruntime-tools (>=1.4.2)", "parameterized", "phonemizer", "protobuf", "psutil", "pyctcdecode (>=0.4.0)", "pydantic", "pytest (>=7.2.0,<8.0.0)", "pytest-rich", "pytest-timeout", "pytest-xdist", "rich", "rjieba", "rouge-score (!=0.0.7,!=0.0.8,!=0.1,!=0.1.1)", "ruff (==0.5.1)", "sacrebleu (>=1.4.12,<2.0.0)", "sacremoses", "scikit-learn", "sentencepiece (>=0.1.91,!=0.1.92)", "tensorboard", "tensorflow (>2.9,<2.16)", "tensorflow-text (<2.16)", "tf2onnx", "timeout-decorator", "tokenizers (>=0.21,<0.22)", "urllib3 (<2.0.0)"] -dev-torch = ["GitPython (<3.1.19)", "Pillow (>=10.0.1,<=15.0)", "accelerate (>=0.26.0)", "beautifulsoup4", "codecarbon (==1.2.0)", "cookiecutter (==1.7.3)", "datasets (!=2.5.0)", "dill (<0.3.5)", "evaluate (>=0.2.0)", "faiss-cpu", "fugashi (>=1.0)", "ipadic (>=1.0.0,<2.0)", "isort (>=5.5.4)", "kenlm", "libcst", "librosa", "nltk (<=3.8.1)", "onnxruntime (>=1.4.0)", "onnxruntime-tools (>=1.4.2)", "optuna", "parameterized", "phonemizer", "protobuf", "psutil", "pyctcdecode (>=0.4.0)", "pydantic", "pytest (>=7.2.0,<8.0.0)", "pytest-rich", "pytest-timeout", "pytest-xdist", "ray[tune] (>=2.7.0)", "rhoknp (>=1.1.0,<1.3.1)", "rich", "rjieba", "rouge-score (!=0.0.7,!=0.0.8,!=0.1,!=0.1.1)", "ruff (==0.5.1)", "sacrebleu (>=1.4.12,<2.0.0)", "sacremoses", "scikit-learn", "sentencepiece (>=0.1.91,!=0.1.92)", "sigopt", "sudachidict-core (>=20220729)", "sudachipy (>=0.6.6)", "tensorboard", "timeout-decorator", "timm (<=1.0.11)", "tokenizers (>=0.21,<0.22)", "torch", "torchaudio", "torchvision", "unidic (>=1.0.2)", "unidic-lite (>=1.0.7)", "urllib3 (<2.0.0)"] -flax = ["flax (>=0.4.1,<=0.7.0)", "jax (>=0.4.1,<=0.4.13)", "jaxlib (>=0.4.1,<=0.4.13)", "optax (>=0.0.8,<=0.1.4)", "scipy (<1.13.0)"] -flax-speech = ["kenlm", "librosa", "phonemizer", "pyctcdecode (>=0.4.0)"] -ftfy = ["ftfy"] -integrations = ["optuna", "ray[tune] (>=2.7.0)", "sigopt"] -ja = ["fugashi (>=1.0)", "ipadic (>=1.0.0,<2.0)", "rhoknp (>=1.1.0,<1.3.1)", "sudachidict-core (>=20220729)", "sudachipy (>=0.6.6)", "unidic (>=1.0.2)", "unidic-lite (>=1.0.7)"] -modelcreation = ["cookiecutter (==1.7.3)"] -natten = ["natten (>=0.14.6,<0.15.0)"] -onnx = ["onnxconverter-common", "onnxruntime (>=1.4.0)", "onnxruntime-tools (>=1.4.2)", "tf2onnx"] -onnxruntime = ["onnxruntime (>=1.4.0)", "onnxruntime-tools (>=1.4.2)"] +chat-template = ["jinja2 (>=3.1.0)"] +codecarbon = ["codecarbon (>=2.8.1)"] +deepspeed = ["accelerate (>=1.1.0)", "deepspeed (>=0.9.3)"] +deepspeed-testing = ["GitPython (<3.1.19)", "accelerate (>=1.1.0)", "accelerate (>=1.1.0)", "beautifulsoup4", "datasets (>=2.15.0)", "datasets (>=2.15.0)", "deepspeed (>=0.9.3)", "dill (<0.3.5)", "evaluate (>=0.4.6)", "faiss-cpu", "fastapi", "filelock", "hf-doc-builder", "libcst", "mistral-common[image] (>=1.11.5)", "nltk (<=3.8.1)", "openai (>=1.98.0)", "optuna", "parameterized (>=0.9)", "protobuf", "protobuf", "psutil", "pydantic (>=2)", "pytest (>=7.2.0,<9.0.0)", "pytest-asyncio (>=1.2.0)", "pytest-env", "pytest-order", "pytest-random-order", "pytest-rerunfailures (<16.0)", "pytest-rich", "pytest-timeout", "pytest-xdist", "rich", "rich", "rjieba", "rouge-score (!=0.0.7,!=0.0.8,!=0.1,!=0.1.1)", "ruff (==0.14.10)", "sacrebleu (>=1.4.12,<2.0.0)", "sacremoses", "sentencepiece (>=0.1.91,!=0.1.92)", "sentencepiece (>=0.1.91,!=0.1.92)", "starlette", "tensorboard", "timeout-decorator", "tomli", "torch (>=2.4)", "transformers-mlinter (==0.1.2)", "ty (==0.0.20)", "urllib3 (<2.0.0)", "uvicorn"] +dev = ["GitPython (<3.1.19)", "Pillow (>=10.0.1,<=15.0)", "accelerate (>=1.1.0)", "accelerate (>=1.1.0)", "av", "beautifulsoup4", "blobfile", "datasets (>=2.15.0)", "datasets (>=2.15.0)", "dill (<0.3.5)", "evaluate (>=0.4.6)", "faiss-cpu", "fastapi", "filelock", "fugashi (>=1.0)", "hf-doc-builder", "ipadic (>=1.0.0,<2.0)", "jinja2 (>=3.1.0)", "kernels (>=0.15.2,<0.16)", "libcst", "librosa", "mistral-common[image] (>=1.11.5)", "mistral-common[image] (>=1.11.5)", "nltk (<=3.8.1)", "num2words", "openai (>=1.98.0)", "parameterized (>=0.9)", "phonemizer", "protobuf", "protobuf", "psutil", "pyctcdecode (>=0.4.0)", "pydantic (>=2)", "pytest (>=7.2.0,<9.0.0)", "pytest-asyncio (>=1.2.0)", "pytest-env", "pytest-order", "pytest-random-order", "pytest-rerunfailures (<16.0)", "pytest-rich", "pytest-timeout", "pytest-xdist", "rhoknp (>=1.1.0,<1.3.1)", "rich", "rich", "rjieba", "rouge-score (!=0.0.7,!=0.0.8,!=0.1,!=0.1.1)", "ruff (==0.14.10)", "sacrebleu (>=1.4.12,<2.0.0)", "sacremoses", "scikit-learn", "sentencepiece (>=0.1.91,!=0.1.92)", "sentencepiece (>=0.1.91,!=0.1.92)", "starlette", "sudachidict_core (>=20220729)", "sudachipy (>=0.6.6)", "tensorboard", "tiktoken", "timeout-decorator", "timm (>=1.0.23)", "tomli", "torch (>=2.4)", "torch (>=2.4)", "torchaudio", "torchvision", "transformers-mlinter (==0.1.2)", "ty (==0.0.20)", "unidic (>=1.0.2)", "unidic_lite (>=1.0.7)", "urllib3 (<2.0.0)", "uvicorn"] +docs = ["hf-doc-builder"] +integrations = ["codecarbon (>=2.8.1)", "kernels (>=0.15.2,<0.16)", "optuna", "ray[tune] (>=2.7.0)"] +ja = ["fugashi (>=1.0)", "ipadic (>=1.0.0,<2.0)", "rhoknp (>=1.1.0,<1.3.1)", "sudachidict_core (>=20220729)", "sudachipy (>=0.6.6)", "unidic (>=1.0.2)", "unidic_lite (>=1.0.7)"] +kernels = ["kernels (>=0.15.2,<0.16)"] +mistral-common = ["mistral-common[image] (>=1.11.5)"] +num2words = ["num2words"] optuna = ["optuna"] -quality = ["GitPython (<3.1.19)", "datasets (!=2.5.0)", "isort (>=5.5.4)", "libcst", "rich", "ruff (==0.5.1)", "urllib3 (<2.0.0)"] +quality = ["GitPython (<3.1.19)", "datasets (>=2.15.0)", "libcst", "rich", "ruff (==0.14.10)", "tomli", "transformers-mlinter (==0.1.2)", "ty (==0.0.20)", "urllib3 (<2.0.0)"] ray = ["ray[tune] (>=2.7.0)"] -retrieval = ["datasets (!=2.5.0)", "faiss-cpu"] -ruff = ["ruff (==0.5.1)"] +retrieval = ["datasets (>=2.15.0)", "faiss-cpu"] sagemaker = ["sagemaker (>=2.31.0)"] sentencepiece = ["protobuf", "sentencepiece (>=0.1.91,!=0.1.92)"] -serving = ["fastapi", "pydantic", "starlette", "uvicorn"] -sigopt = ["sigopt"] +serving = ["accelerate (>=1.1.0)", "fastapi", "openai (>=1.98.0)", "pydantic (>=2)", "rich", "starlette", "torch (>=2.4)", "uvicorn"] sklearn = ["scikit-learn"] -speech = ["kenlm", "librosa", "phonemizer", "pyctcdecode (>=0.4.0)", "torchaudio"] -testing = ["GitPython (<3.1.19)", "beautifulsoup4", "cookiecutter (==1.7.3)", "datasets (!=2.5.0)", "dill (<0.3.5)", "evaluate (>=0.2.0)", "faiss-cpu", "nltk (<=3.8.1)", "parameterized", "psutil", "pydantic", "pytest (>=7.2.0,<8.0.0)", "pytest-rich", "pytest-timeout", "pytest-xdist", "rjieba", "rouge-score (!=0.0.7,!=0.0.8,!=0.1,!=0.1.1)", "ruff (==0.5.1)", "sacrebleu (>=1.4.12,<2.0.0)", "sacremoses", "sentencepiece (>=0.1.91,!=0.1.92)", "tensorboard", "timeout-decorator"] -tf = ["keras-nlp (>=0.3.1,<0.14.0)", "onnxconverter-common", "tensorflow (>2.9,<2.16)", "tensorflow-text (<2.16)", "tf2onnx"] -tf-cpu = ["keras (>2.9,<2.16)", "keras-nlp (>=0.3.1,<0.14.0)", "onnxconverter-common", "tensorflow-cpu (>2.9,<2.16)", "tensorflow-probability (<0.24)", "tensorflow-text (<2.16)", "tf2onnx"] -tf-speech = ["kenlm", "librosa", "phonemizer", "pyctcdecode (>=0.4.0)"] +testing = ["GitPython (<3.1.19)", "accelerate (>=1.1.0)", "beautifulsoup4", "datasets (>=2.15.0)", "datasets (>=2.15.0)", "dill (<0.3.5)", "evaluate (>=0.4.6)", "faiss-cpu", "fastapi", "filelock", "hf-doc-builder", "libcst", "mistral-common[image] (>=1.11.5)", "nltk (<=3.8.1)", "openai (>=1.98.0)", "parameterized (>=0.9)", "protobuf", "psutil", "pydantic (>=2)", "pytest (>=7.2.0,<9.0.0)", "pytest-asyncio (>=1.2.0)", "pytest-env", "pytest-order", "pytest-random-order", "pytest-rerunfailures (<16.0)", "pytest-rich", "pytest-timeout", "pytest-xdist", "rich", "rich", "rjieba", "rouge-score (!=0.0.7,!=0.0.8,!=0.1,!=0.1.1)", "ruff (==0.14.10)", "sacrebleu (>=1.4.12,<2.0.0)", "sacremoses", "sentencepiece (>=0.1.91,!=0.1.92)", "starlette", "tensorboard", "timeout-decorator", "tomli", "torch (>=2.4)", "transformers-mlinter (==0.1.2)", "ty (==0.0.20)", "urllib3 (<2.0.0)", "uvicorn"] tiktoken = ["blobfile", "tiktoken"] -timm = ["timm (<=1.0.11)"] -tokenizers = ["tokenizers (>=0.21,<0.22)"] -torch = ["accelerate (>=0.26.0)", "torch"] -torch-speech = ["kenlm", "librosa", "phonemizer", "pyctcdecode (>=0.4.0)", "torchaudio"] -torch-vision = ["Pillow (>=10.0.1,<=15.0)", "torchvision"] -torchhub = ["filelock", "huggingface-hub (>=0.24.0,<1.0)", "importlib-metadata", "numpy (>=1.17)", "packaging (>=20.0)", "protobuf", "regex (!=2019.12.17)", "requests", "sentencepiece (>=0.1.91,!=0.1.92)", "tokenizers (>=0.21,<0.22)", "torch", "tqdm (>=4.27)"] -video = ["av (==9.2.0)"] -vision = ["Pillow (>=10.0.1,<=15.0)"] +timm = ["timm (>=1.0.23)"] +torch = ["accelerate (>=1.1.0)", "torch (>=2.4)"] +video = ["av"] +vision = ["Pillow (>=10.0.1,<=15.0)", "torchvision"] [[package]] name = "triton" @@ -5625,6 +5701,25 @@ build = ["cmake (>=3.20,<4.0)", "lit"] tests = ["autopep8", "isort", "llnl-hatchet", "numpy", "pytest", "pytest-forked", "pytest-xdist", "scipy (>=1.7.1)"] tutorials = ["matplotlib", "pandas", "tabulate"] +[[package]] +name = "typer" +version = "0.27.1" +description = "Typer, build great CLIs. Easy to code. Based on Python type hints." +optional = true +python-versions = ">=3.10" +groups = ["main"] +markers = "extra == \"huggingface\"" +files = [ + {file = "typer-0.27.1-py3-none-any.whl", hash = "sha256:53150287edd11baeb4e4722c8e394fcdf8181c0ae89485cba8d25c778d5edd56"}, + {file = "typer-0.27.1.tar.gz", hash = "sha256:a79bef8469a79c45498e7b814ecf8d603cc7644e9acbd9e19cac0334240b18df"}, +] + +[package.dependencies] +annotated-doc = ">=0.0.2" +colorama = {version = "*", markers = "platform_system == \"Windows\""} +rich = ">=13.8.0" +shellingham = ">=1.3.0" + [[package]] name = "typing-extensions" version = "4.16.0" @@ -5688,12 +5783,11 @@ version = "1.26.20" description = "HTTP library with thread-safe connection pooling, file post, and more." optional = false python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,>=2.7" -groups = ["main", "dev", "gpu"] +groups = ["main", "dev"] files = [ {file = "urllib3-1.26.20-py2.py3-none-any.whl", hash = "sha256:0ed14ccfbf1c30a9072c7ca157e4319b70d65f623e91e7b32fadb2853431016e"}, {file = "urllib3-1.26.20.tar.gz", hash = "sha256:40c2dc0c681e47eb8f90e7e27bf6ff7df2e677421fd46756da1161c39ca70d32"}, ] -markers = {gpu = "sys_platform == \"win32\" or sys_platform == \"linux\""} [package.extras] brotli = ["brotli (==1.0.9) ; os_name != \"nt\" and python_version < \"3\" and platform_python_implementation == \"CPython\"", "brotli (>=1.0.9) ; python_version >= \"3\" and platform_python_implementation == \"CPython\"", "brotlicffi (>=0.8.0) ; (os_name != \"nt\" or python_version >= \"3\") and platform_python_implementation != \"CPython\"", "brotlipy (>=0.6.0) ; os_name == \"nt\" and python_version < \"3\""] @@ -6123,4 +6217,4 @@ thot = ["sil-thot"] [metadata] lock-version = "2.1" python-versions = ">=3.10,<3.15" -content-hash = "4f76eef02d1daa52c7fcd982f992321c41b61dbe4bedae8d9ecd25d7c484a8dc" +content-hash = "08e0ea2d4df0e5b81730dffe129b0c03300575c129f46e304c2fdfc3a6d2aa49" diff --git a/pyproject.toml b/pyproject.toml index 3de57b31..b4a7c721 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -68,12 +68,12 @@ pep8-naming = "^0.14.1" [tool.poetry.group.gpu.dependencies] torch = { version = "2.9.1", markers = "sys_platform == 'win32' or sys_platform == 'linux'" } -accelerate = { version = "^1.0", markers = "sys_platform == 'win32' or sys_platform == 'linux'" } +accelerate = { version = "^1.14", markers = "sys_platform == 'win32' or sys_platform == 'linux'" } [project.optional-dependencies] sentencepiece = ["sentencepiece~=0.2"] thot = ["sil-thot~=3.5.2"] -huggingface = ["transformers==4.47.1", "datasets~=4.1", "sacremoses>=0.0.53,<1"] +huggingface = ["transformers==5.14.1", "datasets==5.0.1", "sacremoses>=0.0.53,<1"] jobs = [ "clearml[s3]>=1.13.1,<2.0", "json-stream~=1.3", "dynaconf>=3.2.5,<4.0", "eflomal~=2.0; sys_platform == 'linux'", "pydantic>=2.13.4,<3.0"] [build-system] diff --git a/tests/translation/huggingface/test_hugging_face_nmt_engine.py b/tests/translation/huggingface/test_hugging_face_nmt_engine.py index 18ed4369..37653244 100644 --- a/tests/translation/huggingface/test_hugging_face_nmt_engine.py +++ b/tests/translation/huggingface/test_hugging_face_nmt_engine.py @@ -19,32 +19,33 @@ def test_translate_n_batch_beam(output_attentions: bool) -> None: "stas/tiny-m2m_100", src_lang="en", tgt_lang="es", - num_beams=2, + num_beams=5, max_length=10, + do_sample=False, output_attentions=output_attentions, ) as engine: results = engine.translate_n_batch( n=2, segments=["This is a test string", "Hello, world!"], ) - assert results[0][0].translation == "skaberskaber Dollar Dollar ፤ ፤ gerekir gerekir" - assert results[0][0].confidences[0] == approx(1.08e-05, 0.01) + assert results[0][0].translation == "skaberskaberskaber Dollar Dollar ፤ ፤ gerekir gerekir" + assert results[0][0].confidences[0] == approx(1.1e-05, 0.01) assert results[0][0].sequence_confidence == approx(_get_sequence_confidence(results[0][0]), 0.01) - assert str(results[0][0].alignment) == ("2-0 2-1 2-2 2-3 4-4 4-5 4-6 4-7" if output_attentions else "") + assert str(results[0][0].alignment) == ("0-3 2-0 2-1 2-2 2-4 4-5 4-6 4-7 4-8" if output_attentions else "") - assert results[0][1].translation == "skaberskaber Dollar Dollar ፤ ፤ ፤ gerekir" - assert results[0][1].confidences[0] == approx(1.08e-05, 0.01) + assert results[0][1].translation == "skaberskaber Dollar Dollar Dollar ፤ ፤ gerekir gerekir" + assert results[0][1].confidences[0] == approx(1.1e-05, 0.01) assert results[0][1].sequence_confidence == approx(_get_sequence_confidence(results[0][0]), 0.01) - assert str(results[0][1].alignment) == ("2-0 2-1 2-2 2-3 4-4 4-5 4-6 4-7" if output_attentions else "") + assert str(results[0][1].alignment) == ("0-3 2-0 2-1 2-2 2-4 4-5 4-6 4-7 4-8" if output_attentions else "") - assert results[1][0].translation == "skaberskaber Dollar Dollar ፤ ፤ gerekir gerekir" - assert results[1][0].confidences[0] == approx(1.08e-05, 0.01) - assert str(results[1][0].alignment) == ("0-1 0-2 0-7 1-0 3-3 3-4 3-5 3-6" if output_attentions else "") + assert results[1][0].translation == "skaberskaberskaber Dollar Dollar ፤ ፤ gerekir gerekir" + assert results[1][0].confidences[0] == approx(1.1e-05, 0.01) + assert str(results[1][0].alignment) == ("0-1 0-2 0-3 0-4 0-8 1-0 3-5 3-6 3-7" if output_attentions else "") assert results[1][0].sequence_confidence == approx(_get_sequence_confidence(results[0][0]), 0.01) - assert results[1][1].translation == "skaberskaber Dollar Dollar ፤ ፤ ፤ gerekir" - assert results[1][1].confidences[0] == approx(1.08e-05, 0.01) - assert str(results[1][1].alignment) == ("0-1 0-2 0-7 1-0 3-3 3-4 3-5 3-6" if output_attentions else "") + assert results[1][1].translation == "skaberskaber Dollar Dollar Dollar ፤ ፤ gerekir gerekir" + assert results[1][1].confidences[0] == approx(1.1e-05, 0.01) + assert str(results[1][1].alignment) == ("0-1 0-2 0-3 0-4 0-8 1-0 3-5 3-6 3-7" if output_attentions else "") assert results[1][1].sequence_confidence == approx(_get_sequence_confidence(results[0][0]), 0.01) @@ -54,10 +55,10 @@ def test_translate_greedy(output_attentions: bool) -> None: "stas/tiny-m2m_100", src_lang="en", tgt_lang="es", max_length=10, output_attentions=output_attentions ) as engine: result = engine.translate("This is a test string") - assert result.translation == "skaberskaber Dollar Dollar Dollar ፤ gerekir gerekir" - assert result.confidences[0] == approx(1.08e-05, 0.01) + assert result.translation == "skaberskaberskaber Dollar Dollar Dollar ፤ gerekir gerekir" + assert result.confidences[0] == approx(1.1e-05, 0.01) assert result.sequence_confidence == -1.0 - assert str(result.alignment) == ("2-0 2-1 2-2 2-3 4-4 4-5 4-6 4-7" if output_attentions else "") + assert str(result.alignment) == ("0-3 2-0 2-1 2-2 2-4 4-5 4-6 4-7 4-8" if output_attentions else "") @mark.parametrize("output_attentions", [True, False]) @@ -67,5 +68,4 @@ def test_construct_invalid_lang(output_attentions: bool) -> None: def _get_sequence_confidence(result: TranslationResult) -> float: - # Inject a 0 score for the BOS token - return exp(sum([log(c) for c in result.confidences] + [0]) / (len(result.confidences) + 1)) + return exp(sum([log(c) for c in result.confidences]) / len(result.confidences)) diff --git a/tests/translation/huggingface/test_hugging_face_nmt_model_trainer.py b/tests/translation/huggingface/test_hugging_face_nmt_model_trainer.py index f2243342..ae9899af 100644 --- a/tests/translation/huggingface/test_hugging_face_nmt_model_trainer.py +++ b/tests/translation/huggingface/test_hugging_face_nmt_model_trainer.py @@ -11,11 +11,8 @@ from transformers import ( M2M100Tokenizer, MBart50Tokenizer, - MBart50TokenizerFast, MBartTokenizer, - MBartTokenizerFast, NllbTokenizer, - NllbTokenizerFast, PreTrainedTokenizerFast, Seq2SeqTrainingArguments, ) @@ -162,8 +159,8 @@ def test_update_tokenizer_missing_char() -> None: "Ḻ, ḻ, Ṉ, ॽ, " + "‌ and " + "‍" + " are new characters" ) finetuned_result_nochar_composite = finetuned_engine_nochar.tokenizer.encode("Ḏ is a composite character") - norm_result_nochar1 = finetuned_engine_nochar.tokenizer.backend_tokenizer.normalizer.normalize_str("‌ ") - norm_result_nochar2 = finetuned_engine_nochar.tokenizer.backend_tokenizer.normalizer.normalize_str("‍") + norm_result_nochar1 = finetuned_engine_nochar.tokenizer.encode("‌ ") + norm_result_nochar2 = finetuned_engine_nochar.tokenizer.encode("‍") with HuggingFaceNmtModelTrainer( "hf-internal-testing/tiny-random-nllb", @@ -185,14 +182,14 @@ def test_update_tokenizer_missing_char() -> None: "Ḻ, ḻ, Ṉ, ॽ, " + "‌ and " + "‍" + " are new characters" ) finetuned_result_char_composite = finetuned_engine_char.tokenizer.encode("Ḏ is a composite character") - norm_result_char1 = finetuned_engine_char.tokenizer.backend_tokenizer.normalizer.normalize_str("‌ ") - norm_result_char2 = finetuned_engine_char.tokenizer.backend_tokenizer.normalizer.normalize_str("‍") + norm_result_char1 = finetuned_engine_char.tokenizer.encode("‌ ") + norm_result_char2 = finetuned_engine_char.tokenizer.encode("‍") assert norm_result_nochar1 != norm_result_char1 assert norm_result_nochar2 != norm_result_char2 assert finetuned_result_nochar != finetuned_result_char - assert finetuned_result_nochar_composite != finetuned_result_char_composite + assert finetuned_result_nochar_composite == finetuned_result_char_composite def test_update_tokenizer_missing_char_skip() -> None: @@ -497,18 +494,6 @@ def test_nllb_tokenizer_add_lang_code() -> None: return -def test_nllb_tokenizer_fast_add_lang_code() -> None: - with TemporaryDirectory() as temp_dir: - tokenizer = cast(NllbTokenizerFast, NllbTokenizerFast.from_pretrained("facebook/nllb-200-distilled-600M")) - assert "new_lang" not in tokenizer.added_tokens_encoder - add_lang_code_to_tokenizer(tokenizer, "new_lang") - assert "new_lang" in tokenizer.added_tokens_encoder - tokenizer.save_pretrained(temp_dir) - new_tokenizer = cast(NllbTokenizerFast, NllbTokenizerFast.from_pretrained(temp_dir)) - assert "new_lang" in new_tokenizer.added_tokens_encoder - return - - def test_mbart_tokenizer_add_lang_code() -> None: with TemporaryDirectory() as temp_dir: tokenizer = cast(MBartTokenizer, MBartTokenizer.from_pretrained("hf-internal-testing/tiny-random-nllb")) @@ -521,18 +506,6 @@ def test_mbart_tokenizer_add_lang_code() -> None: return -def test_mbart_tokenizer_fast_add_lang_code() -> None: - with TemporaryDirectory() as temp_dir: - tokenizer = cast(MBartTokenizerFast, MBartTokenizerFast.from_pretrained("hf-internal-testing/tiny-random-nllb")) - assert "nl_NS" not in tokenizer.added_tokens_encoder - add_lang_code_to_tokenizer(tokenizer, "nl_NS") - assert "nl_NS" in tokenizer.added_tokens_encoder - tokenizer.save_pretrained(temp_dir) - new_tokenizer = cast(MBartTokenizerFast, MBartTokenizerFast.from_pretrained(temp_dir)) - assert "nl_NS" in new_tokenizer.added_tokens_encoder - return - - def test_mbart_50_tokenizer_add_lang_code() -> None: with TemporaryDirectory() as temp_dir: tokenizer = cast(MBart50Tokenizer, MBart50Tokenizer.from_pretrained("hf-internal-testing/tiny-random-mbart50")) @@ -545,20 +518,6 @@ def test_mbart_50_tokenizer_add_lang_code() -> None: return -def test_mbart_50_tokenizer_fast_add_lang_code() -> None: - with TemporaryDirectory() as temp_dir: - tokenizer = cast( - MBart50TokenizerFast, MBart50TokenizerFast.from_pretrained("hf-internal-testing/tiny-random-mbart50") - ) - assert "nl_NS" not in tokenizer.added_tokens_encoder - add_lang_code_to_tokenizer(tokenizer, "nl_NS") - assert "nl_NS" in tokenizer.added_tokens_encoder - tokenizer.save_pretrained(temp_dir) - new_tokenizer = cast(MBart50TokenizerFast, MBart50TokenizerFast.from_pretrained(temp_dir)) - assert "nl_NS" in new_tokenizer.added_tokens_encoder - return - - def test_m2m_100_tokenizer_add_lang_code() -> None: with TemporaryDirectory() as temp_dir: tokenizer = cast(M2M100Tokenizer, M2M100Tokenizer.from_pretrained("stas/tiny-m2m_100")) From 513e7098fafc4789f3e5ef4ee3c1a4da43f07893 Mon Sep 17 00:00:00 2001 From: Peter Chapman Date: Thu, 10 Sep 2026 10:55:09 +1200 Subject: [PATCH 02/14] Use forward-ported pipeline --- .gitattributes | 1 + machine/translation/huggingface/__init__.py | 10 +- .../huggingface/hugging_face_nmt_engine.py | 96 ++++------- .../hugging_face_nmt_model_trainer.py | 1 + .../huggingface/transformers_compatibility.py | 151 ++++++++++++++++++ .../test_hugging_face_nmt_engine.py | 46 +++--- 6 files changed, 219 insertions(+), 86 deletions(-) create mode 100644 machine/translation/huggingface/transformers_compatibility.py diff --git a/.gitattributes b/.gitattributes index 4e7955e1..1ca09588 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1 +1,2 @@ *.ipynb eol=lf +*.sh eol=lf diff --git a/machine/translation/huggingface/__init__.py b/machine/translation/huggingface/__init__.py index 3342eafb..13c25f4f 100644 --- a/machine/translation/huggingface/__init__.py +++ b/machine/translation/huggingface/__init__.py @@ -6,8 +6,14 @@ if not is_torch_available(): raise RuntimeError("torch is not installed.") -from .hugging_face_nmt_engine import HuggingFaceNmtEngine +from .hugging_face_nmt_engine import HuggingFaceNmtEngine, SilTranslationPipeline from .hugging_face_nmt_model import HuggingFaceNmtModel from .hugging_face_nmt_model_trainer import HuggingFaceNmtModelTrainer, add_lang_code_to_tokenizer -__all__ = ["add_lang_code_to_tokenizer", "HuggingFaceNmtEngine", "HuggingFaceNmtModel", "HuggingFaceNmtModelTrainer"] +__all__ = [ + "add_lang_code_to_tokenizer", + "HuggingFaceNmtEngine", + "HuggingFaceNmtModel", + "HuggingFaceNmtModelTrainer", + "SilTranslationPipeline", +] diff --git a/machine/translation/huggingface/hugging_face_nmt_engine.py b/machine/translation/huggingface/hugging_face_nmt_engine.py index e087f692..69e33ab3 100644 --- a/machine/translation/huggingface/hugging_face_nmt_engine.py +++ b/machine/translation/huggingface/hugging_face_nmt_engine.py @@ -15,13 +15,12 @@ AutoTokenizer, M2M100Tokenizer, NllbTokenizer, - Pipeline, PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast, ) from transformers.generation.utils import GenerateBeamEncoderDecoderOutput, GenerateEncoderDecoderOutput -from transformers.tokenization_utils_base import TruncationStrategy +from transformers.tokenization_utils_base import BatchEncoding, TruncationStrategy from ...annotations.range import Range from ...corpora.aligned_word_pair import AlignedWordPair @@ -30,6 +29,7 @@ from ..translation_result_builder import TranslationResultBuilder from ..translation_sources import TranslationSources from ..word_alignment_matrix import WordAlignmentMatrix +from .transformers_compatibility import TranslationPipeline logger = logging.getLogger(__name__) @@ -106,7 +106,7 @@ def __init__( self._oom_batch_size_backoff_mult = oom_batch_size_backoff_mult - self._pipeline = _TranslationPipeline( + self._pipeline = SilTranslationPipeline( model=self._model, tokenizer=cast(PreTrainedTokenizer, self.tokenizer), mpn=self._mpn, @@ -146,7 +146,7 @@ def translate_n_batch( raise self._batch_size = max(int(round(self._batch_size * self._oom_batch_size_backoff_mult)), 1) logger.warning(f"Out of memory error caught. Reducing batch size to {self._batch_size} and retrying.") - self._pipeline = _TranslationPipeline( + self._pipeline = SilTranslationPipeline( model=self._model, tokenizer=self._tokenizer, batch_size=self._batch_size, @@ -196,44 +196,17 @@ def close(self) -> None: torch.cuda.empty_cache() -class _TranslationPipeline(Pipeline): +class SilTranslationPipeline(TranslationPipeline): def __init__( self, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, + batch_size: int, mpn: Optional[MosesPunctNormalizer] = None, **kwargs, ) -> None: - if isinstance(model, (str, Path)): - model = cast(PreTrainedModel, AutoModelForSeq2SeqLM.from_pretrained(model)) - self.device: torch.device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu") - self.model: Any = cast(PreTrainedModel, cast(torch.nn.Module, model).to(self.device)) - self.tokenizer = tokenizer + super().__init__(model=model, tokenizer=tokenizer, batch_size=batch_size, **kwargs) self._mpn = mpn - self._kwargs = kwargs - - def __call__(self, sentences: Iterable[Any], **generate_kwargs) -> List[List[dict]]: - model_inputs = self.preprocess(*sentences).to(self.device) - with torch.no_grad(): - model_outputs = self._forward(model_inputs, **generate_kwargs) - return self.postprocess(model_outputs) - - def _sanitize_parameters(self, **pipeline_parameters): - return pipeline_parameters - - def check_inputs(self, input_length: int, min_length: Optional[int], max_length: Optional[int]) -> None: - """ - Checks whether there might be something wrong with given input with regard to the model. - """ - if min_length is not None and max_length is not None and max_length < min_length: - logger.warning(f"Your min_length={min_length} must be inferior than your max_length={max_length}.") - - if max_length is not None and input_length < max_length: - logger.warning( - f"Your max_length is set to {max_length}, but your input_length is only {input_length}. Since this is " - "a summarization task, where outputs shorter than the input are typically wanted, you might " - f"consider decreasing max_length manually, e.g. summarizer('...', max_length={input_length // 2})" - ) def preprocess(self, *args, truncation=TruncationStrategy.DO_NOT_TRUNCATE, src_lang=None, tgt_lang=None): if self.tokenizer is None: @@ -256,17 +229,22 @@ def preprocess(self, *args, truncation=TruncationStrategy.DO_NOT_TRUNCATE, src_l ) for s in args ] - - # The source language prefix and the forced target language token are configured on the tokenizer and the - # generation config in _configure_model, so a plain tokenizer call is all that is needed here. - return self.tokenizer( - cast(list[str] | list[list[str]], sentences), - return_tensors="pt", - truncation=truncation, - padding=True, + inputs = cast( + BatchEncoding, + super().preprocess(*sentences, truncation=truncation, src_lang=src_lang, tgt_lang=tgt_lang), ) + if inputs.encodings is not None: + inputs["input_tokens"] = [ + _get_encoding_fast_tokens(inputs.encodings[i]) if isinstance(args[i], str) else args[i] + for i in range(len(args)) + ] + else: + inputs["input_tokens"] = [self.tokenizer.tokenize(s) if isinstance(s, str) else s for s in args] + return inputs def _forward(self, model_inputs, **generate_kwargs): + if self.tokenizer is None: + raise RuntimeError("No tokenizer is specified.") in_b, input_length = model_inputs["input_ids"].shape if "input_tokens" in model_inputs: @@ -274,24 +252,11 @@ def _forward(self, model_inputs, **generate_kwargs): else: input_tokens = [self.tokenizer.convert_ids_to_tokens(seq) for seq in model_inputs["input_ids"]] - if hasattr(self.model, "generation_config") and self.model.generation_config is not None: - config = self.model.generation_config - else: - config = self.model.config - generate_kwargs["min_length"] = generate_kwargs.get( - "min_length", self._kwargs.get("min_length", config.min_length or 0) - ) - generate_kwargs["max_length"] = generate_kwargs.get( - "max_length", self._kwargs.get("max_length", config.max_length or 200) - ) - generate_kwargs["output_attentions"] = generate_kwargs.get( - "output_attentions", self._kwargs.get("output_attentions", True) - ) - generate_kwargs["num_beams"] = self._kwargs.get("num_beams", config.num_beams or 1) - self.check_inputs(input_length, generate_kwargs["min_length"], generate_kwargs["max_length"]) - output = self.model.generate( + self.check_inputs(input_length, self.generation_config.min_length, self.generation_config.max_length) + output = cast(Any, self.model).generate( **model_inputs, **generate_kwargs, + generation_config=self.generation_config, output_scores=True, return_dict_in_generate=True, ) @@ -316,7 +281,7 @@ def _forward(self, model_inputs, **generate_kwargs): try: transition_scores = cast( torch.Tensor, - self.model.compute_transition_scores( + cast(Any, self.model).compute_transition_scores( output_ids, scores, beam_indices, @@ -329,14 +294,13 @@ def _forward(self, model_inputs, **generate_kwargs): beam_indices = beam_indices.to("cpu") if beam_indices is not None else None transition_scores = cast( torch.Tensor, - self.model.compute_transition_scores( + cast(Any, self.model).compute_transition_scores( output_ids, scores, beam_indices, normalize_logits=True, ), ) - sequences_scores = getattr(output, "sequences_scores", None) if beam_indices is None: beam_indices = torch.zeros_like(output_ids) @@ -365,7 +329,7 @@ def _forward(self, model_inputs, **generate_kwargs): start_index = 0 if self.model.config.decoder_start_token_id is not None: start_index = 1 - if generate_kwargs["output_attentions"] is True: + if self.generation_config.output_attentions: assert attentions is not None num_heads = attentions[0][0].shape[1] @@ -421,7 +385,7 @@ def _forward(self, model_inputs, **generate_kwargs): "attentions": attentions, } - def postprocess(self, model_outputs, clean_up_tokenization_spaces=False) -> List[List[dict]]: + def postprocess(self, model_outputs, clean_up_tokenization_spaces=False): if self.tokenizer is None: raise RuntimeError("No tokenizer is specified.") all_special_ids = set(self.tokenizer.all_special_ids) @@ -478,7 +442,7 @@ def postprocess(self, model_outputs, clean_up_tokenization_spaces=False) -> List records.append(record) - return [records] + return records def torch_gather_nd(params: torch.Tensor, indices: torch.Tensor, batch_dim: int = 0) -> torch.Tensor: @@ -511,3 +475,7 @@ def torch_gather_nd(params: torch.Tensor, indices: torch.Tensor, batch_dim: int out = torch.gather(params, dim=batch_dim, index=indices) return out.reshape(*index_shape, *tail_sizes) + + +def _get_encoding_fast_tokens(encoding) -> List[str]: + return [token for (token, mask) in zip(encoding.tokens, encoding.special_tokens_mask) if not mask] diff --git a/machine/translation/huggingface/hugging_face_nmt_model_trainer.py b/machine/translation/huggingface/hugging_face_nmt_model_trainer.py index 0f69e55b..ee110708 100644 --- a/machine/translation/huggingface/hugging_face_nmt_model_trainer.py +++ b/machine/translation/huggingface/hugging_face_nmt_model_trainer.py @@ -535,6 +535,7 @@ def decorator(*args, **kwargs): try: return function(batch_size, *args, **kwargs) except Exception as e: + logger.error(f"Attempt with batch_size={batch_size} failed with error: {e}", exc_info=True) if should_reduce_batch_size(e): gc.collect() torch.cuda.empty_cache() diff --git a/machine/translation/huggingface/transformers_compatibility.py b/machine/translation/huggingface/transformers_compatibility.py new file mode 100644 index 00000000..3b0b60a6 --- /dev/null +++ b/machine/translation/huggingface/transformers_compatibility.py @@ -0,0 +1,151 @@ +import enum +import warnings +from typing import Any, Callable, Optional, Sequence, Union, cast + +from transformers import GenerationConfig, Pipeline +from transformers.tokenization_utils_base import TruncationStrategy + +# The following classes are a port of the same classes found in transformers v4 + + +class ReturnType(enum.Enum): + TENSORS = 0 + TEXT = 1 + + +class TranslationPipeline(Pipeline): + + _pipeline_calls_generate = True + _load_processor = False + _load_image_processor = False + _load_feature_extractor = False + _load_tokenizer = True + # Make sure the docstring is updated when the default generation config is changed (in all pipelines in this file) + _default_generation_config = GenerationConfig( + max_new_tokens=256, + num_beams=4, + ) + + def __init__( + self, + framework: Optional[str] = "pt", + **kwargs, + ): + super().__init__(**kwargs) + self.framework = framework + + def _sanitize_parameters( + self, + src_lang=None, + tgt_lang=None, + return_tensors=None, + return_text=None, + return_type=None, + clean_up_tokenization_spaces=None, + truncation=None, + stop_sequence=None, + **generate_kwargs, + ): + preprocess_params = {} + if truncation is not None: + preprocess_params["truncation"] = truncation + + forward_params = generate_kwargs + + postprocess_params = {} + if return_tensors is not None and return_type is None: + return_type = ReturnType.TENSORS if return_tensors else ReturnType.TEXT + if return_type is not None: + postprocess_params["return_type"] = return_type + + if clean_up_tokenization_spaces is not None: + postprocess_params["clean_up_tokenization_spaces"] = clean_up_tokenization_spaces + + if stop_sequence is not None and self.tokenizer is not None: + stop_sequence_ids = self.tokenizer.encode(stop_sequence, add_special_tokens=False) + if len(stop_sequence_ids) > 1: + warnings.warn( + "Stopping on a multiple token sequence is not yet supported on transformers. The first token of" + " the stop sequence will be used as the stop sequence string in the interim." + ) + generate_kwargs["eos_token_id"] = stop_sequence_ids[0] + + if self.assistant_model is not None: + forward_params["assistant_model"] = self.assistant_model + if self.assistant_tokenizer is not None: + forward_params["tokenizer"] = self.tokenizer + forward_params["assistant_tokenizer"] = self.assistant_tokenizer + + if src_lang is not None: + preprocess_params["src_lang"] = src_lang + if tgt_lang is not None: + preprocess_params["tgt_lang"] = tgt_lang + if src_lang is None and tgt_lang is None: + # Backward compatibility, direct arguments use is preferred. + task = generate_kwargs.get("task", self.task) + items = task.split("_") + if task and len(items) == 4: + # translation, XX, to YY + preprocess_params["src_lang"] = items[1] + preprocess_params["tgt_lang"] = items[3] + return preprocess_params, forward_params, postprocess_params + + def _parse_and_tokenize(self, *args, truncation): + if self.tokenizer is None: + raise RuntimeError("No tokenizer is specified.") + prefix = self.prefix if self.prefix is not None else "" + if isinstance(args[0], list): + if self.tokenizer.pad_token_id is None: + raise ValueError("Please make sure that the tokenizer has a pad_token_id when using a batch input") + args = ([prefix + arg for arg in args[0]],) + padding = True + + elif isinstance(args[0], str): + args = (prefix + args[0],) + padding = False + else: + raise TypeError( + f" `args[0]`: {args[0]} have the wrong format. The should be either of type `str` or type `list`" + ) + inputs = self.tokenizer(*args, padding=padding, truncation=truncation, return_tensors=self.framework) + # This is produced by tokenizers but is an invalid generate kwargs + if "token_type_ids" in inputs: + del inputs["token_type_ids"] + return inputs + + def __call__(self, *args: Sequence[Union[str, Sequence[str]]], **kwargs: Any) -> list[dict[str, str]]: + result = super().__call__(*args, **kwargs) + if not isinstance(result, list) or not args: + return cast(list[dict[str, str]], result or []) + if ( + isinstance(args[0], list) + and all(isinstance(el, str) for el in args[0]) + and all(isinstance(res, (list, tuple)) and len(res) == 1 for res in result) + ): + return [res[0] for res in result] # type: ignore + + return cast(list[dict[str, str]], result) + + def check_inputs(self, input_length: int, min_length: int, max_length: int): + if input_length > 0.9 * max_length: + warnings.warn( + f"Your input_length: {input_length} is bigger than 0.9 * max_length: {max_length}. You might consider " + "increasing your max_length manually, e.g. translator('...', max_length=400)" + ) + return True + + def preprocess( + self, + *args, + truncation=TruncationStrategy.DO_NOT_TRUNCATE, + src_lang: str | None = None, + tgt_lang: str | None = None, + ): + if self.tokenizer: + build_inputs = getattr(self.tokenizer, "_build_translation_inputs", None) + if callable(build_inputs): + build_inputs_fn = cast(Callable[..., Any], build_inputs) + return build_inputs_fn( + *args, return_tensors=self.framework, truncation=truncation, src_lang=src_lang, tgt_lang=tgt_lang + ) + return self._parse_and_tokenize(*args, truncation=truncation) diff --git a/tests/translation/huggingface/test_hugging_face_nmt_engine.py b/tests/translation/huggingface/test_hugging_face_nmt_engine.py index 37653244..44dffb7e 100644 --- a/tests/translation/huggingface/test_hugging_face_nmt_engine.py +++ b/tests/translation/huggingface/test_hugging_face_nmt_engine.py @@ -19,46 +19,51 @@ def test_translate_n_batch_beam(output_attentions: bool) -> None: "stas/tiny-m2m_100", src_lang="en", tgt_lang="es", - num_beams=5, + num_beams=2, max_length=10, - do_sample=False, + device="cpu", # Keep test results consistent across platforms output_attentions=output_attentions, ) as engine: results = engine.translate_n_batch( n=2, segments=["This is a test string", "Hello, world!"], ) - assert results[0][0].translation == "skaberskaberskaber Dollar Dollar ፤ ፤ gerekir gerekir" - assert results[0][0].confidences[0] == approx(1.1e-05, 0.01) + assert results[0][0].translation == "skaberskaber Dollar Dollar ፤ ፤ gerekir gerekir" + assert results[0][0].confidences[0] == approx(1.08e-05, 0.01) assert results[0][0].sequence_confidence == approx(_get_sequence_confidence(results[0][0]), 0.01) - assert str(results[0][0].alignment) == ("0-3 2-0 2-1 2-2 2-4 4-5 4-6 4-7 4-8" if output_attentions else "") + assert str(results[0][0].alignment) == ("0-2 2-0 2-1 2-3 4-4 4-5 4-6 4-7" if output_attentions else "") - assert results[0][1].translation == "skaberskaber Dollar Dollar Dollar ፤ ፤ gerekir gerekir" - assert results[0][1].confidences[0] == approx(1.1e-05, 0.01) + assert results[0][1].translation == "skaberskaber Dollar Dollar ፤ ፤ ፤ gerekir" + assert results[0][1].confidences[0] == approx(1.08e-05, 0.01) assert results[0][1].sequence_confidence == approx(_get_sequence_confidence(results[0][0]), 0.01) - assert str(results[0][1].alignment) == ("0-3 2-0 2-1 2-2 2-4 4-5 4-6 4-7 4-8" if output_attentions else "") + assert str(results[0][1].alignment) == ("0-2 2-0 2-1 2-3 4-4 4-5 4-6 4-7" if output_attentions else "") - assert results[1][0].translation == "skaberskaberskaber Dollar Dollar ፤ ፤ gerekir gerekir" - assert results[1][0].confidences[0] == approx(1.1e-05, 0.01) - assert str(results[1][0].alignment) == ("0-1 0-2 0-3 0-4 0-8 1-0 3-5 3-6 3-7" if output_attentions else "") + assert results[1][0].translation == "skaberskaber Dollar Dollar ፤ ፤ gerekir gerekir" + assert results[1][0].confidences[0] == approx(1.08e-05, 0.01) + assert str(results[1][0].alignment) == ("0-0 0-1 0-2 0-3 0-7 3-4 3-5 3-6" if output_attentions else "") assert results[1][0].sequence_confidence == approx(_get_sequence_confidence(results[0][0]), 0.01) - assert results[1][1].translation == "skaberskaber Dollar Dollar Dollar ፤ ፤ gerekir gerekir" - assert results[1][1].confidences[0] == approx(1.1e-05, 0.01) - assert str(results[1][1].alignment) == ("0-1 0-2 0-3 0-4 0-8 1-0 3-5 3-6 3-7" if output_attentions else "") + assert results[1][1].translation == "skaberskaber Dollar Dollar ፤ ፤ ፤ gerekir" + assert results[1][1].confidences[0] == approx(1.08e-05, 0.01) + assert str(results[1][1].alignment) == ("0-0 0-1 0-2 0-3 0-7 3-4 3-5 3-6" if output_attentions else "") assert results[1][1].sequence_confidence == approx(_get_sequence_confidence(results[0][0]), 0.01) @mark.parametrize("output_attentions", [True, False]) def test_translate_greedy(output_attentions: bool) -> None: with HuggingFaceNmtEngine( - "stas/tiny-m2m_100", src_lang="en", tgt_lang="es", max_length=10, output_attentions=output_attentions + "stas/tiny-m2m_100", + src_lang="en", + tgt_lang="es", + max_length=10, + device="cpu", # Keep test results consistent across platforms + output_attentions=output_attentions, ) as engine: result = engine.translate("This is a test string") - assert result.translation == "skaberskaberskaber Dollar Dollar Dollar ፤ gerekir gerekir" - assert result.confidences[0] == approx(1.1e-05, 0.01) - assert result.sequence_confidence == -1.0 - assert str(result.alignment) == ("0-3 2-0 2-1 2-2 2-4 4-5 4-6 4-7 4-8" if output_attentions else "") + assert result.translation == "skaberskaber Dollar Dollar ፤ ፤ gerekir gerekir" + assert result.confidences[0] == approx(1.08e-05, 0.01) + assert result.sequence_confidence == approx(_get_sequence_confidence(result), 0.01) + assert str(result.alignment) == ("0-2 2-0 2-1 2-3 4-4 4-5 4-6 4-7" if output_attentions else "") @mark.parametrize("output_attentions", [True, False]) @@ -68,4 +73,5 @@ def test_construct_invalid_lang(output_attentions: bool) -> None: def _get_sequence_confidence(result: TranslationResult) -> float: - return exp(sum([log(c) for c in result.confidences]) / len(result.confidences)) + # Inject a 0 score for the BOS token + return exp(sum([log(c) for c in result.confidences] + [0]) / (len(result.confidences) + 1)) From a7185d82bd042a83dec8046d9e98239691ccf439 Mon Sep 17 00:00:00 2001 From: Peter Chapman Date: Mon, 14 Sep 2026 10:25:18 +1200 Subject: [PATCH 03/14] Use auto_find_batch_size --- machine/jobs/settings.yaml | 1 + .../hugging_face_nmt_model_trainer.py | 82 +------------------ 2 files changed, 3 insertions(+), 80 deletions(-) diff --git a/machine/jobs/settings.yaml b/machine/jobs/settings.yaml index 18174033..50bd3be7 100644 --- a/machine/jobs/settings.yaml +++ b/machine/jobs/settings.yaml @@ -7,6 +7,7 @@ default: huggingface: parent_model_name: facebook/nllb-200-distilled-1.3B train_params: + auto_find_batch_size: true do_train: true optim: adamw_torch warmup_steps: 1000 diff --git a/machine/translation/huggingface/hugging_face_nmt_model_trainer.py b/machine/translation/huggingface/hugging_face_nmt_model_trainer.py index ee110708..0733f60a 100644 --- a/machine/translation/huggingface/hugging_face_nmt_model_trainer.py +++ b/machine/translation/huggingface/hugging_face_nmt_model_trainer.py @@ -4,24 +4,18 @@ import os import re from pathlib import Path -from typing import Any, Callable, Dict, List, Optional, Tuple, Union, cast +from typing import Any, Callable, List, Optional, Union, cast import torch # pyright: ignore[reportMissingImports] -from accelerate import Accelerator # pyright: ignore[reportMissingImports] -from accelerate.utils.memory import should_reduce_batch_size # pyright: ignore[reportMissingImports] from datasets.arrow_dataset import Dataset from sacremoses import MosesPunctNormalizer from torch import Tensor # pyright: ignore[reportMissingImports] -from torch.nn import Module # pyright: ignore[reportMissingImports] -from torch.optim.lr_scheduler import LambdaLR # pyright: ignore[reportMissingImports] -from torch.optim.optimizer import Optimizer # pyright: ignore[reportMissingImports] from torch.utils.checkpoint import checkpoint # pyright: ignore[reportMissingImports] # noqa: F401 from transformers import ( AutoConfig, AutoModelForSeq2SeqLM, AutoTokenizer, DataCollatorForSeq2Seq, - EvalPrediction, M2M100ForConditionalGeneration, M2M100Tokenizer, MBart50Tokenizer, @@ -29,7 +23,6 @@ NllbTokenizer, PreTrainedModel, PreTrainedTokenizer, - PreTrainedTokenizerBase, PreTrainedTokenizerFast, Seq2SeqTrainer, Seq2SeqTrainingArguments, @@ -390,11 +383,10 @@ def preprocess_function(examples): pad_to_multiple_of=8 if self._training_args.fp16 else None, ) - self._trainer = AutoGradientAccumulationStepsSeq2SeqTrainer( + self._trainer = Seq2SeqTrainer( model=model, args=self._training_args, train_dataset=cast(Any, train_dataset), - tokenizer=tokenizer, data_collator=data_collator, callbacks=[ _ProgressCallback( @@ -478,76 +470,6 @@ def on_step_end(self, args: TrainingArguments, state: TrainerState, control: Tra ) -class AutoGradientAccumulationStepsSeq2SeqTrainer(Seq2SeqTrainer): - def __init__( - self, - model: Union[PreTrainedModel, Module], - args: Seq2SeqTrainingArguments, - data_collator: Any, - train_dataset: Optional[Dataset] = None, - eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None, - tokenizer: Optional[PreTrainedTokenizerBase] = None, - model_init: Optional[Callable[[], PreTrainedModel]] = None, - compute_loss_func: Callable | None = None, - compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None, - callbacks: Optional[List[TrainerCallback]] = None, - optimizers: Tuple[Optional[Optimizer], Optional[LambdaLR]] = (None, None), - preprocess_logits_for_metrics: Optional[Callable[[Tensor, Tensor], Tensor]] = None, - ): - super().__init__( - model, - args, - data_collator, - train_dataset, - eval_dataset, # type: ignore - tokenizer, - model_init, - compute_loss_func, - compute_metrics, - callbacks, - optimizers, - preprocess_logits_for_metrics, - ) - - def _inner_training_loop( - self, batch_size=None, args=None, resume_from_checkpoint=None, trial=None, ignore_keys_for_eval=None - ): - inner_training_loop = find_executable_batch_size(super()._inner_training_loop, batch_size, self.accelerator) - return inner_training_loop( - args=args, - resume_from_checkpoint=resume_from_checkpoint, - trial=trial, - ignore_keys_for_eval=ignore_keys_for_eval, - ) - - -def find_executable_batch_size(function: Callable, starting_batch_size, accelerator: Accelerator): - batch_size = starting_batch_size - - def decorator(*args, **kwargs): - nonlocal batch_size - gc.collect() - torch.cuda.empty_cache() - - while True: - if batch_size == 0: - raise RuntimeError("No executable batch size found, reached zero.") - try: - return function(batch_size, *args, **kwargs) - except Exception as e: - logger.error(f"Attempt with batch_size={batch_size} failed with error: {e}", exc_info=True) - if should_reduce_batch_size(e): - gc.collect() - torch.cuda.empty_cache() - batch_size //= 2 - accelerator.gradient_accumulation_steps = accelerator.gradient_accumulation_steps * 2 - kwargs["args"].gradient_accumulation_steps = accelerator.gradient_accumulation_steps - else: - raise - - return decorator - - def add_lang_code_to_tokenizer(tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast], lang_code: str): if isinstance(tokenizer, M2M100Tokenizer): lang_token = "__" + lang_code + "__" From 62216420e164ee5f5da8d317107c14191ff9d6b5 Mon Sep 17 00:00:00 2001 From: Peter Chapman Date: Thu, 17 Sep 2026 11:36:14 +1200 Subject: [PATCH 04/14] Revert Use auto_find_batch_size --- machine/jobs/settings.yaml | 1 - .../hugging_face_nmt_model_trainer.py | 65 ++++++++++++++++++- 2 files changed, 64 insertions(+), 2 deletions(-) diff --git a/machine/jobs/settings.yaml b/machine/jobs/settings.yaml index 50bd3be7..18174033 100644 --- a/machine/jobs/settings.yaml +++ b/machine/jobs/settings.yaml @@ -7,7 +7,6 @@ default: huggingface: parent_model_name: facebook/nllb-200-distilled-1.3B train_params: - auto_find_batch_size: true do_train: true optim: adamw_torch warmup_steps: 1000 diff --git a/machine/translation/huggingface/hugging_face_nmt_model_trainer.py b/machine/translation/huggingface/hugging_face_nmt_model_trainer.py index 0733f60a..348798c8 100644 --- a/machine/translation/huggingface/hugging_face_nmt_model_trainer.py +++ b/machine/translation/huggingface/hugging_face_nmt_model_trainer.py @@ -7,9 +7,12 @@ from typing import Any, Callable, List, Optional, Union, cast import torch # pyright: ignore[reportMissingImports] +from accelerate import Accelerator # pyright: ignore[reportMissingImports] +from accelerate.utils.memory import release_memory, should_reduce_batch_size # pyright: ignore[reportMissingImports] from datasets.arrow_dataset import Dataset from sacremoses import MosesPunctNormalizer from torch import Tensor # pyright: ignore[reportMissingImports] +from torch.nn import Module # pyright: ignore[reportMissingImports] from torch.utils.checkpoint import checkpoint # pyright: ignore[reportMissingImports] # noqa: F401 from transformers import ( AutoConfig, @@ -383,7 +386,7 @@ def preprocess_function(examples): pad_to_multiple_of=8 if self._training_args.fp16 else None, ) - self._trainer = Seq2SeqTrainer( + self._trainer = AutoGradientAccumulationStepsSeq2SeqTrainer( model=model, args=self._training_args, train_dataset=cast(Any, train_dataset), @@ -470,6 +473,66 @@ def on_step_end(self, args: TrainingArguments, state: TrainerState, control: Tra ) +class AutoGradientAccumulationStepsSeq2SeqTrainer(Seq2SeqTrainer): + def __init__( + self, + model: Union[PreTrainedModel, Module], + args: Seq2SeqTrainingArguments, + data_collator: Any, + train_dataset: Optional[Dataset] = None, + callbacks: Optional[List[TrainerCallback]] = None, + ): + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + callbacks=callbacks, + ) + + def _inner_training_loop( + self, batch_size=None, args=None, resume_from_checkpoint=None, trial=None, ignore_keys_for_eval=None + ): + inner_training_loop = self.find_executable_batch_size( + super()._inner_training_loop, batch_size, self.accelerator + ) + return inner_training_loop( + args=args, + resume_from_checkpoint=resume_from_checkpoint, + trial=trial, + ignore_keys_for_eval=ignore_keys_for_eval, + ) + + def find_executable_batch_size(self, function: Callable, starting_batch_size, accelerator: Accelerator): + batch_size = starting_batch_size + + def decorator(*args, **kwargs): + nonlocal batch_size + self.accelerator.free_memory() + gc.collect() + torch.cuda.empty_cache() + + while True: + if batch_size == 0: + raise RuntimeError("No executable batch size found, reached zero.") + try: + return function(batch_size, *args, **kwargs) + except Exception as e: + logger.error(f"Attempt with batch_size={batch_size} failed with error: {e}", exc_info=True) + if should_reduce_batch_size(e): + self.accelerator.free_memory() + gc.collect() + torch.cuda.empty_cache() + batch_size //= 2 + self._update_auto_batch_size(batch_size) + accelerator.gradient_accumulation_steps = accelerator.gradient_accumulation_steps * 2 + kwargs["args"].gradient_accumulation_steps = accelerator.gradient_accumulation_steps + else: + raise + + return decorator + + def add_lang_code_to_tokenizer(tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast], lang_code: str): if isinstance(tokenizer, M2M100Tokenizer): lang_token = "__" + lang_code + "__" From fe2945b812de75e2f54a6759d9c1ecee8c641ace Mon Sep 17 00:00:00 2001 From: Peter Chapman Date: Thu, 17 Sep 2026 14:52:38 +1200 Subject: [PATCH 05/14] Fixes from code review --- .../hugging_face_nmt_model_factory.py | 10 ++++ machine/jobs/nmt_build_options.py | 1 + machine/jobs/settings.yaml | 2 +- .../huggingface/hugging_face_nmt_engine.py | 50 +++++++----------- .../huggingface/hugging_face_nmt_model.py | 6 ++- .../hugging_face_nmt_model_trainer.py | 50 +++++++----------- .../huggingface/transformers_compatibility.py | 51 ++++--------------- samples/machine_translation.ipynb | 4 +- 8 files changed, 63 insertions(+), 111 deletions(-) diff --git a/machine/jobs/huggingface/hugging_face_nmt_model_factory.py b/machine/jobs/huggingface/hugging_face_nmt_model_factory.py index 5a0a23ef..56065aa5 100644 --- a/machine/jobs/huggingface/hugging_face_nmt_model_factory.py +++ b/machine/jobs/huggingface/hugging_face_nmt_model_factory.py @@ -1,5 +1,6 @@ import logging import tarfile +import warnings from pathlib import Path from typing import Any, cast @@ -26,6 +27,15 @@ def __init__(self, config: Any) -> None: self._config = config args = config.huggingface.train_params.to_dict() args["output_dir"] = str(self._model_dir) + # Allow group_by_length backwards compatibility + if "group_by_length" in args: + warnings.warn( + "'group_by_length' is deprecated and will be removed in a future release.", + category=DeprecationWarning, + stacklevel=2, + ) + if args.pop("group_by_length") is True: + args["train_sampling_strategy"] = "group_by_length" # Use "max_steps" from root for backward compatibility if "max_steps" in self._config.huggingface: args["max_steps"] = self._config.huggingface.max_steps diff --git a/machine/jobs/nmt_build_options.py b/machine/jobs/nmt_build_options.py index d0954de6..21039dba 100644 --- a/machine/jobs/nmt_build_options.py +++ b/machine/jobs/nmt_build_options.py @@ -15,6 +15,7 @@ class TrainParams(BaseModel): per_device_train_batch_size: int | None = None gradient_accumulation_steps: int | None = None label_smoothing_factor: float | None = None + group_by_length: bool | None = None train_sampling_strategy: str | None = None gradient_checkpointing: bool | None = None lr_scheduler_type: str | None = None diff --git a/machine/jobs/settings.yaml b/machine/jobs/settings.yaml index 18174033..c121343b 100644 --- a/machine/jobs/settings.yaml +++ b/machine/jobs/settings.yaml @@ -13,7 +13,7 @@ default: per_device_train_batch_size: 64 gradient_accumulation_steps: 1 label_smoothing_factor: 0.2 - train_sampling_strategy : group_by_length + train_sampling_strategy: group_by_length gradient_checkpointing: true lr_scheduler_type: cosine learning_rate: 0.0002 diff --git a/machine/translation/huggingface/hugging_face_nmt_engine.py b/machine/translation/huggingface/hugging_face_nmt_engine.py index 69e33ab3..aa161a28 100644 --- a/machine/translation/huggingface/hugging_face_nmt_engine.py +++ b/machine/translation/huggingface/hugging_face_nmt_engine.py @@ -17,6 +17,7 @@ NllbTokenizer, PreTrainedModel, PreTrainedTokenizer, + PreTrainedTokenizerBase, PreTrainedTokenizerFast, ) from transformers.generation.utils import GenerateBeamEncoderDecoderOutput, GenerateEncoderDecoderOutput @@ -42,15 +43,19 @@ def __init__( **pipeline_kwargs, ) -> None: self._pipeline_kwargs = pipeline_kwargs + if self._pipeline_kwargs.get("output_attentions") is None: + self._pipeline_kwargs["output_attentions"] = True if isinstance(model, PreTrainedModel): self._model = model + if self._pipeline_kwargs["output_attentions"]: + model.set_attn_implementation("eager") self._model.eval() self._is_model_owned = False else: model_config = AutoConfig.from_pretrained(str(model), label2id={}, id2label={}, num_labels=0) # If output_attentions is True or None, we need to set the attn_implementation to eager to get the attentions - attn_implementation = "eager" if self._pipeline_kwargs.get("output_attentions", True) else "sdpa" + attn_implementation = "eager" if self._pipeline_kwargs["output_attentions"] else "sdpa" self._model = cast( PreTrainedModel, @@ -108,7 +113,7 @@ def __init__( self._pipeline = SilTranslationPipeline( model=self._model, - tokenizer=cast(PreTrainedTokenizer, self.tokenizer), + tokenizer=self.tokenizer, mpn=self._mpn, batch_size=self._batch_size, **self._pipeline_kwargs, @@ -200,7 +205,7 @@ class SilTranslationPipeline(TranslationPipeline): def __init__( self, model: PreTrainedModel, - tokenizer: PreTrainedTokenizer, + tokenizer: PreTrainedTokenizerBase, batch_size: int, mpn: Optional[MosesPunctNormalizer] = None, **kwargs, @@ -247,10 +252,7 @@ def _forward(self, model_inputs, **generate_kwargs): raise RuntimeError("No tokenizer is specified.") in_b, input_length = model_inputs["input_ids"].shape - if "input_tokens" in model_inputs: - input_tokens = model_inputs.pop("input_tokens") - else: - input_tokens = [self.tokenizer.convert_ids_to_tokens(seq) for seq in model_inputs["input_ids"]] + input_tokens = model_inputs.pop("input_tokens") self.check_inputs(input_length, self.generation_config.min_length, self.generation_config.max_length) output = cast(Any, self.model).generate( @@ -278,29 +280,15 @@ def _forward(self, model_inputs, **generate_kwargs): else: raise RuntimeError("Cannot postprocess the output of the model.") - try: - transition_scores = cast( - torch.Tensor, - cast(Any, self.model).compute_transition_scores( - output_ids, - scores, - beam_indices, - normalize_logits=True, - ), - ) - except Exception: - output_ids = output_ids.to("cpu") - scores = tuple(score.to("cpu") for score in scores) - beam_indices = beam_indices.to("cpu") if beam_indices is not None else None - transition_scores = cast( - torch.Tensor, - cast(Any, self.model).compute_transition_scores( - output_ids, - scores, - beam_indices, - normalize_logits=True, - ), - ) + transition_scores = cast( + torch.Tensor, + cast(Any, self.model).compute_transition_scores( + output_ids, + scores, + beam_indices, + normalize_logits=False, + ), + ) if beam_indices is None: beam_indices = torch.zeros_like(output_ids) @@ -329,7 +317,7 @@ def _forward(self, model_inputs, **generate_kwargs): start_index = 0 if self.model.config.decoder_start_token_id is not None: start_index = 1 - if self.generation_config.output_attentions: + if self.generation_config.output_attentions is not False: assert attentions is not None num_heads = attentions[0][0].shape[1] diff --git a/machine/translation/huggingface/hugging_face_nmt_model.py b/machine/translation/huggingface/hugging_face_nmt_model.py index 26d65897..7f82823d 100644 --- a/machine/translation/huggingface/hugging_face_nmt_model.py +++ b/machine/translation/huggingface/hugging_face_nmt_model.py @@ -17,14 +17,16 @@ class HuggingFaceNmtModel(TranslationModel): def __init__( self, - model: Union[PreTrainedModel, Path], + model: Union[PreTrainedModel, Path, str], parent_model_name: str, training_args: Optional[Seq2SeqTrainingArguments] = None, **pipeline_kwargs, ) -> None: self._model = model if isinstance(model, PreTrainedModel): - self._model_path = Path(str(model.name_or_path)) + if model.name_or_path is None: + raise ValueError("No model name_or_path specified.") + self._model_path = Path(model.name_or_path) else: self._model_path = Path(model) self._parent_model_name = parent_model_name diff --git a/machine/translation/huggingface/hugging_face_nmt_model_trainer.py b/machine/translation/huggingface/hugging_face_nmt_model_trainer.py index 348798c8..2858f330 100644 --- a/machine/translation/huggingface/hugging_face_nmt_model_trainer.py +++ b/machine/translation/huggingface/hugging_face_nmt_model_trainer.py @@ -8,7 +8,7 @@ import torch # pyright: ignore[reportMissingImports] from accelerate import Accelerator # pyright: ignore[reportMissingImports] -from accelerate.utils.memory import release_memory, should_reduce_batch_size # pyright: ignore[reportMissingImports] +from accelerate.utils.memory import should_reduce_batch_size # pyright: ignore[reportMissingImports] from datasets.arrow_dataset import Dataset from sacremoses import MosesPunctNormalizer from torch import Tensor # pyright: ignore[reportMissingImports] @@ -35,7 +35,6 @@ ) from transformers.tokenization_utils_base import BatchEncoding from transformers.trainer_callback import TrainerControl, TrainerState -from transformers.trainer_utils import get_last_checkpoint from transformers.training_args import TrainingArguments from ...corpora.parallel_text_corpus import ParallelTextCorpus @@ -114,24 +113,6 @@ def train( progress: Optional[Callable[[ProgressStatus], None]] = None, check_canceled: Optional[Callable[[], None]] = None, ) -> None: - last_checkpoint = None - if ( - self._training_args.output_dir is not None - and os.path.isdir(self._training_args.output_dir) - and self._training_args.resume_from_checkpoint is not None - ): - last_checkpoint = get_last_checkpoint(self._training_args.output_dir) - if last_checkpoint is None and any(os.path.isfile(p) for p in os.listdir(self._training_args.output_dir)): - raise ValueError( - f"Output directory ({self._training_args.output_dir}) already exists and is not empty. " - "Remove --resume_from_checkpoint to overcome." - ) - elif last_checkpoint is not None and self._training_args.resume_from_checkpoint is None: - logger.info( - f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change " - "the `--output_dir` or remove `--resume_from_checkpoint` to train from scratch." - ) - # Set seed before initializing model. set_seed(self._training_args.seed) @@ -170,7 +151,7 @@ def train( def find_missing_characters(tokenizer: Any, train_dataset: Dataset, lang_codes: List[str]) -> List[str]: vocab = tokenizer.get_vocab().keys() charset = set() - mpn_normalize = True if isinstance(tokenizer, NllbTokenizer) else False + mpn_normalize = isinstance(tokenizer, NllbTokenizer) for ex in train_dataset["translation"]: for lang_code in lang_codes: ex_text = ex[lang_code] @@ -299,12 +280,20 @@ def batch_prepare_for_model( batch_tokens: List[List[str]], return_tensors: Optional[Union[str, TensorType]] = None, ) -> BatchEncoding: - return tokenizer( - batch_tokens, - is_split_into_words=True, - add_special_tokens=False, - return_tensors=return_tensors, - ) + input_ids: List[List[int]] = [] + attention_mask: List[List[int]] = [] + + for tokens in batch_tokens: + ids = cast(List[int], tokenizer.convert_tokens_to_ids(tokens)) + input_ids.append(ids) + attention_mask.append([1] * len(ids)) + + batch_outputs = { + "input_ids": input_ids, + "attention_mask": attention_mask, + } + + return BatchEncoding(batch_outputs, tensor_type=return_tensors) def preprocess_function(examples): # Add one to the content_type in order to convert back from ClassLabels which are enumerated from 0, not 1 @@ -401,13 +390,8 @@ def preprocess_function(examples): ) logger.info("Train NMT model") - ckpt = None - if self._training_args.resume_from_checkpoint is not None: - ckpt = self._training_args.resume_from_checkpoint - elif last_checkpoint is not None: - ckpt = last_checkpoint train_result = self._trainer.train( - resume_from_checkpoint=ckpt, + resume_from_checkpoint=self._training_args.resume_from_checkpoint, ) self._metrics = train_result.metrics diff --git a/machine/translation/huggingface/transformers_compatibility.py b/machine/translation/huggingface/transformers_compatibility.py index 3b0b60a6..45d6fdb4 100644 --- a/machine/translation/huggingface/transformers_compatibility.py +++ b/machine/translation/huggingface/transformers_compatibility.py @@ -1,6 +1,6 @@ import enum import warnings -from typing import Any, Callable, Optional, Sequence, Union, cast +from typing import Any, Callable, Sequence, Union, cast from transformers import GenerationConfig, Pipeline from transformers.tokenization_utils_base import TruncationStrategy @@ -28,19 +28,14 @@ class TranslationPipeline(Pipeline): def __init__( self, - framework: Optional[str] = "pt", **kwargs, ): super().__init__(**kwargs) - self.framework = framework def _sanitize_parameters( self, src_lang=None, tgt_lang=None, - return_tensors=None, - return_text=None, - return_type=None, clean_up_tokenization_spaces=None, truncation=None, stop_sequence=None, @@ -53,11 +48,6 @@ def _sanitize_parameters( forward_params = generate_kwargs postprocess_params = {} - if return_tensors is not None and return_type is None: - return_type = ReturnType.TENSORS if return_tensors else ReturnType.TEXT - if return_type is not None: - postprocess_params["return_type"] = return_type - if clean_up_tokenization_spaces is not None: postprocess_params["clean_up_tokenization_spaces"] = clean_up_tokenization_spaces @@ -90,29 +80,6 @@ def _sanitize_parameters( preprocess_params["tgt_lang"] = items[3] return preprocess_params, forward_params, postprocess_params - def _parse_and_tokenize(self, *args, truncation): - if self.tokenizer is None: - raise RuntimeError("No tokenizer is specified.") - prefix = self.prefix if self.prefix is not None else "" - if isinstance(args[0], list): - if self.tokenizer.pad_token_id is None: - raise ValueError("Please make sure that the tokenizer has a pad_token_id when using a batch input") - args = ([prefix + arg for arg in args[0]],) - padding = True - - elif isinstance(args[0], str): - args = (prefix + args[0],) - padding = False - else: - raise TypeError( - f" `args[0]`: {args[0]} have the wrong format. The should be either of type `str` or type `list`" - ) - inputs = self.tokenizer(*args, padding=padding, truncation=truncation, return_tensors=self.framework) - # This is produced by tokenizers but is an invalid generate kwargs - if "token_type_ids" in inputs: - del inputs["token_type_ids"] - return inputs - def __call__(self, *args: Sequence[Union[str, Sequence[str]]], **kwargs: Any) -> list[dict[str, str]]: result = super().__call__(*args, **kwargs) if not isinstance(result, list) or not args: @@ -141,11 +108,11 @@ def preprocess( src_lang: str | None = None, tgt_lang: str | None = None, ): - if self.tokenizer: - build_inputs = getattr(self.tokenizer, "_build_translation_inputs", None) - if callable(build_inputs): - build_inputs_fn = cast(Callable[..., Any], build_inputs) - return build_inputs_fn( - *args, return_tensors=self.framework, truncation=truncation, src_lang=src_lang, tgt_lang=tgt_lang - ) - return self._parse_and_tokenize(*args, truncation=truncation) + if self.tokenizer is None: + raise RuntimeError("No tokenizer is specified.") + build_inputs = getattr(self.tokenizer, "_build_translation_inputs", None) + if callable(build_inputs): + build_inputs_fn = cast(Callable[..., Any], build_inputs) + return build_inputs_fn( + *args, return_tensors="pt", truncation=truncation, src_lang=src_lang, tgt_lang=tgt_lang + ) diff --git a/samples/machine_translation.ipynb b/samples/machine_translation.ipynb index 89103f23..2eaa9eee 100644 --- a/samples/machine_translation.ipynb +++ b/samples/machine_translation.ipynb @@ -16,7 +16,7 @@ "outputs": [], "source": [ "%pip install sil-machine[thot]\n", - "%pip install transformers==4.34.0 datasets sacremoses accelerate==0.26.1 gcsfs==2023.3.0" + "%pip install transformers==5.14.1 datasets sacremoses accelerate==1.14.0 gcsfs==2023.3.0" ] }, { @@ -382,7 +382,7 @@ "from machine.translation.huggingface import HuggingFaceNmtModelTrainer\n", "\n", "training_args = Seq2SeqTrainingArguments(\n", - " output_dir=\"out/sp-en-nmt\", overwrite_output_dir=True, num_train_epochs=1, report_to=[], disable_tqdm=False\n", + " output_dir=\"out/sp-en-nmt\", num_train_epochs=1, report_to=[], disable_tqdm=False\n", ")\n", "\n", "with HuggingFaceNmtModelTrainer(\n", From 0dc46cbdcf3e714079a2cb3daa11cf64dd43b6ef Mon Sep 17 00:00:00 2001 From: Damien Daspit Date: Thu, 17 Sep 2026 17:57:03 -0400 Subject: [PATCH 06/14] More memory efficient method for retrieving token scores --- .../huggingface/hugging_face_nmt_engine.py | 50 +++++++++++++++---- 1 file changed, 41 insertions(+), 9 deletions(-) diff --git a/machine/translation/huggingface/hugging_face_nmt_engine.py b/machine/translation/huggingface/hugging_face_nmt_engine.py index aa161a28..b7d32526 100644 --- a/machine/translation/huggingface/hugging_face_nmt_engine.py +++ b/machine/translation/huggingface/hugging_face_nmt_engine.py @@ -270,6 +270,8 @@ def _forward(self, model_inputs, **generate_kwargs): assert scores is not None and beam_indices is not None sequences_scores = output.sequences_scores attentions = output.cross_attentions + # Beam search scores are already log probabilities. + normalize_logits = False elif isinstance(output, GenerateEncoderDecoderOutput): output_ids = output.sequences beam_indices = None @@ -277,18 +279,12 @@ def _forward(self, model_inputs, **generate_kwargs): scores = output.scores sequences_scores = None attentions = output.cross_attentions + # Greedy search scores are unnormalized logits. + normalize_logits = True else: raise RuntimeError("Cannot postprocess the output of the model.") - transition_scores = cast( - torch.Tensor, - cast(Any, self.model).compute_transition_scores( - output_ids, - scores, - beam_indices, - normalize_logits=False, - ), - ) + transition_scores = _compute_transition_scores(output_ids, scores, beam_indices, normalize_logits) if beam_indices is None: beam_indices = torch.zeros_like(output_ids) @@ -465,5 +461,41 @@ def torch_gather_nd(params: torch.Tensor, indices: torch.Tensor, batch_dim: int return out.reshape(*index_shape, *tail_sizes) +def _compute_transition_scores( + sequences: torch.Tensor, + scores: Tuple[torch.Tensor, ...], + beam_indices: Optional[torch.Tensor], + normalize_logits: bool, +) -> torch.Tensor: + """ + Compute the log probability of each generated token. + + This is equivalent to PreTrainedModel.compute_transition_scores, but gathers one step at a time instead of + stacking the scores for every step into a single tensor, which requires several gigabytes of memory for a large + vocabulary. + """ + if beam_indices is None: + # Greedy search is equivalent to a beam search where the first (and only) beam is always selected. + beam_indices = torch.arange(scores[0].shape[0], device=sequences.device).view(-1, 1).expand(-1, len(scores)) + + # Cut the beam indices to the longest beam length. Beams that finished early are masked out below. + beam_indices_mask = beam_indices < 0 + max_beam_length = int((1 - beam_indices_mask.long()).sum(-1).max().item()) + beam_indices_mask = beam_indices_mask[:, :max_beam_length] + beam_indices = beam_indices[:, :max_beam_length].masked_fill(beam_indices_mask, 0) + + # The token generated at step i is at cut_idx + i in the sequence. + cut_idx = sequences.shape[-1] - max_beam_length + transition_scores = torch.zeros(sequences.shape[0], max_beam_length, device=sequences.device, dtype=scores[0].dtype) + for i in range(max_beam_length): + step_scores = scores[i] + if normalize_logits: + step_scores = torch.nn.functional.log_softmax(step_scores, dim=-1) + transition_scores[:, i] = step_scores[beam_indices[:, i], sequences[:, cut_idx + i]] + + transition_scores[beam_indices_mask] = 0 + return transition_scores + + def _get_encoding_fast_tokens(encoding) -> List[str]: return [token for (token, mask) in zip(encoding.tokens, encoding.special_tokens_mask) if not mask] From 39573871d53a65325f6461f26aff6ea0df435c0a Mon Sep 17 00:00:00 2001 From: Damien Daspit Date: Thu, 17 Sep 2026 23:15:24 -0400 Subject: [PATCH 07/14] Remove default generation config - fix test_update_tokenizer_missing_char --- .../huggingface/transformers_compatibility.py | 7 +--- .../test_hugging_face_nmt_engine.py | 5 ++- .../test_hugging_face_nmt_model_trainer.py | 42 ++++++++++++++----- 3 files changed, 36 insertions(+), 18 deletions(-) diff --git a/machine/translation/huggingface/transformers_compatibility.py b/machine/translation/huggingface/transformers_compatibility.py index 45d6fdb4..34438e4d 100644 --- a/machine/translation/huggingface/transformers_compatibility.py +++ b/machine/translation/huggingface/transformers_compatibility.py @@ -2,7 +2,7 @@ import warnings from typing import Any, Callable, Sequence, Union, cast -from transformers import GenerationConfig, Pipeline +from transformers import Pipeline from transformers.tokenization_utils_base import TruncationStrategy # The following classes are a port of the same classes found in transformers v4 @@ -20,11 +20,6 @@ class TranslationPipeline(Pipeline): _load_image_processor = False _load_feature_extractor = False _load_tokenizer = True - # Make sure the docstring is updated when the default generation config is changed (in all pipelines in this file) - _default_generation_config = GenerationConfig( - max_new_tokens=256, - num_beams=4, - ) def __init__( self, diff --git a/tests/translation/huggingface/test_hugging_face_nmt_engine.py b/tests/translation/huggingface/test_hugging_face_nmt_engine.py index 44dffb7e..a4b7c5c9 100644 --- a/tests/translation/huggingface/test_hugging_face_nmt_engine.py +++ b/tests/translation/huggingface/test_hugging_face_nmt_engine.py @@ -60,9 +60,10 @@ def test_translate_greedy(output_attentions: bool) -> None: output_attentions=output_attentions, ) as engine: result = engine.translate("This is a test string") - assert result.translation == "skaberskaber Dollar Dollar ፤ ፤ gerekir gerekir" + assert result.translation == "skaberskaber Dollar Dollar Dollar ፤ gerekir gerekir" assert result.confidences[0] == approx(1.08e-05, 0.01) - assert result.sequence_confidence == approx(_get_sequence_confidence(result), 0.01) + # Greedy search does not produce a sequence score + assert result.sequence_confidence == -1.0 assert str(result.alignment) == ("0-2 2-0 2-1 2-3 4-4 4-5 4-6 4-7" if output_attentions else "") diff --git a/tests/translation/huggingface/test_hugging_face_nmt_model_trainer.py b/tests/translation/huggingface/test_hugging_face_nmt_model_trainer.py index ae9899af..cbbc95e7 100644 --- a/tests/translation/huggingface/test_hugging_face_nmt_model_trainer.py +++ b/tests/translation/huggingface/test_hugging_face_nmt_model_trainer.py @@ -9,12 +9,15 @@ from typing import cast from transformers import ( + M2M100Config, + M2M100ForConditionalGeneration, M2M100Tokenizer, MBart50Tokenizer, MBartTokenizer, NllbTokenizer, PreTrainedTokenizerFast, Seq2SeqTrainingArguments, + set_seed, ) from machine.corpora import DictionaryTextCorpus, MemoryText, TextRow @@ -99,7 +102,8 @@ def test_train_non_empty_corpus() -> None: def test_update_tokenizer_missing_char() -> None: - with TemporaryDirectory() as temp_dir: + with TemporaryDirectory() as model_dir, TemporaryDirectory() as temp_dir: + _create_tiny_nllb_model(model_dir) source_corpus = DictionaryTextCorpus( [ MemoryText( @@ -138,11 +142,11 @@ def test_update_tokenizer_missing_char() -> None: ) with HuggingFaceNmtModelTrainer( - "hf-internal-testing/tiny-random-nllb", + model_dir, training_args, corpus, - src_lang="en_XX", - tgt_lang="es_XX", + src_lang="eng_Latn", + tgt_lang="spa_Latn", max_src_length=20, max_tgt_length=20, add_unk_src_tokens=False, @@ -152,7 +156,7 @@ def test_update_tokenizer_missing_char() -> None: trainer_nochar.save() with HuggingFaceNmtEngine( - temp_dir, src_lang="en_XX", tgt_lang="es_XX", max_length=20 + temp_dir, src_lang="eng_Latn", tgt_lang="spa_Latn", max_length=20 ) as finetuned_engine_nochar: assert isinstance(finetuned_engine_nochar.tokenizer, PreTrainedTokenizerFast) finetuned_result_nochar = finetuned_engine_nochar.tokenizer.encode( @@ -163,11 +167,11 @@ def test_update_tokenizer_missing_char() -> None: norm_result_nochar2 = finetuned_engine_nochar.tokenizer.encode("‍") with HuggingFaceNmtModelTrainer( - "hf-internal-testing/tiny-random-nllb", + model_dir, training_args, corpus, - src_lang="en_XX", - tgt_lang="es_XX", + src_lang="eng_Latn", + tgt_lang="spa_Latn", max_src_length=20, max_tgt_length=20, add_unk_src_tokens=True, @@ -176,7 +180,9 @@ def test_update_tokenizer_missing_char() -> None: trainer_char.train() trainer_char.save() - with HuggingFaceNmtEngine(temp_dir, src_lang="en_XX", tgt_lang="es_XX", max_length=20) as finetuned_engine_char: + with HuggingFaceNmtEngine( + temp_dir, src_lang="eng_Latn", tgt_lang="spa_Latn", max_length=20 + ) as finetuned_engine_char: assert isinstance(finetuned_engine_char.tokenizer, PreTrainedTokenizerFast) finetuned_result_char = finetuned_engine_char.tokenizer.encode( "Ḻ, ḻ, Ṉ, ॽ, " + "‌ and " + "‍" + " are new characters" @@ -189,7 +195,7 @@ def test_update_tokenizer_missing_char() -> None: assert norm_result_nochar2 != norm_result_char2 assert finetuned_result_nochar != finetuned_result_char - assert finetuned_result_nochar_composite == finetuned_result_char_composite + assert finetuned_result_nochar_composite != finetuned_result_char_composite def test_update_tokenizer_missing_char_skip() -> None: @@ -535,3 +541,19 @@ def test_m2m_100_tokenizer_add_lang_code() -> None: def _row(row_ref: int, text: str) -> TextRow: return TextRow("text1", row_ref, segment=[text]) + + +def _create_tiny_nllb_model(model_dir: str) -> None: + # hf-internal-testing/tiny-random-nllb declares MBartTokenizer, which in transformers 5 discards the normalizer in + # tokenizer.json. Pair its tiny config with the real NLLB tokenizer so that input normalization is exercised. + tokenizer = cast(NllbTokenizer, NllbTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")) + config = cast(M2M100Config, M2M100Config.from_pretrained("hf-internal-testing/tiny-random-nllb")) + config.vocab_size = len(tokenizer) + config.bos_token_id = cast(int, tokenizer.bos_token_id) + config.pad_token_id = cast(int, tokenizer.pad_token_id) + config.eos_token_id = cast(int, tokenizer.eos_token_id) + config.decoder_start_token_id = config.eos_token_id + set_seed(0) + model = M2M100ForConditionalGeneration(config) + model.save_pretrained(model_dir) + tokenizer.save_pretrained(model_dir) From 2e1239f52fd1820ae6dd843b856a79bd3ca1e7a7 Mon Sep 17 00:00:00 2001 From: Damien Daspit Date: Fri, 18 Sep 2026 09:58:08 -0400 Subject: [PATCH 08/14] Fix type checking issues on Mac --- machine/translation/huggingface/hugging_face_nmt_engine.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/machine/translation/huggingface/hugging_face_nmt_engine.py b/machine/translation/huggingface/hugging_face_nmt_engine.py index b7d32526..1937a897 100644 --- a/machine/translation/huggingface/hugging_face_nmt_engine.py +++ b/machine/translation/huggingface/hugging_face_nmt_engine.py @@ -464,7 +464,7 @@ def torch_gather_nd(params: torch.Tensor, indices: torch.Tensor, batch_dim: int def _compute_transition_scores( sequences: torch.Tensor, scores: Tuple[torch.Tensor, ...], - beam_indices: Optional[torch.Tensor], + output_beam_indices: Optional[torch.Tensor], normalize_logits: bool, ) -> torch.Tensor: """ @@ -474,9 +474,11 @@ def _compute_transition_scores( stacking the scores for every step into a single tensor, which requires several gigabytes of memory for a large vocabulary. """ - if beam_indices is None: + if output_beam_indices is None: # Greedy search is equivalent to a beam search where the first (and only) beam is always selected. beam_indices = torch.arange(scores[0].shape[0], device=sequences.device).view(-1, 1).expand(-1, len(scores)) + else: + beam_indices = output_beam_indices # Cut the beam indices to the longest beam length. Beams that finished early are masked out below. beam_indices_mask = beam_indices < 0 From c2d6a2274abb43475abf3ebaecdc24190ce5b7c8 Mon Sep 17 00:00:00 2001 From: Peter Chapman Date: Thu, 24 Sep 2026 10:50:59 +1200 Subject: [PATCH 09/14] Fixes for issues raised by Claude --- .../hugging_face_nmt_model_factory.py | 2 ++ .../hugging_face_nmt_model_trainer.py | 5 ++- .../huggingface/transformers_compatibility.py | 31 +++++++++++++++++++ 3 files changed, 37 insertions(+), 1 deletion(-) diff --git a/machine/jobs/huggingface/hugging_face_nmt_model_factory.py b/machine/jobs/huggingface/hugging_face_nmt_model_factory.py index 56065aa5..5aedc7b5 100644 --- a/machine/jobs/huggingface/hugging_face_nmt_model_factory.py +++ b/machine/jobs/huggingface/hugging_face_nmt_model_factory.py @@ -36,6 +36,8 @@ def __init__(self, config: Any) -> None: ) if args.pop("group_by_length") is True: args["train_sampling_strategy"] = "group_by_length" + else: + args["train_sampling_strategy"] = "random" # Use "max_steps" from root for backward compatibility if "max_steps" in self._config.huggingface: args["max_steps"] = self._config.huggingface.max_steps diff --git a/machine/translation/huggingface/hugging_face_nmt_model_trainer.py b/machine/translation/huggingface/hugging_face_nmt_model_trainer.py index 2858f330..1564c9fa 100644 --- a/machine/translation/huggingface/hugging_face_nmt_model_trainer.py +++ b/machine/translation/huggingface/hugging_face_nmt_model_trainer.py @@ -33,7 +33,7 @@ TrainerCallback, set_seed, ) -from transformers.tokenization_utils_base import BatchEncoding +from transformers.tokenization_utils_base import BatchEncoding, PreTrainedTokenizerBase from transformers.trainer_callback import TrainerControl, TrainerState from transformers.training_args import TrainingArguments @@ -378,6 +378,7 @@ def preprocess_function(examples): self._trainer = AutoGradientAccumulationStepsSeq2SeqTrainer( model=model, args=self._training_args, + processing_class=tokenizer, train_dataset=cast(Any, train_dataset), data_collator=data_collator, callbacks=[ @@ -463,6 +464,7 @@ def __init__( model: Union[PreTrainedModel, Module], args: Seq2SeqTrainingArguments, data_collator: Any, + processing_class: Optional[PreTrainedTokenizerBase] = None, train_dataset: Optional[Dataset] = None, callbacks: Optional[List[TrainerCallback]] = None, ): @@ -471,6 +473,7 @@ def __init__( args=args, data_collator=data_collator, train_dataset=train_dataset, + processing_class=processing_class, callbacks=callbacks, ) diff --git a/machine/translation/huggingface/transformers_compatibility.py b/machine/translation/huggingface/transformers_compatibility.py index 34438e4d..8f81c3b6 100644 --- a/machine/translation/huggingface/transformers_compatibility.py +++ b/machine/translation/huggingface/transformers_compatibility.py @@ -111,3 +111,34 @@ def preprocess( return build_inputs_fn( *args, return_tensors="pt", truncation=truncation, src_lang=src_lang, tgt_lang=tgt_lang ) + else: + return self._parse_and_tokenize(*args, truncation=truncation) + + def _parse_and_tokenize(self, *args, truncation): + prefix = self.prefix if self.prefix is not None else "" + + tokenizer = self.tokenizer + if tokenizer is None: + raise ValueError("A tokenizer is required for translation") + + if isinstance(args[0], list): + if tokenizer.pad_token_id is None: + raise ValueError("Please make sure that the tokenizer has a pad_token_id when using a batch input") + args = ([prefix + arg for arg in args[0]],) + padding = True + + elif isinstance(args[0], str): + args = (prefix + args[0],) + padding = False + else: + raise TypeError( + f" `args[0]`: {args[0]} have the wrong format. The should be either of type `str` or type `list`" + ) + + inputs = tokenizer(*args, padding=padding, truncation=truncation, return_tensors="pt") + + # This is produced by tokenizers but is an invalid generate kwargs + if "token_type_ids" in inputs: + del inputs["token_type_ids"] + + return inputs From 130bd27a103d700130579876b04c4bf39de1c516 Mon Sep 17 00:00:00 2001 From: Peter Chapman Date: Wed, 30 Sep 2026 10:33:32 +1300 Subject: [PATCH 10/14] Further fixes for issues raised by Claude --- .../translation/huggingface/hugging_face_nmt_engine.py | 6 +++++- .../huggingface/transformers_compatibility.py | 10 ++++------ 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/machine/translation/huggingface/hugging_face_nmt_engine.py b/machine/translation/huggingface/hugging_face_nmt_engine.py index 1937a897..9836ccc6 100644 --- a/machine/translation/huggingface/hugging_face_nmt_engine.py +++ b/machine/translation/huggingface/hugging_face_nmt_engine.py @@ -47,6 +47,7 @@ def __init__( self._pipeline_kwargs["output_attentions"] = True if isinstance(model, PreTrainedModel): self._model = model + self._model_attn_implementation = model.config._attn_implementation if self._pipeline_kwargs["output_attentions"]: model.set_attn_implementation("eager") self._model.eval() @@ -196,6 +197,9 @@ def close(self) -> None: del self._pipeline if self._is_model_owned: del self._model + elif self._model_attn_implementation is not None: + # Restore the attn implementation to the model + self._model.set_attn_implementation(self._model_attn_implementation) gc.collect() with torch.no_grad(): torch.cuda.empty_cache() @@ -313,7 +317,7 @@ def _forward(self, model_inputs, **generate_kwargs): start_index = 0 if self.model.config.decoder_start_token_id is not None: start_index = 1 - if self.generation_config.output_attentions is not False: + if self.generation_config.output_attentions is not False or generate_kwargs.get("output_attentions"): assert attentions is not None num_heads = attentions[0][0].shape[1] diff --git a/machine/translation/huggingface/transformers_compatibility.py b/machine/translation/huggingface/transformers_compatibility.py index 8f81c3b6..285e074f 100644 --- a/machine/translation/huggingface/transformers_compatibility.py +++ b/machine/translation/huggingface/transformers_compatibility.py @@ -1,4 +1,3 @@ -import enum import warnings from typing import Any, Callable, Sequence, Union, cast @@ -8,11 +7,6 @@ # The following classes are a port of the same classes found in transformers v4 -class ReturnType(enum.Enum): - TENSORS = 0 - TEXT = 1 - - class TranslationPipeline(Pipeline): _pipeline_calls_generate = True @@ -34,8 +28,12 @@ def _sanitize_parameters( clean_up_tokenization_spaces=None, truncation=None, stop_sequence=None, + prefix=None, **generate_kwargs, ): + if prefix is not None: + self.prefix = prefix + preprocess_params = {} if truncation is not None: preprocess_params["truncation"] = truncation From e77bcc909093cf638ea9b61be207f9443b296a78 Mon Sep 17 00:00:00 2001 From: Damien Daspit Date: Wed, 30 Sep 2026 11:51:11 -0400 Subject: [PATCH 11/14] Fix T5 loading, attention handling, and group_by_length mapping Let transformers choose the attention implementation when attentions are not needed, since forcing "sdpa" fails to load T5/mT5. Decide whether to build alignments from the attentions generate returns, so SilTranslationPipeline no longer crashes when output_attentions is unset or overridden per call. Ignore a null group_by_length, let an explicit train_sampling_strategy take precedence, and log the deprecation so it appears in job logs. Switch a caller-owned model to eager attention only at the end of the engine constructor, so a constructor that raises leaves the model unchanged. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../hugging_face_nmt_model_factory.py | 19 +++---- .../huggingface/hugging_face_nmt_engine.py | 19 ++++--- .../test_hugging_face_nmt_model_factory.py | 46 ++++++++++++++++ .../test_hugging_face_nmt_engine.py | 53 ++++++++++++++++++- 4 files changed, 118 insertions(+), 19 deletions(-) create mode 100644 tests/jobs/test_hugging_face_nmt_model_factory.py diff --git a/machine/jobs/huggingface/hugging_face_nmt_model_factory.py b/machine/jobs/huggingface/hugging_face_nmt_model_factory.py index 5aedc7b5..a680a8e1 100644 --- a/machine/jobs/huggingface/hugging_face_nmt_model_factory.py +++ b/machine/jobs/huggingface/hugging_face_nmt_model_factory.py @@ -1,6 +1,5 @@ import logging import tarfile -import warnings from pathlib import Path from typing import Any, cast @@ -27,17 +26,13 @@ def __init__(self, config: Any) -> None: self._config = config args = config.huggingface.train_params.to_dict() args["output_dir"] = str(self._model_dir) - # Allow group_by_length backwards compatibility - if "group_by_length" in args: - warnings.warn( - "'group_by_length' is deprecated and will be removed in a future release.", - category=DeprecationWarning, - stacklevel=2, - ) - if args.pop("group_by_length") is True: - args["train_sampling_strategy"] = "group_by_length" - else: - args["train_sampling_strategy"] = "random" + # Allow group_by_length backwards compatibility. The settings default for train_sampling_strategy is + # group_by_length, so any other value was set explicitly and takes precedence over the legacy option. + group_by_length = args.pop("group_by_length", None) + if group_by_length is not None: + logger.warning("'group_by_length' is deprecated. Use 'train_sampling_strategy' instead.") + if args.get("train_sampling_strategy", "group_by_length") == "group_by_length": + args["train_sampling_strategy"] = "group_by_length" if group_by_length else "random" # Use "max_steps" from root for backward compatibility if "max_steps" in self._config.huggingface: args["max_steps"] = self._config.huggingface.max_steps diff --git a/machine/translation/huggingface/hugging_face_nmt_engine.py b/machine/translation/huggingface/hugging_face_nmt_engine.py index 9836ccc6..1e9e8992 100644 --- a/machine/translation/huggingface/hugging_face_nmt_engine.py +++ b/machine/translation/huggingface/hugging_face_nmt_engine.py @@ -42,21 +42,24 @@ def __init__( oom_batch_size_backoff_mult: float = 1.0, **pipeline_kwargs, ) -> None: + """ + A model passed as a `PreTrainedModel` stays owned by the caller. When `output_attentions` is enabled, the + engine switches that model to eager attention and restores its original implementation in `close()`. + """ self._pipeline_kwargs = pipeline_kwargs if self._pipeline_kwargs.get("output_attentions") is None: self._pipeline_kwargs["output_attentions"] = True if isinstance(model, PreTrainedModel): self._model = model self._model_attn_implementation = model.config._attn_implementation - if self._pipeline_kwargs["output_attentions"]: - model.set_attn_implementation("eager") self._model.eval() self._is_model_owned = False else: model_config = AutoConfig.from_pretrained(str(model), label2id={}, id2label={}, num_labels=0) - # If output_attentions is True or None, we need to set the attn_implementation to eager to get the attentions - attn_implementation = "eager" if self._pipeline_kwargs["output_attentions"] else "sdpa" + # Only eager attention returns attention weights. Otherwise, let transformers choose, since forcing "sdpa" + # fails to load architectures without SDPA support (e.g. T5). + attn_implementation = "eager" if self._pipeline_kwargs["output_attentions"] else None self._model = cast( PreTrainedModel, @@ -120,6 +123,10 @@ def __init__( **self._pipeline_kwargs, ) + # Last, so that a constructor that raises leaves the caller's model unchanged. + if not self._is_model_owned and self._pipeline_kwargs["output_attentions"]: + self._model.set_attn_implementation("eager") + @property def tokenizer(self) -> PreTrainedTokenizer | PreTrainedTokenizerFast: return self._tokenizer @@ -317,8 +324,8 @@ def _forward(self, model_inputs, **generate_kwargs): start_index = 0 if self.model.config.decoder_start_token_id is not None: start_index = 1 - if self.generation_config.output_attentions is not False or generate_kwargs.get("output_attentions"): - assert attentions is not None + # output_attentions can be unset or overridden per call, so rely on what generate actually returned. + if attentions is not None: num_heads = attentions[0][0].shape[1] # Truncate/Pad beam_indices to match output_ids length exact slice diff --git a/tests/jobs/test_hugging_face_nmt_model_factory.py b/tests/jobs/test_hugging_face_nmt_model_factory.py new file mode 100644 index 00000000..40753c2f --- /dev/null +++ b/tests/jobs/test_hugging_face_nmt_model_factory.py @@ -0,0 +1,46 @@ +import sys + +if sys.platform == "darwin": + from pytest import skip + + skip("skipping Hugging Face tests on MacOS", allow_module_level=True) + +from typing import Optional + +from pytest import mark +from testutils.mock_settings import MockSettings + +from machine.jobs.huggingface.hugging_face_nmt_model_factory import HuggingFaceNmtModelFactory + + +@mark.parametrize( + "group_by_length, train_sampling_strategy, expected", + [ + (None, "group_by_length", "group_by_length"), + (True, "group_by_length", "group_by_length"), + (False, "group_by_length", "random"), + (False, "sequential", "sequential"), + (True, "random", "random"), + ], +) +def test_group_by_length_backwards_compatibility( + group_by_length: Optional[bool], train_sampling_strategy: str, expected: str +) -> None: + config = MockSettings( + { + "data_dir": "/tmp", + "shared_file_folder": "test", + "build_id": "build1", + "clearml": False, + "huggingface": { + "train_params": { + "group_by_length": group_by_length, + "train_sampling_strategy": train_sampling_strategy, + } + }, + } + ) + + factory = HuggingFaceNmtModelFactory(config) + + assert factory._training_args.train_sampling_strategy == expected diff --git a/tests/translation/huggingface/test_hugging_face_nmt_engine.py b/tests/translation/huggingface/test_hugging_face_nmt_engine.py index a4b7c5c9..26ab81e5 100644 --- a/tests/translation/huggingface/test_hugging_face_nmt_engine.py +++ b/tests/translation/huggingface/test_hugging_face_nmt_engine.py @@ -6,10 +6,12 @@ skip("skipping Hugging Face tests on MacOS", allow_module_level=True) from math import exp, log +from typing import Any, Dict, Optional from pytest import approx, mark, raises +from transformers import AutoModelForSeq2SeqLM, AutoTokenizer -from machine.translation.huggingface import HuggingFaceNmtEngine +from machine.translation.huggingface import HuggingFaceNmtEngine, SilTranslationPipeline from machine.translation.translation_result import TranslationResult @@ -73,6 +75,55 @@ def test_construct_invalid_lang(output_attentions: bool) -> None: HuggingFaceNmtEngine("stas/tiny-m2m_100", src_lang="qaa", tgt_lang="es", output_attentions=output_attentions) +def test_construct_invalid_lang_leaves_caller_model_unchanged() -> None: + model = AutoModelForSeq2SeqLM.from_pretrained("stas/tiny-m2m_100", attn_implementation="sdpa") + with raises(ValueError): + HuggingFaceNmtEngine(model, src_lang="en", tgt_lang="qaa", output_attentions=True) + assert model.config._attn_implementation == "sdpa" + + +def test_close_restores_caller_model_attn_implementation() -> None: + model = AutoModelForSeq2SeqLM.from_pretrained("stas/tiny-m2m_100", attn_implementation="sdpa") + with HuggingFaceNmtEngine(model, src_lang="en", tgt_lang="es", max_length=10, output_attentions=True) as engine: + assert model.config._attn_implementation == "eager" + assert str(engine.translate("This is a test string").alignment) != "" + assert model.config._attn_implementation == "sdpa" + + +@mark.parametrize("output_attentions", [True, False]) +def test_translate_model_without_sdpa_support(output_attentions: bool) -> None: + with HuggingFaceNmtEngine( + "hf-internal-testing/tiny-random-T5ForConditionalGeneration", + max_length=10, + device="cpu", + output_attentions=output_attentions, + ) as engine: + result = engine.translate("This is a test string") + assert result.source_tokens == ["▁This", "▁is", "▁", "a", "▁test", "▁string"] + + +@mark.parametrize( + "construct_output_attentions, call_output_attentions, expect_attentions", + [(None, None, False), (True, False, False), (False, True, True)], +) +def test_pipeline_output_attentions( + construct_output_attentions: Optional[bool], call_output_attentions: Optional[bool], expect_attentions: bool +) -> None: + model = AutoModelForSeq2SeqLM.from_pretrained("stas/tiny-m2m_100", attn_implementation="eager") + tokenizer = AutoTokenizer.from_pretrained("stas/tiny-m2m_100") + construct_kwargs: Dict[str, Any] = ( + {} if construct_output_attentions is None else {"output_attentions": construct_output_attentions} + ) + call_kwargs = {} if call_output_attentions is None else {"output_attentions": call_output_attentions} + pipeline = SilTranslationPipeline( + model=model, tokenizer=tokenizer, batch_size=1, src_lang="en", tgt_lang="es", max_length=10, **construct_kwargs + ) + + outputs = pipeline(["This is a test string"], **call_kwargs) + + assert ("token_attentions" in outputs[0]) == expect_attentions + + def _get_sequence_confidence(result: TranslationResult) -> float: # Inject a 0 score for the BOS token return exp(sum([log(c) for c in result.confidences] + [0]) / (len(result.confidences) + 1)) From 5ec388be6591f9455b723c2893e55501c53af03b Mon Sep 17 00:00:00 2001 From: Peter Chapman Date: Thu, 1 Oct 2026 07:52:55 +1300 Subject: [PATCH 12/14] Relax datasets version to a range --- poetry.lock | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/poetry.lock b/poetry.lock index a119e1a2..ac444522 100644 --- a/poetry.lock +++ b/poetry.lock @@ -6217,4 +6217,4 @@ thot = ["sil-thot"] [metadata] lock-version = "2.1" python-versions = ">=3.10,<3.15" -content-hash = "08e0ea2d4df0e5b81730dffe129b0c03300575c129f46e304c2fdfc3a6d2aa49" +content-hash = "79f394b1c95e82dbd1aa7a8fc7b2c82293060ea46a10f3cbca2dd204c3e2fb19" diff --git a/pyproject.toml b/pyproject.toml index b4a7c721..8671a08d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,7 +73,7 @@ accelerate = { version = "^1.14", markers = "sys_platform == 'win32' or sys_plat [project.optional-dependencies] sentencepiece = ["sentencepiece~=0.2"] thot = ["sil-thot~=3.5.2"] -huggingface = ["transformers==5.14.1", "datasets==5.0.1", "sacremoses>=0.0.53,<1"] +huggingface = ["transformers==5.14.1", "datasets>=5.0.1,<6", "sacremoses>=0.0.53,<1"] jobs = [ "clearml[s3]>=1.13.1,<2.0", "json-stream~=1.3", "dynaconf>=3.2.5,<4.0", "eflomal~=2.0; sys_platform == 'linux'", "pydantic>=2.13.4,<3.0"] [build-system] From 915ca5aa3d1d1b8bd369e0c38c5349f17945d3c4 Mon Sep 17 00:00:00 2001 From: Damien Daspit Date: Wed, 30 Sep 2026 13:57:50 -0400 Subject: [PATCH 13/14] Require eager attention for output_attentions on a caller-owned model A caller-owned model was switched to eager attention for the life of the engine, so closing one engine broke another engine that shared the model. The engine no longer modifies a model it does not own. Instead, the pipeline raises a ValueError when output_attentions is enabled, at construction or per call, and the model does not use eager attention. Models the engine loads itself are still loaded with eager attention. Add tests that check greedy alignments against the cross-attention directly and _compute_transition_scores against compute_transition_scores. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../huggingface/hugging_face_nmt_engine.py | 33 +++--- .../test_hugging_face_nmt_engine.py | 100 ++++++++++++++++-- 2 files changed, 113 insertions(+), 20 deletions(-) diff --git a/machine/translation/huggingface/hugging_face_nmt_engine.py b/machine/translation/huggingface/hugging_face_nmt_engine.py index 1e9e8992..80219d6a 100644 --- a/machine/translation/huggingface/hugging_face_nmt_engine.py +++ b/machine/translation/huggingface/hugging_face_nmt_engine.py @@ -43,15 +43,14 @@ def __init__( **pipeline_kwargs, ) -> None: """ - A model passed as a `PreTrainedModel` stays owned by the caller. When `output_attentions` is enabled, the - engine switches that model to eager attention and restores its original implementation in `close()`. + A model passed as a `PreTrainedModel` stays owned by the caller and is not modified. `output_attentions`, + which defaults to True, requires that model to use eager attention. Otherwise, a `ValueError` is raised. """ self._pipeline_kwargs = pipeline_kwargs if self._pipeline_kwargs.get("output_attentions") is None: self._pipeline_kwargs["output_attentions"] = True if isinstance(model, PreTrainedModel): self._model = model - self._model_attn_implementation = model.config._attn_implementation self._model.eval() self._is_model_owned = False else: @@ -123,10 +122,6 @@ def __init__( **self._pipeline_kwargs, ) - # Last, so that a constructor that raises leaves the caller's model unchanged. - if not self._is_model_owned and self._pipeline_kwargs["output_attentions"]: - self._model.set_attn_implementation("eager") - @property def tokenizer(self) -> PreTrainedTokenizer | PreTrainedTokenizerFast: return self._tokenizer @@ -204,9 +199,6 @@ def close(self) -> None: del self._pipeline if self._is_model_owned: del self._model - elif self._model_attn_implementation is not None: - # Restore the attn implementation to the model - self._model.set_attn_implementation(self._model_attn_implementation) gc.collect() with torch.no_grad(): torch.cuda.empty_cache() @@ -222,6 +214,8 @@ def __init__( **kwargs, ) -> None: super().__init__(model=model, tokenizer=tokenizer, batch_size=batch_size, **kwargs) + if self.generation_config.output_attentions: + _check_eager_attention(self.model) self._mpn = mpn def preprocess(self, *args, truncation=TruncationStrategy.DO_NOT_TRUNCATE, src_lang=None, tgt_lang=None): @@ -266,6 +260,9 @@ def _forward(self, model_inputs, **generate_kwargs): input_tokens = model_inputs.pop("input_tokens") self.check_inputs(input_length, self.generation_config.min_length, self.generation_config.max_length) + if generate_kwargs.get("output_attentions", self.generation_config.output_attentions): + _check_eager_attention(self.model) + output = cast(Any, self.model).generate( **model_inputs, **generate_kwargs, @@ -324,8 +321,10 @@ def _forward(self, model_inputs, **generate_kwargs): start_index = 0 if self.model.config.decoder_start_token_id is not None: start_index = 1 - # output_attentions can be unset or overridden per call, so rely on what generate actually returned. - if attentions is not None: + # Rely on what generate returned, since output_attentions can be unset or overridden per call. + if not attentions or not attentions[0]: + attentions = None + else: num_heads = attentions[0][0].shape[1] # Truncate/Pad beam_indices to match output_ids length exact slice @@ -510,5 +509,15 @@ def _compute_transition_scores( return transition_scores +def _check_eager_attention(model: PreTrainedModel) -> None: + # Other attention implementations do not return attention weights, so alignments would silently be empty. + attn_implementation = model.config._attn_implementation + if attn_implementation != "eager": + raise ValueError( + f"output_attentions requires a model that uses eager attention, but the model uses " + f"'{attn_implementation}'. Load the model with attn_implementation='eager' or set output_attentions=False." + ) + + def _get_encoding_fast_tokens(encoding) -> List[str]: return [token for (token, mask) in zip(encoding.tokens, encoding.special_tokens_mask) if not mask] diff --git a/tests/translation/huggingface/test_hugging_face_nmt_engine.py b/tests/translation/huggingface/test_hugging_face_nmt_engine.py index 26ab81e5..862c755b 100644 --- a/tests/translation/huggingface/test_hugging_face_nmt_engine.py +++ b/tests/translation/huggingface/test_hugging_face_nmt_engine.py @@ -8,11 +8,14 @@ from math import exp, log from typing import Any, Dict, Optional +import torch from pytest import approx, mark, raises from transformers import AutoModelForSeq2SeqLM, AutoTokenizer from machine.translation.huggingface import HuggingFaceNmtEngine, SilTranslationPipeline +from machine.translation.huggingface.hugging_face_nmt_engine import _compute_transition_scores from machine.translation.translation_result import TranslationResult +from machine.translation.word_alignment_matrix import WordAlignmentMatrix @mark.parametrize("output_attentions", [True, False]) @@ -75,21 +78,89 @@ def test_construct_invalid_lang(output_attentions: bool) -> None: HuggingFaceNmtEngine("stas/tiny-m2m_100", src_lang="qaa", tgt_lang="es", output_attentions=output_attentions) -def test_construct_invalid_lang_leaves_caller_model_unchanged() -> None: +@mark.parametrize("output_attentions", [True, None]) +def test_construct_output_attentions_non_eager_caller_model(output_attentions: Optional[bool]) -> None: model = AutoModelForSeq2SeqLM.from_pretrained("stas/tiny-m2m_100", attn_implementation="sdpa") - with raises(ValueError): - HuggingFaceNmtEngine(model, src_lang="en", tgt_lang="qaa", output_attentions=True) + with raises(ValueError, match="eager attention"): + HuggingFaceNmtEngine(model, src_lang="en", tgt_lang="es", output_attentions=output_attentions) assert model.config._attn_implementation == "sdpa" -def test_close_restores_caller_model_attn_implementation() -> None: +def test_translate_non_eager_caller_model_without_output_attentions() -> None: model = AutoModelForSeq2SeqLM.from_pretrained("stas/tiny-m2m_100", attn_implementation="sdpa") - with HuggingFaceNmtEngine(model, src_lang="en", tgt_lang="es", max_length=10, output_attentions=True) as engine: - assert model.config._attn_implementation == "eager" - assert str(engine.translate("This is a test string").alignment) != "" + with HuggingFaceNmtEngine(model, src_lang="en", tgt_lang="es", max_length=10, output_attentions=False) as engine: + result = engine.translate("This is a test string") + assert result.translation == "skaberskaber Dollar Dollar Dollar ፤ gerekir gerekir" + assert str(result.alignment) == "" assert model.config._attn_implementation == "sdpa" +def test_engines_share_caller_model() -> None: + model = AutoModelForSeq2SeqLM.from_pretrained("stas/tiny-m2m_100", attn_implementation="eager") + engine1 = HuggingFaceNmtEngine(model, src_lang="en", tgt_lang="es", max_length=10, output_attentions=True) + with HuggingFaceNmtEngine(model, src_lang="en", tgt_lang="es", max_length=10, output_attentions=True) as engine2: + engine1.close() + assert str(engine2.translate("This is a test string").alignment) == "0-2 2-0 2-1 2-3 4-4 4-5 4-6 4-7" + + +def test_translate_greedy_alignment_matches_cross_attention() -> None: + model = AutoModelForSeq2SeqLM.from_pretrained("stas/tiny-m2m_100", attn_implementation="eager") + tokenizer = AutoTokenizer.from_pretrained("stas/tiny-m2m_100", src_lang="en", tgt_lang="es") + with HuggingFaceNmtEngine(model, src_lang="en", tgt_lang="es", max_length=10, device="cpu") as engine: + result = engine.translate("This is a test string") + + inputs = tokenizer("This is a test string", return_tensors="pt") + output = model.generate( + **inputs, + forced_bos_token_id=tokenizer.get_lang_id("es"), + max_length=10, + output_attentions=True, + return_dict_in_generate=True, + ) + assert tokenizer.decode(output.sequences[0], skip_special_tokens=True) == result.translation + + # The engine aligns each target token to the source token with the most cross-attention, averaged over the + # heads of the layer two thirds of the way through the decoder. + special_ids = set(tokenizer.all_special_ids) + src_indices = [i for i, id in enumerate(inputs["input_ids"][0].tolist()) if id not in special_ids] + layer = (2 * len(output.cross_attentions[0])) // 3 + word_pairs = [] + for i, id in enumerate(output.sequences[0].tolist()): + if id in special_ids: + continue + # The sequence starts with the decoder start token, so token i is generated at step i - 1. + attention = output.cross_attentions[i - 1][layer][0].mean(dim=0)[-1, src_indices] + word_pairs.append((int(attention.argmax()), len(word_pairs))) + expected = WordAlignmentMatrix.from_word_pairs(len(src_indices), len(word_pairs), word_pairs) + assert str(result.alignment) == str(expected) + + +@mark.parametrize("num_beams", [1, 2]) +def test_compute_transition_scores_matches_transformers(num_beams: int) -> None: + model = AutoModelForSeq2SeqLM.from_pretrained("stas/tiny-m2m_100") + tokenizer = AutoTokenizer.from_pretrained("stas/tiny-m2m_100", src_lang="en", tgt_lang="es") + inputs = tokenizer(["This is a test string", "Hello, world!"], return_tensors="pt", padding=True) + output = model.generate( + **inputs, + forced_bos_token_id=tokenizer.get_lang_id("es"), + max_length=10, + num_beams=num_beams, + num_return_sequences=num_beams, + output_scores=True, + return_dict_in_generate=True, + ) + beam_indices = output.beam_indices if num_beams > 1 else None + # Beam search scores are already log probabilities, and greedy search scores are logits. + normalize_logits = num_beams == 1 + + expected = model.compute_transition_scores( + output.sequences, output.scores, beam_indices, normalize_logits=normalize_logits + ) + actual = _compute_transition_scores(output.sequences, output.scores, beam_indices, normalize_logits) + + assert torch.allclose(actual, expected) + + @mark.parametrize("output_attentions", [True, False]) def test_translate_model_without_sdpa_support(output_attentions: bool) -> None: with HuggingFaceNmtEngine( @@ -104,7 +175,7 @@ def test_translate_model_without_sdpa_support(output_attentions: bool) -> None: @mark.parametrize( "construct_output_attentions, call_output_attentions, expect_attentions", - [(None, None, False), (True, False, False), (False, True, True)], + [(None, None, False), (True, None, True), (True, False, False), (False, True, True)], ) def test_pipeline_output_attentions( construct_output_attentions: Optional[bool], call_output_attentions: Optional[bool], expect_attentions: bool @@ -124,6 +195,19 @@ def test_pipeline_output_attentions( assert ("token_attentions" in outputs[0]) == expect_attentions +def test_pipeline_output_attentions_non_eager_model() -> None: + model = AutoModelForSeq2SeqLM.from_pretrained("stas/tiny-m2m_100", attn_implementation="sdpa") + tokenizer = AutoTokenizer.from_pretrained("stas/tiny-m2m_100") + with raises(ValueError, match="eager attention"): + SilTranslationPipeline(model=model, tokenizer=tokenizer, batch_size=1, output_attentions=True) + + pipeline = SilTranslationPipeline( + model=model, tokenizer=tokenizer, batch_size=1, src_lang="en", tgt_lang="es", max_length=10 + ) + with raises(ValueError, match="eager attention"): + pipeline(["This is a test string"], output_attentions=True) + + def _get_sequence_confidence(result: TranslationResult) -> float: # Inject a 0 score for the BOS token return exp(sum([log(c) for c in result.confidences] + [0]) / (len(result.confidences) + 1)) From 6d4cd7f19d11ab61f5f0823b2c2a44c0fa813678 Mon Sep 17 00:00:00 2001 From: Damien Daspit Date: Wed, 30 Sep 2026 15:48:39 -0400 Subject: [PATCH 14/14] Avoid importing torch in the engine tests Pyright cannot resolve torch on macOS with Python 3.14, which failed the CI build. The transition score test now compares tensors with the Tensor.allclose method instead. Co-Authored-By: Claude Opus 5.5 (1M context) --- tests/translation/huggingface/test_hugging_face_nmt_engine.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/translation/huggingface/test_hugging_face_nmt_engine.py b/tests/translation/huggingface/test_hugging_face_nmt_engine.py index 862c755b..d1d34be4 100644 --- a/tests/translation/huggingface/test_hugging_face_nmt_engine.py +++ b/tests/translation/huggingface/test_hugging_face_nmt_engine.py @@ -8,7 +8,6 @@ from math import exp, log from typing import Any, Dict, Optional -import torch from pytest import approx, mark, raises from transformers import AutoModelForSeq2SeqLM, AutoTokenizer @@ -158,7 +157,7 @@ def test_compute_transition_scores_matches_transformers(num_beams: int) -> None: ) actual = _compute_transition_scores(output.sequences, output.scores, beam_indices, normalize_logits) - assert torch.allclose(actual, expected) + assert actual.allclose(expected) @mark.parametrize("output_attentions", [True, False])