Ng,ddlmZddlZddlZddlmZddlmZddlm Z ddl m Z m Z m Z ddlZddlmZddlmZdd lmZmZmZmZdd lmZmZmZmZdd lmZdd lm Z dd l!m"Z"ddl#m$Z$ddl%m&Z&ddl'm(Z(ddl)m*Z*m+Z+ddl,m-Z-ddl.m/Z/ddl0m1Z1ddl2m3Z3ddl4m5Z5m6Z6m7Z7m8Z8m9Z9ddl:m;Z;mm?Z?m@Z@mAZAe@rddlBmCZCmDZDmEZEmFZFmGZGejHeIZJe rddlKmLZLGddeZMdS)) annotationsN) OrderedDict) nullcontext)partial) TYPE_CHECKINGAnyCallable)parse)nn) BatchSampler ConcatDataset DataLoaderSubsetRandomSampler)EvalPredictionPreTrainedTokenizerBaseTrainerTrainerCallback) __version__) DataCollator) WandbCallback)TRAINING_ARGS_NAME)EvalLoopOutput)SentenceTransformerDataCollator)SentenceEvaluatorSequentialEvaluator) CoSENTLoss)ModelCardCallback)Pooling) Transformer)DefaultBatchSamplerGroupByLabelBatchSamplerNoDuplicatesBatchSamplerProportionalBatchSamplerRoundRobinBatchSampler) BatchSamplersMultiDatasetBatchSamplers$SentenceTransformerTrainingArguments)disable_loggingis_datasets_availableis_training_available)Dataset DatasetDictIterableDatasetIterableDatasetDictValueSentenceTransformerceZdZdZ ddfd Zdd"Zddfd$ Zdd&Zdd(Zdd+Z ddd2Z dd4Z ddfd< Z ddfdC Z dfdD ZdddIZ dddQZ dddYZddZZddd\Zdd_ZdddaZddcZddeZddfZe dddoZ dddrZ dddvZ dddZ ddfd ZxZS)SentenceTransformerTraineru SentenceTransformerTrainer is a simple but feature-complete training and eval loop for PyTorch based on the 🤗 Transformers :class:`~transformers.Trainer`. This trainer integrates support for various :class:`transformers.TrainerCallback` subclasses, such as: - :class:`~transformers.integrations.WandbCallback` to automatically log training metrics to W&B if `wandb` is installed - :class:`~transformers.integrations.TensorBoardCallback` to log training metrics to TensorBoard if `tensorboard` is accessible. - :class:`~transformers.integrations.CodeCarbonCallback` to track the carbon emissions of your model during training if `codecarbon` is installed. - Note: These carbon emissions will be included in your automatically generated model card. See the Transformers `Callbacks `_ documentation for more information on the integrated callbacks and how to write your own callbacks. Args: model (:class:`~sentence_transformers.SentenceTransformer`, *optional*): The model to train, evaluate or use for predictions. If not provided, a `model_init` must be passed. args (:class:`~sentence_transformers.training_args.SentenceTransformerTrainingArguments`, *optional*): The arguments to tweak for training. Will default to a basic instance of :class:`~sentence_transformers.training_args.SentenceTransformerTrainingArguments` with the `output_dir` set to a directory named *tmp_trainer* in the current directory if not provided. train_dataset (Union[:class:`datasets.Dataset`, :class:`datasets.DatasetDict`, :class:`datasets.IterableDataset`, Dict[str, :class:`datasets.Dataset`]], *optional*): The dataset to use for training. Must have a format accepted by your loss function, see `Training Overview > Dataset Format <../../../docs/sentence_transformer/training_overview.html#dataset-format>`_. eval_dataset (Union[:class:`datasets.Dataset`, :class:`datasets.DatasetDict`, :class:`datasets.IterableDataset`, Dict[str, :class:`datasets.Dataset`]], *optional*): The dataset to use for evaluation. Must have a format accepted by your loss function, see `Training Overview > Dataset Format <../../../docs/sentence_transformer/training_overview.html#dataset-format>`_. loss (Optional[Union[:class:`torch.nn.Module`, Dict[str, :class:`torch.nn.Module`], Callable[[:class:`~sentence_transformers.SentenceTransformer`], :class:`torch.nn.Module`], Dict[str, Callable[[:class:`~sentence_transformers.SentenceTransformer`]]]], *optional*): The loss function to use for training. Can either be a loss class instance, a dictionary mapping dataset names to loss class instances, a function that returns a loss class instance given a model, or a dictionary mapping dataset names to functions that return a loss class instance given a model. In practice, the latter two are primarily used for hyper-parameter optimization. Will default to :class:`~sentence_transformers.losses.CoSENTLoss` if no ``loss`` is provided. evaluator (Union[:class:`~sentence_transformers.evaluation.SentenceEvaluator`, List[:class:`~sentence_transformers.evaluation.SentenceEvaluator`]], *optional*): The evaluator instance for useful evaluation metrics during training. You can use an ``evaluator`` with or without an ``eval_dataset``, and vice versa. Generally, the metrics that an ``evaluator`` returns are more useful than the loss value returned from the ``eval_dataset``. A list of evaluators will be wrapped in a :class:`~sentence_transformers.evaluation.SequentialEvaluator` to run them sequentially. callbacks (List of [:class:`transformers.TrainerCallback`], *optional*): A list of callbacks to customize the training loop. Will add those to the list of default callbacks detailed in [here](callback). If you want to remove one of the default callbacks used, use the [`Trainer.remove_callback`] method. optimizers (`Tuple[:class:`torch.optim.Optimizer`, :class:`torch.optim.lr_scheduler.LambdaLR`]`, *optional*, defaults to `(None, None)`): A tuple containing the optimizer and the scheduler to use. Will default to an instance of :class:`torch.optim.AdamW` on your model and a scheduler given by :func:`transformers.get_linear_schedule_with_warmup` controlled by `args`. Important attributes: - **model** -- Always points to the core model. If using a transformers model, it will be a [`PreTrainedModel`] subclass. - **model_wrapped** -- Always points to the most external model in case one or more other modules wrap the original model. This is the model that should be used for the forward pass. For example, under `DeepSpeed`, the inner model is wrapped in `DeepSpeed` and then again in `torch.nn.DistributedDataParallel`. If the inner model hasn't been wrapped, then `self.model_wrapped` is the same as `self.model`. - **is_model_parallel** -- Whether or not a model has been switched to a model parallel mode (different from data parallelism, this means some of the model layers are split on different GPUs). - **place_model_on_device** -- Whether or not to automatically place the model on the device - it will be set to `False` if model parallel or deepspeed is used, or if the default `TrainingArguments.place_model_on_device` is overridden to return `False` . - **is_in_train** -- Whether or not a model is currently running `train` (e.g. when `evaluate` is called while in `train`) NNNmodelSentenceTransformer | Noneargsr' train_datasetCDataset | DatasetDict | IterableDataset | dict[str, Dataset] | None eval_datasetlossnn.Module | dict[str, nn.Module] | Callable[[SentenceTransformer], torch.nn.Module] | dict[str, Callable[[SentenceTransformer], torch.nn.Module]] | None evaluator2SentenceEvaluator | list[SentenceEvaluator] | None data_collatorDataCollator | None tokenizer)PreTrainedTokenizerBase | Callable | None model_init(Callable[[], SentenceTransformer] | Nonecompute_metrics'Callable[[EvalPrediction], dict] | None callbackslist[TrainerCallback] | None optimizers?tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]preprocess_logits_for_metrics;Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | NonereturnNonec  tstd|1d}td|dt |}n$t |tst d-| | _n2td| t d| _| t d t d  }|j r+j j sj |j |!t jt rj}|t#j }t'd d g||gD]\}}t |t(r}|jvt-t/|}t0dt2dt4dt6difd|D}t d|d|d|d|d t |t:r$t |t<st=|}t |t:r$t |t<st=|}jrdn||||||nd| | | | | d }t?t@t?dkr||d<n||d<tCj"d.i|j#dkrd_#d_$i_%tMdj'j(Dr tRj*+dd |.td!tYj-}t |t:rfd"|D_.t'd d g||gD]\}}|t |t:st d#|d$t_|0t_|0z x}r>t d%|d&tc|d'te|d(krd)nd*d+|d, n3|_.|$t |thstk|}|_6j7"8||j9d -_7j#"8||j9d -_#:|dS)/NzTo train a SentenceTransformer model, you need to install the `accelerate` and `datasets` modules. You can do so with the `train` extra: pip install -U "sentence-transformers[train]" tmp_trainerz1No `TrainingArguments` passed, using `output_dir=z`.) output_dirzEPlease use `TrainingArguments` imported from `sentence_transformers`.z<`Trainer` requires either a `model` or `model_init` argumentz`Trainer` requires either a `model` or `model_init` argument, but not both. `model_init` will overwrite your model when calling the `train` method.z`compute_metrics` is currently not compatible with the SentenceTransformerTrainer. Please use the `evaluator` argument instead for detailed evaluation metrics, or the `eval_dataset` argument for the evaluation loss.unused) tokenize_fntrainevalstringint64float32boolc xi|]6\}}|tt|d7S)null)r/gettype).0keyvaluenaive_type_mappings Y/var/www/html/ai-engine/env/lib/python3.11/site-packages/sentence_transformers/trainer.py z7SentenceTransformerTrainer.__init__..sM$$$PZPSUZC155d5kk6JJKK$$$zThe provided `z6_dataset` must have Features. Specify them with e.g.: z _dataset = z_dataset.cast(Features(z)) or by providing the Features to the IterableDataset initialization method. See the Datasets documentation for more information on dataset Features: https://huggingface.co/docs/datasets/en/about_dataset_featuresdummy) r5r7r?r8r:rCrErGrIrK4.46.0processing_classrATc8g|]}t|tS) isinstancer)r^callbacks rb z7SentenceTransformerTrainer.__init__..s"ddd 8]33dddrd WANDB_PROJECTzsentence-transformersz@No `loss` passed, using `losses.CoSENTLoss` as a default option.cDi|]\}}||Sri) prepare_loss)r^ dataset_nameloss_fnr5selfs rbrcz7SentenceTransformerTrainer.__init__..s1sssMb\[bt'8'8%'H'Hsssrdz,If the provided `loss` is a dict, then the `z"_dataset` must be a `DatasetDict`.z:If the provided `loss` is a dict, then all keys from the `z;_dataset` dictionary must occur in `loss` also. Currently, z occursz in `z_dataset` but not in `loss`.rpri);r* RuntimeErrorloggerinfor'rj ValueErrorrCcall_model_initwarningto_dict hub_model_idmodel_card_datamodel_id set_model_idrArrtokenizezipr- column_namesnextiterstrintfloatrYitemsdictr, parse_versiontransformers_versionsuper__init__r:can_return_loss_prompt_length_mappinganycallback_handlerrGosenviron setdefaultrr5r;setkeyssortedlenrorrr=r8(maybe_add_prompts_or_dataset_name_columnpromptsadd_model_card_callback)rrr5r7r8r:r;r=r?rArCrErGrIrKrQdefault_args_dictrpdatasetsampleexample_features super_kwargsmissingra __class__s`` @rbrz#SentenceTransformerTrainer.__init__xsL(%&& @  <&J KKZJZZZ [ [ [7:NNNDDD"FGG fdee e =%",,,.."#abbb%M)DO  & NN'   AHUUU]]__   BU%:%C B  ! . .t/@ A A A  EO=T!U!U I  ;WWWM%('6):]LP\@])^)^  % g?!'400$w|www"',,..11C 4D4DDD7$OUaOO&,WooOODGLLTUDUDUSS[]OOdpOOO ))$66DI  I?P)Q)Q +I66I"   )!%!N!Nt|'"O""D    ( $ M Mdl!N!!D  $$%677777rdrdict[str, Any]ct||}||||j|j|j|jdS)ah Add a callback responsible for automatically tracking data required for the automatic model card generation This method is called in the ``__init__`` method of the :class:`~sentence_transformers.trainer.SentenceTransformerTrainer` class. Args: default_args_dict (Dict[str, Any]): A dictionary of the default training arguments, so we can determine which arguments have been changed for the model card. .. note:: This method can be overriden by subclassing the trainer to remove/customize this callback in custom uses cases N)r add_callback on_init_endr7statecontrolr5)rrrmodel_card_callbacks rbrz2SentenceTransformerTrainer.add_model_card_callbacksR 06GHH -...'' 4:t|TZXXXXXrdr1ct|}t|ds|St|jt r|jD]h\}}t|tjj s|||j|<8t|dr | |j||_intt|jtjj s|||_n5t|jdr | |j||_|S)N)trialr;r5) rr{hasattrrjr;rrtorchr Moduleoverride_model_in_loss)rrrr5r_rqrs rbr{z*SentenceTransformerTrainer.call_model_init.s0''e'44tV$$ L di & & F $  1 1 N N W!'58?;;N%,WU^^DIcNNWg..N $ ; ;DIu M MDI  NDIux77 F %((DIITY ( ( F33DIuEEDI rdtorch.nn.Modulec  ddlm}|D]g\}}|dkrt||r||_#t|t jjr%t||| ||h|S)Nrr0r5) sentence_transformersr1named_childrenrjr5rr rsetattrr)rrr;r5r1namechilds rbrz1SentenceTransformerTrainer.override_model_in_lossGs======..00 O OKD%w:e5H#I#I" E58?33 OdD$?$?u$M$MNNN rdBCallable[[SentenceTransformer], torch.nn.Module] | torch.nn.Modulect|tjjr||jS|||jSN)rjrr rtodevice)rrr;r5s rbroz'SentenceTransformerTrainer.prepare_lossQsJ dEHO , , )775<(( (tE{{~~el+++rd dataset_dictr,c|D]8\}}d|jvr*|d|gt|z||<9|S)Nrp)rr add_columnr)rrrr_rs rbadd_dataset_name_columnz2SentenceTransformerTrainer.add_dataset_name_columnZs`(..00 ] ]LCW%999$+$6$6~usSZ||G[$\$\ S!rdFinputsdict[str, torch.Tensor | Any]return_outputsrY2torch.Tensor | tuple[torch.Tensor, dict[str, Any]]c`|dd}||\}}|j}t|tr |r||}||jkr<||jkr1t|dr!|j|kr|||}|||} |r| ifS| S)a Computes the loss for the SentenceTransformer model. It uses ``self.loss`` to compute the loss, which can be a single loss function or a dictionary of loss functions for different datasets. If the loss is a dictionary, the dataset name is expected to be passed in the inputs under the key "dataset_name". This is done automatically in the ``add_dataset_name_column`` method. Note that even if ``return_outputs = True``, the outputs will be empty, as the SentenceTransformers losses do not return outputs. Args: model (SentenceTransformer): The SentenceTransformer model. inputs (Dict[str, Union[torch.Tensor, Any]]): The input data for the model. return_outputs (bool, optional): Whether to return the outputs along with the loss. Defaults to False. num_items_in_batch (int, optional): The number of items in the batch. Defaults to None. Unused, but required by the transformers Trainer. Returns: Union[torch.Tensor, Tuple[torch.Tensor, Dict[str, Any]]]: The computed loss. If `return_outputs` is True, returns a tuple of loss and outputs. Otherwise, returns only the loss. rpNr5) popcollect_featuresr;rjr model_wrappedr5rr) rrr5rrnum_items_in_batchrpfeatureslabelsrqr;s rb compute_lossz'SentenceTransformerTrainer.compute_loss`s2zz.$77 0088&) gt $ $ , ,l+G T' ' '##))$ &&11'5AAGwx((   8O rd9tuple[list[dict[str, torch.Tensor]], torch.Tensor | None]cg}|D]}|dr|dtd n]|dr|dtd n/|dr|dtd n|fd|D|d d}||fS) aTurn the inputs from the dataloader into the separate model inputs & the labels. Example:: >>> list(inputs.keys()) ['return_loss', 'label', 'sentence_0_input_ids', 'sentence_0_token_type_ids', 'sentence_0_attention_mask', 'sentence_1_input_ids', 'sentence_1_token_type_ids', 'sentence_1_attention_mask'] >>> features, labels = self.collect_features(inputs) >>> len(features) 2 >>> list(features[0].keys()) ['input_ids', 'token_type_ids', 'attention_mask'] >>> list(features[1].keys()) ['input_ids', 'token_type_ids', 'attention_mask'] >>> torch.equal(labels, inputs["label"]) True _input_idsN input_ids_sentence_embeddingsentence_embedding _pixel_values pixel_valuescpi|]2\}}||td|3Sr) startswithr)r^r_r`prefixs rbrcz?SentenceTransformerTrainer.collect_features..sCrrr:3[^[i[ijp[q[qrSV/rrrrdlabel)endswithrappendrr\)rrrrcolumnrrs @rbrz+SentenceTransformerTrainer.collect_featuress* t tF|,,  33{#3#3"3 34!677  <3';#<#<"< <=11  63~#6#6"6 67 OOrrrrrrr s s s sGT**rdrU#Dataset | dict[str, Dataset] | None ignore_keyslist[str] | Nonemetric_key_prefixrdict[str, float]c|r#|||jjd}n|j}t |||S)NrUrv)rr7rr:revaluate)rrr:rrrs rbrz#SentenceTransformerTrainer.evaluates]  -HHdi/fILL ,Lww k;LMMMrd dataloaderr descriptionprediction_loss_only bool | Nonerc>t|||||}|j|S|jrlt |jt rR|dr=|ddt|j dkrd}n|S| rtnttj5||j}dddn #1swxYwYt |t sd|i}t| D]7}||ds||||d|<8|j||S)N)rrrrreval_rrUr=_)revaluation_loopr= is_in_trainrjr:rrlistris_local_process_zerorr(loggingINFOr5rmetricsupdate) rrrrrrroutputevaluator_metricsr_rs rbrz*SentenceTransformerTrainer.evaluation_loops((!#!5#/ )   > !M    4+-C-C-E-E(F(Fq(III$*!! "88:: ][]]]PWP\@]@] ; ; $tz : :  ; ; ; ; ; ; ; ; ; ; ; ; ; ; ;+T22 A!,.? @ )..0011 ] ]C>>%6"9"9"9:: ]BSBWBWX[B\B\!%6">">">">?/000 sDD  D cTt|jdtstddS |jjx}rH|ddd}|jj t|n#t$rYnwxYw|j}|jdj |_ t|j}||_||jd_ S#|j}||_||jd_ wxYw)NrzLCould not load best model, as the model is not a `transformers`-based model.-rs)rjr5rrxryrbest_model_checkpointrsplitrset_best_model_stepr Exception auto_modelr_load_best_model)rr checkpointstep full_modelloaded_auto_modelrs rbrz+SentenceTransformerTrainer._load_best_models$$*Q-55  KKf g g g F !Z==z J!((a004 *>>s4yyIII    D  Z Z]-  977++-- $ #DJ'8DJqM $ $!% #DJ'8DJqM $ 8 8 8 8sAB B#"B#D"D'rr+rp str | Nonec t|tr3|D]\}}|||dSt |jddhzx}r*t d|r|dznddt|ddS) Nrv return_lossrpz/The following column names are invalid in your  ruz dataset: zH. Avoid using these column names, as they are reserved for internal use.)rjrrvalidate_column_namesrrrzr)rrrrpoverlaps rbrz0SentenceTransformerTrainer.validate_column_namess gt $ $ )0 O O% g**7*NNNN F'.//=.2QQ Q7 ZXdBl,QTBTBTjlZZw{}DxExEZZZ   rd batch_sizer drop_lastvalid_label_columns generatortorch.Generator | NoneBatchSampler | Nonec t|tr6|jjtjkrt ddS|jjtjkrt|||||S|jjtj krt||||S|jjtjkr;tttt||||SdS)a Returns the appropriate batch sampler based on the ``batch_sampler`` argument in ``self.args``. This batch sampler class supports ``__len__`` and ``__iter__`` methods, and is used as the ``batch_sampler`` to create the :class:`torch.utils.data.DataLoader`. .. note:: Override this method to provide a custom batch sampler. Args: dataset (Dataset): The dataset to sample from. batch_size (int): Number of samples per batch. drop_last (bool): If True, drop the last incomplete batch if the dataset size is not divisible by the batch size. valid_label_columns (List[str]): List of column names to check for labels. The first column name from ``valid_label_columns`` found in the dataset will be used as the label column. generator (torch.Generator, optional): Optional random number generator for shuffling the indices. BWhen using an IterableDataset, you cannot specify a batch sampler.N)rr r r r )rr r r )r r r )rjr-r7 batch_samplerr% BATCH_SAMPLERrxr| NO_DUPLICATESr"GROUP_BY_LABELr!r rranger)rrrr r r r s rbget_batch_samplerz,SentenceTransformerTrainer.get_batch_samplers 6 g / / y&-*EEEcddd4 9 "m&A A A+%#$7#   9 "m&B B B+%#$7   9 "m&A A A&#E#g,,$7$79MMM%#  B Ardrr batch_samplerslist[BatchSampler]seed int | Noner c|jjtjkrt ||||S|jjtjkrt ||||SdS)a/ Returns the appropriate multi-dataset batch sampler based on the ``multi_dataset_batch_sampler`` argument in ``self.args``. This batch sampler class supports ``__len__`` and ``__iter__`` methods, and is used as the ``batch_sampler`` to create the :class:`torch.utils.data.DataLoader`. .. note:: Override this method to provide a custom multi-dataset batch sampler. Args: dataset (ConcatDataset): The concatenation of all datasets. batch_samplers (List[BatchSampler]): List of batch samplers for each dataset in the concatenated dataset. generator (torch.Generator, optional): Optional random number generator for shuffling the indices. seed (int, optional): Optional seed for the random number generator rrr rN)r7multi_dataset_batch_samplerr& ROUND_ROBINr$ PROPORTIONALr#)rrrrr rs rbget_multi_dataset_batch_samplerz:SentenceTransformerTrainer.get_multi_dataset_batch_samplerKs}* 9 04M4Y Y Y)-#   9 04M4Z Z Z+-#   [ Zrdcjtdj}jtjjjrjjjjjj jj jj d}t|trb|jjjjdjjt$jkrt(dn=t|t,rtdt|t.r|D]&}t|trtd'fd|D}t3|}||jj }||d <n]t|t6r9|jjjjj }||d <ntd d j_j tC|fi|_"j"S)a@ Returns the training [`~torch.utils.data.DataLoader`]. Will use no sampler if `train_dataset` does not implement `__len__`, a random sampler (adapted to distributed training if necessary) otherwise. Subclass and override this method if you want to inject some custom behavior. NzOTraining requires specifying a train_dataset to the SentenceTransformerTrainer. collate_fn num_workers pin_memorypersistent_workersprefetch_factorrrcSentence Transformers is not compatible with IterableDatasetDict. Please use a DatasetDict instead.YSentence Transformers is not compatible with a DatasetDict containing an IterableDataset.c vg|]5}|jjjjj6Sr r r r )rr7train_batch_sizedataloader_drop_lastr r^rr?r rrs rbrlzCSentenceTransformerTrainer.get_train_dataloader..s\   &&#y9"i<(5(I' '   rdrrr-z^Unsupported `train_dataset` type. Use a Dataset, DatasetDict, or IterableDataset for training.F)#r8rzr?r Generatorr7r manual_seeddataloader_num_workersdataloader_pin_memorydataloader_persistent_workersdataloader_prefetch_factorrjr-rr.r/rr%rrxr|r.r,valuesr r!r+rr  accelerator even_batchespreparer_train_dataloader)rrr8dataloader_paramsrrrr?r s` @@rbget_train_dataloaderz/SentenceTransformerTrainer.get_train_dataloaderps   %noo o* * O%% 9> 2  ! !$). 1 1 1(9;)9"&)"I#yC    m_ 5 56   $ $"&)"*>*@*@AAM @@%-#Y^ AM 2? o . .  w / /  2295)8$1$E# 3M2? o . .p ).%!%!1!9!9*]:h:hVg:h:h!i!i%%rd.Dataset | DatasetDict | IterableDataset | Nonec|,j%jtgStd||nj}jt jjjr jjjj jj jj jj d}t|tr.|jjjjdn=t|t&rtdt|t(r|D]&}t|trtd'fd|D}t-|}||jj}||d <n]t|t0r9|jjjjj }||d <ntd d j_jt|fi|S) a Returns the evaluation [`~torch.utils.data.DataLoader`]. Subclass and override this method if you want to inject some custom behavior. Args: eval_dataset (`torch.utils.data.Dataset`, *optional*): If provided, will override `self.eval_dataset`. If it is a [`~datasets.Dataset`], columns not accepted by the `model.forward()` method are automatically removed. It must implement `__len__`. NzQEvaluation requires specifying an eval_dataset to the SentenceTransformerTrainer.r#rr)r*c vg|]5}|jjjjj6Sr,rr7eval_batch_sizer/r r0s rbrlzBSentenceTransformerTrainer.get_eval_dataloader..\   &&#y8"i<(5(I' '   rdrrr-z_Unsupported `eval_dataset` type. Use a Dataset, DatasetDict, or IterableDataset for evaluation.T)r:r=rrzr?rr1r7rr2r3r4r5r6rjr-rrBr/r.r,r7r r!r+rr r8r9r:)rrr:r<rrrr?r s` @@rbget_eval_dataloaderz.SentenceTransformerTrainer.get_eval_dataloaders  D$5$=~)!"~~%pqq q'3'?||TEV * O%% 9> 2  ! !$). 1 1 1(9;)9"&)"I#yC    lO 4 45   $ $"&)";!%!?     &9 : :- u  k 2 2( '..00  g77$s       ,2244   N))<)<)>)>??L @@$-#Y^ AM 2? o . .  g . .  2294)8$1$E# 3M2? o . .q )-%'' <(U(UCT(U(UVVVrd test_dataset'Dataset | DatasetDict | IterableDatasetcTjtjjjrjjjjjjjjjj d}t|tr.| jj jjdn=t|trt!dt|t"r|D]&}t|trt!d'fd|D}t'|}||jj}||d<n]t|t*r9|jj jjj}||d<nt!d d j_jt7|fi|S) a Returns the training [`~torch.utils.data.DataLoader`]. Subclass and override this method if you want to inject some custom behavior. Args: test_dataset (`torch.utils.data.Dataset`, *optional*): The test dataset to use. If it is a [`~datasets.Dataset`], columns not accepted by the `model.forward()` method are automatically removed. It must implement `__len__`. r#rr)r*c vg|]5}|jjjjj6Sr,rAr0s rbrlzBSentenceTransformerTrainer.get_test_dataloader..SrCrdrrr-z\Unsupported `test_dataset` type. Use a Dataset, DatasetDict, or IterableDataset for testing.T)r?rr1r7rr2r3r4r5r6rjr-rrBr/r.rzr,r7r r!r+rr r8r9r:r)rrrEr<rrrr?r s` @@rbget_test_dataloaderz.SentenceTransformerTrainer.get_test_dataloader&s|* O%% 9> 2  ! !$). 1 1 1(9;)9"&)"I#yC    lO 4 45   $ $"&)";!%!?     &9 : :- u  k 2 2( '..00  g77$s       ,2244   N))<)<)>)>??L @@$-#Y^ AM 2? o . .  g . .  2294)8$1$E# 3M2? o . .n )-%'' <(U(UCT(U(UVVVrdrQc"||n |jj}tj|dtd||j||jjtttdkr"|j |j |n!|j |j |tj|jtj|t"dS)NT)exist_okzSaving model checkpoint to )safe_serializationrf)r7rQrmakedirsrxryr5save_pretrainedsave_safetensorsrrrgrArsavepathjoinr)rrrQ state_dicts rb_savez SentenceTransformerTrainer._save|s#-#9ZZty?S  J.... >*>>??? "":$)B\"]]] - . .-2I2I I I$0%55jAAA~)..z:::  49bgll:7IJJKKKKKrdcheckpoint_pathcddlm}|||jj}|j|dS)Nrr0)trust_remote_code)rr1r5rWload_state_dictrS)rrrUr1 loaded_models rb_load_from_checkpointz0SentenceTransformerTrainer._load_from_checkpointsV======**?djNjkkk  ""<#:#:#<#<=====rdpromptc |j|S#t$r>|j|gdjddz }||j|<|cYSwxYw)Nrrrs)rKeyErrorr5rshape)rrr[ prompt_lengths rb_get_prompt_lengthz-SentenceTransformerTrainer._get_prompt_lengthss !.v6 6 ! ! ! J//99+FLRPSTTM2?D ' /  !s AAAcX|jD]!}t|tr |j cS"dS)a Return whether the prompt length should be passed to the model's forward method. True if the model does not include the prompt in the pooling layer. Can be overridden by the user if it's useful to include the prompt length. F)r5rjrinclude_prompt)rrmodules rb_include_prompt_lengthz1SentenceTransformerTrainer._include_prompt_lengthsCj 1 1F&'** 1!00000 1urdbatchdict[str, list[Any]]rdict[str, str] | str | Noneprompt_lengthsdict[str, int] | int | None transform6Callable[[dict[str, list[Any]]], dict[str, list[Any]]]c  |r ||}|s|St|d}||s|Sttrnt|D]L\}}t|dtr,fd|D||<||gt |z||d<Mtt r\D]G\} ||vr> fd||D||<|r%| gt ||z||d<H|r|gt ||z|d<|S)a~A transform/map function that adds prompts or dataset names to the batch. Args: batch (dict[str, list[Any]]): The batch of data, where each key is a column name and each value is a list of values. prompts (dict[str, str] | str | None, optional): An optional mapping of column names to string prompts, or a string prompt for all columns. Defaults to None. prompt_lengths (dict[str, int] | int | None, optional): An optional mapping of prompts names to prompt token length, or a prompt token length if the prompt is a string. Defaults to None. dataset_name (str | None, optional): The name of this dataset, only if there are multiple datasets that use a different loss. Defaults to None. transform (Callable[[dict[str, list[Any]]], dict[str, list[Any]]], optional): An optional transform function to apply on the batch before adding prompts, etc. Defaults to None. Returns: dict[str, list[Any]]: The "just-in-time" transformed batch with prompts and/or dataset names added. rcg|]}|zSriri)r^r`rs rbrlzTSentenceTransformerTrainer.add_prompts_or_dataset_name_transform..s)N)N)Ne'E/)N)N)NrdN_prompt_lengthcg|]}|zSriri)r^r`r[s rbrlzTSentenceTransformerTrainer.add_prompts_or_dataset_name_transform..s)Y)Y)YU&5.)Y)Y)Yrdrp)rrrjrrrr) rerrhrprjkwargs first_column column_namerr[s ` @rb%add_prompts_or_dataset_name_transformz@SentenceTransformerTrainer.add_prompts_or_dataset_name_transforms8  %Ie$$E LEJJLL))!, \" L gs # # _'+EKKMM':': _ _# VfQi--_)N)N)N)Nv)N)N)NE+&%1AO@PSVW]S^S^@^<<<= gt $ $ s'.}} s s# V%'')Y)Y)Y)YeKFX)Y)Y)YE+&%sAOPVAW@X[^_dep_q[r[r@r<<<=  N%1NS|9L5M5M$ME. ! rdDatasetDict | Dataset | None7dict[str, dict[str, str]] | dict[str, str] | str | Nonec|dSt|jt}t|dr|S||||s|r-|}|||||}|S)a Maybe add prompts or dataset names to the dataset. We add the dataset_name column to the dataset if: 1. The loss is a dictionary and the dataset is a DatasetDict, or 2. The prompts contain a mapping to dataset names. There are 4 cases for the prompts: 1. `str`: One prompt for all datasets and columns. 2. `dict[str, str]`: A column to prompt mapping. 3. `dict[str, str]`: A dataset to prompt mapping. 4. `dict[str, dict[str, str]]`: A dataset to column to prompt mapping. And 2 cases for the dataset: A. `Dataset`: A single dataset. B. `DatasetDict`: A dictionary of datasets. 3A is not allowed, and 2A doesn't make sense. Args: dataset_dict (DatasetDict | Dataset | None): The dataset to add prompts or dataset names to. Returns: DatasetDict | Dataset | None: The dataset with prompts or dataset names added. N#_sentence_transformers_preprocessedrv)rinclude_prompt_lengthsinclude_dataset_name)rjr;rrrrd"add_prompts_or_dataset_name_column)rrrrrpryrxs rbrzCSentenceTransformerTrainer.maybe_add_prompts_or_dataset_name_columns@  4)$)T:: nnt|trY|D]D\} } | | } | |vr3t| t&r| jdvrt'd | | d <Etj|||d | ntd d _S)N)rrrprxryra=The prompts provided to the trainer are a nested dictionary. In this setting, the first level of the dictionary should map to dataset names and the second level to column names. However, as the provided dataset is a not a DatasetDict, no dataset names can be inferred. The keys to the provided prompts dictionary are cTi|]$\}}|jv||%Sri)rr`)r^rrr[rrrs rbrczQSentenceTransformerTrainer.add_prompts_or_dataset_name_column..GsE&&&/K&,*CCC 7 7 ? ?CCCrd)rrhrprVrp)rV large_stringint16rnT)batchedrzUnsupported dataset type.)rjr.r,rrr\rzrr`rrrzr+ set_transformrrs_format_kwargsr-rr/rdtypemaprw)rrrrrprxryrnested_promptsrh first_keyrrrfeaturer[s`` rbrzz=SentenceTransformerTrainer.add_prompts_or_dataset_name_column s l%8+$F G G )5););)=)=  % gGQRY[_G`G`!m\7!C!C!Cfm-1-T-T!(*1E!O4+A)= .U.. \**   '3'' )F%)%<%##1!-  #1      o 6 6! :#,H ;+0??( Vgs++ V'3'@VV "*;"7%gu55V'-Ke:e:eGLW~~H %C%C%CDV ..V/6}}VV+ V"*;"7"n44 *7E : :5 ' 1K K KGLW~~H %C%C%CD'++>##1!-  !,  LL899 9<@ 8rdlanguagelicensetagsstr | list[str] | None model_namefinetuned_fromtasks dataset_tags dataset_argsc  D|sdS|r|jj||r|jj||r|jj||j|jj|dS)N)r) is_world_process_zeror5r set_language set_licenseadd_tags_create_model_cardr7rQ) rrrrrrrrrrrrps rbcreate_model_cardz,SentenceTransformerTrainer.create_model_cards))++  F  > J & 3 3H = = =  < J & 2 27 ; ; ;  6 J & / / 5 5 5 %%di&:z%RRRRRrdtuple[Any, Any]ct|jtr'tjt |j}n|j}t ||\}}hdt| zsh| |fd| D|j j dfd| Dddg|d<||fS)a? We have to override the optimizer_grouped_parameters because the Trainer superclass bases it on the `model` itself, but the SentenceTransformer losses can have weights that should be updated as well, e.g. SoftmaxLoss (see #2872). This method requires `transformers` >= 4.43.0. >r5paramsoptimizer_dictc0g|]\}}|v |j|Sri requires_gradr^npdecay_parameterss rbrlzKSentenceTransformerTrainer.get_optimizer_cls_and_kwargs..s8"aM]H]H]bcbqH]H]H]H]rd)r weight_decayc0g|]\}}|v |j|Srirrs rbrlzKSentenceTransformerTrainer.get_optimizer_cls_and_kwargs..s8"aQaHaHafgfuHaHaHaHardgr)rjr;rr Sequentialrrget_optimizer_cls_and_kwargsrrget_decay_parameter_namesnamed_parametersr7r)rrr7r5 loss_model optimizer_clsoptimizer_kwargsrrs @rbrz7SentenceTransformerTrainer.get_optimizer_cls_and_kwargss= di & & #{49'='=>>JJJ*/''*N*NtU_*`*`' '544s;K;P;P;R;R7S7SS #==jII &0&A&A&C&C%)I$: &0&A&A&C&C%(  2 - ....rd) NNNNNNNNNNNr4N)r5r6r7r'r8r9r:r9r;r<r=r>r?r@rArBrCrDrErFrGrHrIrJrKrLrMrN)rrrMrNr)rMr1)r;rr5r1rMr)r;rr5r1rMr)rr,rMr,)FN)r5r1rrrrYrMr)rrrMr)NNrU)r:rrrrrrMr) rrrrrrrrrrrMr)rMrN)rr+rprrMrN) rr+r rr rYr rr r rMr)Nr) rr rrr r rrrMr )rMr)r:r>rMr)rErFrMr)rQrrMrN)rUrrMrN)r[rrMr)rMrY)NNNN) rerfrrgrhrirprrjrkrMrf)rrtrrurprrMrt)NNFF) rr{rrgrprrxrYryrYrMrt) NNNNNNNNN)rrrrrrrrrrrrrrrrrrrMrN)r7r'r5r6rMr) __name__ __module__ __qualname____doc__rrr{rrorrrrrrrrr!r=rDrIrTrZr`rd staticmethodrsrrzrr __classcell__)rs@rbr3r32sCCN-159]a\` HL-1?C?CCG26Vbei%`8`8`8`8`8`8`8DYYYY(2,,,, % 00000d! ! ! ! J=A(,!' N N N N N N N$-1(,!' )))))))V9999992     "15,0 66666x-1 #####JX&X&X&X&tZWZWZWZWZWxTWTWTWTWlLLLLL&>>>> !!!!    046:#'LP ====\=DLP#' 55555t04#'',%* `````H $"'+!%%)(,/3*./3SSSSS4_c$/$/$/$/$/$/$/$/$/$/$/rdr3)N __future__rrr collectionsr contextlibr functoolsrtypingrrr rpackaging.versionr rr torch.utils.datar r rr transformersrrrrrrtransformers.data.data_collatorrtransformers.integrationsrtransformers.trainerrtransformers.trainer_utilsr#sentence_transformers.data_collatorr sentence_transformers.evaluationrr'sentence_transformers.losses.CoSENTLossr sentence_transformers.model_cardrsentence_transformers.modelsr(sentence_transformers.models.Transformerrsentence_transformers.samplerr r!r"r#r$#sentence_transformers.training_argsr%r&r'sentence_transformers.utilr(r)r*datasetsr+r,r-r.r/ getLoggerrrx)sentence_transformers.SentenceTransformerr1r3rirdrbrs"""""" ######""""""////////// 444444YYYYYYYYYYYYZZZZZZZZZZZZ<<<<<<888888333333333333555555OOOOOOSSSSSSSS>>>>>>>>>>>>000000@@@@@@ eddddddddd[ZZZZZZZZZZZZZZ  8 $ $NMMMMMMM/M/M/M/M/M/M/M/M/M/rd