Ch+SSKJr SSKrSSKJr SSKJr SSKrSSKJ s J r SSKJ r J r SSK Jr SSKJrJrJr SSjr"S S 5r"S S 5r"S S\ R,5rg)) annotationsN)Iterable)Any)Tensornn)SentenceTransformer)CachedGISTEmbedLoss"CachedMultipleNegativesRankingLoss+CachedMultipleNegativesSymmetricRankingLosscURSnX:a[SUSU35eUSSU24n[R"USSS9nU$)Nz Dimension zL in matryoshka_dims cannot be greater than the model's embedding dimension: .)pdim)shape ValueErrorF normalize)tensorr tensor_dims c/var/www/html/shao/venv/lib/python3.13/site-packages/sentence_transformers/losses/MatryoshkaLoss.pyshrinkrsab!J ijtiu v  C#I F [[1" -F Mc6\rSrSrSrSSjrSSjrS SjrSrg) ForwardDecoratora? This decorator is used to cache the output of the Sentence Transformer's forward pass, so that it can be shrank and reused for multiple loss calculations. This prevents the model from recalculating the embeddings for each desired Matryoshka dimensionality. This decorator is applied to `SentenceTransformer.forward`. cHXlSUl/UlSUlSUlgNr)fnrcache cache_dimidx)selfrs r__init__ForwardDecorator.__init__'s# rcXlSUlgr)rr")r#rs rset_dimForwardDecorator.set_dim/srcURbURUR:Xa>URU5nURR U5 URUlOURUR nSU;a[ USUR5US'[ USUR5US'U=R S- slU$)Ntoken_embeddingssentence_embedding)r!rrr appendr"r)r#featuresoutputs r__call__ForwardDecorator.__call__3s >> !T^^txx%?WWX&F JJ  f %!XXDNZZ)F  ')/7I0JDHH)UF% &'-f5I.JDHH'U#$ A  r)r r!rrr"N)returnNone)r.dict[str, Tensor]r2r4) __name__ __module__ __qualname____firstlineno____doc__r$r'r0__static_attributes__rrrrs rrc@\rSrSrSrSSSjjrS SjrSrg) CachedLossDecoratorCa This decorator is used with the Cached... losses to compute the underlying loss function for each Matryoshka dimensionality. This is done by shrinking the pre-computed embeddings to the desired dimensionality and then passing them to the underlying loss function once for each desired dimensionality. This decorator is applied to the `calculate_loss` method of the Cached... losses. c4XlX lX0lX@lg)Nrmatryoshka_dimsmatryoshka_weightsn_dims_per_step)r#rrArBrCs rr$CachedLossDecorator.__init__Ms."4.rc :[[UR55nURS:a9UR[U5:a [R "X@R5nSnUGHnURUnUR UnUV V s/sHoV s/sHn [X5PM sn PM n n n [R"5n U (a>U V V s/sH/oV s/sH oR5R5PM" sn PM1 n n n OU n XXUR"U /UQ70UD6-- nU (dM[X5H8up[X5H$unnURUUR-5 M& M: GM U$s sn fs sn n fs sn fs sn n fNrg)rangelenrArCrandomsamplerBrtorchis_grad_enableddetachrequires_grad_rzipbackwardgrad)r#repsargskwargs dim_indiceslossr"rweight minibatchr truncatedcompute_gradientsmatryoshka_reps t_minibatch d_minibatchtds rr0CachedLossDecorator.__call__UsiC 4 456   ! #(<(1P!& =S $T"os0F F/F F"&FFF Fr@N)r )rA list[int]rBzlist[float | int]rCintr2r3)rRzlist[list[Tensor]]r2r)r5r6r7r8r9r$r0r:r;rrr=r=Cs8mo/#,/BS/fi/ /rr=cr^\rSrSrSSU4SjjjrS SjrS Sjr\S Sj5rSr U=r $) MatryoshkaLossqc>[TU]5 XlX lUcS/[ U5-n[ X45n[ [ USSS96uUlUlXPl [[[4Ul [X R5(a1[UR URUR5Ulgg)a* The MatryoshkaLoss can be seen as a loss *modifier* that allows you to use other loss functions at various different embedding dimensions. This is useful for when you want to train a model where users have the option to lower the embedding dimension to improve their embedding comparison speed and costs. This loss is also compatible with the Cached... losses, which are in-batch negative losses that allow for higher batch sizes. The higher batch sizes allow for more negatives, and often result in a stronger model. Args: model: SentenceTransformer model loss: The loss function to be used, e.g. :class:`MultipleNegativesRankingLoss`, :class:`CoSENTLoss`, etc. matryoshka_dims: A list of embedding dimensions to be used for the loss function, e.g. [768, 512, 256, 128, 64]. matryoshka_weights: A list of weights to be used for the loss function, e.g. [1, 1, 1, 1, 1]. If None, then the weights will be set to 1 for all dimensions. n_dims_per_step: The number of dimensions to use per step. If -1, then all dimensions are used. If > 0, then a random sample of n_dims_per_step dimensions are used per step. The default value is -1. References: - The concept was introduced in this paper: https://arxiv.org/abs/2205.13147 - `Matryoshka Embeddings <../../../examples/sentence_transformer/training/matryoshka/README.html>`_ Inputs: +---------------------------------------+--------+ | Texts | Labels | +=======================================+========+ | any | any | +---------------------------------------+--------+ Relations: - :class:`Matryoshka2dLoss` uses this loss in combination with :class:`AdaptiveLayerLoss` which allows for layer reduction for faster inference. Example: :: from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses from datasets import Dataset model = SentenceTransformer("microsoft/mpnet-base") train_dataset = Dataset.from_dict({ "anchor": ["It's nice weather outside today.", "He drove to work."], "positive": ["It's so sunny.", "He took the car to the office."], }) loss = losses.MultipleNegativesRankingLoss(model) loss = losses.MatryoshkaLoss(model, loss, [768, 512, 256, 128, 64]) trainer = SentenceTransformerTrainer( model=model, train_dataset=train_dataset, loss=loss, ) trainer.train() Nr,c US$rr;)xs r)MatryoshkaLoss.__init__..s `abc`drT)keyreverse)superr$modelrVrHrOsortedrArBrCr r r cached_losses isinstancer=calculate_loss)r#rorVrArBrC dims_weights __class__s rr$MatryoshkaLoss.__init__rsF    %"#s?';!; ?? 8;VLVdnr=s8t5d5. /  7  d.. / /"5##T%9%94;R;R#D  0rc[URUR5(aURX5$URRn[ U5nX@Rl[ [UR55nURS:aIUR[U5:a0[R"XPR5nUR5 SnUHHnURUnURUn URU5 XiURX5-- nMJ X0RlU$!X0Rlf=frF)rrrVrqroforwardrrGrHrArCrIrJsortrBr') r#sentence_featureslabelsoriginal_forwarddecorated_forwardrUrVr"rrWs rrxMatryoshkaLoss.forwards  dii!3!3 4 499.7 7 ::-- 2 01A B !2JJ D$8$8 9:K##a'D,@,@3{CS,S$mmK9M9MN   "D"**3/005!))#.+rsW" $ 5""J++\KRYYKr