ó
    — óhª  ã                  ó€   • S SK Jr  S SKJr  S SKJs  Jr  S SKJ	r	Jr  S SK
Jr  SSKJr   " S S	\R                  5      rg)
é    )Úannotations)ÚIterableN)ÚTensorÚnn)ÚSentenceTransformeré   )ÚSiameseDistanceMetricc                  ó`   ^ • \ rS rSr\R
                  S4     SU 4S jjjrSSS jjrSrU =r	$ )	ÚOnlineContrastiveLossé   g      à?c                óF   >• [         TU ]  5         Xl        X0l        X l        g)ay  
This Online Contrastive loss is similar to :class:`ConstrativeLoss`, but it selects hard positive (positives that
are far apart) and hard negative pairs (negatives that are close) and computes the loss only for these pairs.
This loss often yields better performances than ContrastiveLoss.

Args:
    model: SentenceTransformer model
    distance_metric: Function that returns a distance between
        two embeddings. The class SiameseDistanceMetric contains
        pre-defined metrics that can be used
    margin: Negative samples (label == 0) should have a distance
        of at least the margin value.

References:
    - `Training Examples > Quora Duplicate Questions <../../../examples/sentence_transformer/training/quora_duplicate_questions/README.html>`_

Requirements:
    1. (anchor, positive/negative) pairs
    2. Data should include hard positives and hard negatives

Inputs:
    +-----------------------------------------------+------------------------------+
    | Texts                                         | Labels                       |
    +===============================================+==============================+
    | (anchor, positive/negative) pairs             | 1 if positive, 0 if negative |
    +-----------------------------------------------+------------------------------+

Relations:
    - :class:`ContrastiveLoss` is similar, but does not use hard positive and hard negative pairs.
      :class:`OnlineContrastiveLoss` often yields better results.

Example:
    ::

        from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
        from datasets import Dataset

        model = SentenceTransformer("microsoft/mpnet-base")
        train_dataset = Dataset.from_dict({
            "sentence1": ["It's nice weather outside today.", "He drove to work."],
            "sentence2": ["It's so sunny.", "She walked to the store."],
            "label": [1, 0],
        })
        loss = losses.OnlineContrastiveLoss(model)

        trainer = SentenceTransformerTrainer(
            model=model,
            train_dataset=train_dataset,
            loss=loss,
        )
        trainer.train()
N)ÚsuperÚ__init__ÚmodelÚmarginÚdistance_metric)Úselfr   r   r   Ú	__class__s       €Úl/home/james-whalen/.local/lib/python3.13/site-packages/sentence_transformers/losses/OnlineContrastiveLoss.pyr   ÚOnlineContrastiveLoss.__init__   s!   ø€ ôn 	‰ÑÔØŒ
ØŒØ.Õó    c                ó8  • U Vs/ s H  o@R                  U5      S   PM     nnU R                  US   US   5      nXbS:H     nXbS:H     nXw[        U5      S:”  a  UR                  5       OUR	                  5       :     n	Xˆ[        U5      S:”  a  UR                  5       OUR	                  5       :„     n
U
R                  S5      R                  5       n[        R                  " U R                  U	-
  5      R                  S5      R                  5       nX¼-   nU$ s  snf )NÚsentence_embeddingr   r   é   )r   r   ÚlenÚmaxÚmeanÚminÚpowÚsumÚFÚrelur   )r   Úsentence_featuresÚlabelsÚsize_averageÚsentence_featureÚ
embeddingsÚdistance_matrixÚnegsÚpossÚnegative_pairsÚpositive_pairsÚpositive_lossÚnegative_lossÚlosss                 r   ÚforwardÚOnlineContrastiveLoss.forwardJ   sü   € ÙarÓsÒarÐM]—j‘jÐ!1Ó2Ð3GÔHÑarˆ
Ðsà×.Ñ.¨z¸!©}¸jÈ¹mÓLˆØ¨™{Ñ+ˆØ¨™{Ñ+ˆð ´C¸³IÀ³M d§h¡h¤jÀtÇyÁyÃ{ÑSÑTˆØ´C¸³IÀ³M d§h¡h¤jÀtÇyÁyÃ{ÑSÑTˆà&×*Ñ*¨1Ó-×1Ñ1Ó3ˆÜŸš˜tŸ{™{¨^Ñ;Ó<×@Ñ@ÀÓC×GÑGÓIˆØÑ,ˆØˆùò ts   …D)r   r   r   )r   r   r   ÚfloatÚreturnÚNone)F)r#   zIterable[dict[str, Tensor]]r$   r   r3   r   )
Ú__name__Ú
__module__Ú__qualname__Ú__firstlineno__r	   ÚCOSINE_DISTANCEr   r0   Ú__static_attributes__Ú__classcell__)r   s   @r   r   r      s<   ø† à:O×:_Ñ:_Ðqtð:/Ø(ð:/Øinð:/à	÷:/ð :/÷xô r   r   )Ú
__future__r   Úcollections.abcr   Útorch.nn.functionalr   Ú
functionalr!   Útorchr   Ú)sentence_transformers.SentenceTransformerr   ÚContrastiveLossr	   ÚModuler   © r   r   Ú<module>rE      s,   ðÝ "å $ç Ð ß å Iå 2ôK˜BŸI™Iõ Kr   