
    h?                        S SK Jr  S SKJrJrJrJrJrJrJ	r	J
r
JrJr  SSKJr  SSKJrJr  SSKJrJr  \
" S5      r\
" S5      r\
" S	5      r\
" S
5      r\\\4   r\\\\\   \\   4   r " S S\\\\\4   5      r " S S\5      r\R>                  " S5      SSSS.S\ S\\\      S\\\\4      S\4S jj5       r!\R>                  " S5      SSSSS.S\ S\\\      S\\\\4      S\\   S\4
S jj5       r"\R>                  " S5      SSSSSS.S\ S\\\      S\\\\4      S\\   S\#S\4S  jj5       r$ " S! S"\5      r%\R>                  " S#5      SSS$.S\ S\\\      S\%4S% jj5       r&\R>                  " S&5      SSSS'.S\ S\\\      S\\   S\%4S( jj5       r'\R>                  " S)5      SSSSSS.S\ S\\\      S\\\\4      S\\   S\#S\%4S* jj5       r( " S+ S,\5      r)\R>                  " S-5      SS..S\ S\)4S/ jj5       r* " S0 S1\5      r+\R>                  " S25      SS3S4.S\ S5\ S\+4S6 jj5       r,S\4S7 jr-S\4S8 jr./ S9Qr/g):    )abstractmethod)
AnyDictGenericListOptionalSequenceTupleTypeVarUnioncast   )registry)Floats2dInts1d)get_array_moduleto_categoricalLossTGradTGuessTTruthTc                       \ rS rSrSrS\SS4S jrS\S\S\	\
\4   4S	 jr\S\S\S\
4S
 j5       r\S\S\S\4S j5       rSrg)Loss   a	  Base class for classes computing the loss / gradient. The class can
be initialized with settings if needed. It provides get_loss and
get_grad as separate methods to allow calculating them separately. It
also provides a __call__ method that returns a tuple of both.
kwargsreturnNc                     g N )selfr   s     D/home/james-whalen/.local/lib/python3.13/site-packages/thinc/loss.py__init__Loss.__init__"   s        guessestruthsc                 F    U R                  X5      U R                  X5      4$ r   get_gradget_lossr    r%   r&   s      r!   __call__Loss.__call__%       }}W-t}}W/MMMr$   c                     g r   r   r+   s      r!   r)   Loss.get_grad(       r$   c                     g r   r   r+   s      r!   r*   Loss.get_loss,   r1   r$   r   )__name__
__module____qualname____firstlineno____doc__r   r"   r   r   r
   r   r   r,   r   r)   r*   __static_attributes__r   r$   r!   r   r      s      N N N5;N N   5     5  r$   r   c                   $   \ rS rSr% \\\      \S'   \\\\	4      \S'   \
\\	4   \S'   SSSSSS.S	\S\\\      S\\\\	4      S
\\   S\4
S jjrS\S\\\4   4S jrS\S\S\\\4   4S jrS\S\S\4S jrS\S\S\4S jrS\S\4S jrSrg)CategoricalCrossentropy1   namesmissing_value
_name_to_iTN        	normalizer=   r>   
neg_prefixlabel_smoothingrB   rC   rD   c                    Xl         X l        X0l        X@l        XPl        Ub&  [        U5       VVs0 s H  u  pgXv_M	     snnU l        g 0 U l        g s  snnf r   )rB   r=   r>   rC   rD   	enumerater?   )r    rB   r=   r>   rC   rD   inames           r!   r"    CategoricalCrossentropy.__init__6   sS     #
*$.6?6FG6F71tw6FGDO DO Hs   Ar%   r   c                    [        U5      n/ nS nU R                  (       a.  UR                  [        U5      [        U R                  5      4SS9nU R                  n[        U[        5      (       Gaf  [        U5      n[        U5      (       Ga.  [        US   [        5      (       a-  [        U5       H  u  pxX:X  d  M  UR                  U5        M     OU R                  c  Sn	[        U	5      e[        U5       H  u  pxX:X  a$  U R                  S   X'   UR                  U5        M.  U(       d  M7  U R                  (       d  MJ  UR                  U R                  5      (       d  Ml  U[        U R                  5      S  X'   U R                  X      n
SXW'   SXW   U
'   M     U Vs/ s H  oR                  U   PM     nnUR                  USS9n[        X$5      nO[!        XU5      nUR"                  UR"                  :w  a1  [%        ['        [(        U5      UR*                  S   U R,                  S9nOU R,                  (       a  [        S5      eUb  X-  nSXS:H  '   S	XUS:H  '   X-  nX4$ s  snf )
Nfdtyper   zCannot calculate loss from list of strings without names. You can pass the names as a keyword argument when you create the loss object, e.g. CategoricalCrossentropy(names=['dog', 'cat'])rG   )	n_classesrD   zLabel smoothing is only applied, when truths have type List[str], List[int] or Ints1d, but it seems like Floats2d was provided.r   )r   r=   oneslenr>   
isinstancelistintrF   append
ValueErrorrC   
startswithr?   asarray
_make_mask_make_mask_by_valuendimr   r   r   shaperD   )r    r&   r%   xpmissingnegatives_maskr>   rG   valuemsg	neg_indexrH   masks                r!   convert_truths&CategoricalCrossentropy.convert_truthsI   s/   g&::WWc&k3tzz?%C3WON**fd##&\F6{{fQi--$-f$5 1#NN1- %6 zz)Q  )o-$-f$5 1(,

1FI#NN1-!E $ % 0 0 A A(-c$//.B.D(EFI(,	(BI01N-;=N-i8 %6 AGGood3FGZZcZ2Fg/D&vFD;;',,&#VV$!--+ $ 4 4F ## $  %$F#$FR< 34NR/0"D|5 Hs   !I&r&   c                 H    U R                  X5      nX0R                  U5      4$ r   r)   _get_loss_from_gradr    r%   r&   d_truths       r!   r,    CategoricalCrossentropy.__call__   s&     --011':;;r$   c                    U R                  X!5      u  p4[        U5      nUR                  UR                  :w  a(  SUR                   SUR                   S3n[        U5      eUR	                  US:  5      (       d  UR	                  US:  5      (       a  Sn[        U5      eUR	                  US:  5      (       d  UR	                  US:  5      (       a  Sn[        U5      eX-
  nXt-  nU R
                  (       a  XqR                  S   -  nU$ )NzBCannot calculate CategoricalCrossentropy loss: mismatched shapes:  vs .r   r   zVCannot calculate CategoricalCrossentropy loss with guesses outside the [0,1] interval.z[Cannot calculate CategoricalCrossentropy loss with truth values outside the [0,1] interval.)rd   r   r\   rV   anyrB   )r    r%   r&   targetrc   r]   err
differences           r!   r)    CategoricalCrossentropy.get_grad   s    **6;f%==FLL(VW^WdWdVeeijpjvjviwwxyCS/!66'A+"&&1"5"5jCS/!66&1*
!3!3oCS/!%

>>#mmA&66Jr$   c                 F    U R                  X5      nU R                  U5      $ r   rg   ri   s       r!   r*    CategoricalCrossentropy.get_loss   s!    --0''00r$   rj   c                 (    US-  R                  5       $ )N   )sum)r    rj   s     r!   rh   +CategoricalCrossentropy._get_loss_from_grad   s    
!!r$   )r?   rD   r>   r=   rC   rB   )r4   r5   r6   r7   r   r	   str__annotations__r   rT   r   boolfloatr"   r   r
   rd   IntsOrFloatsOrStrsr,   r)   r*   rh   r9   r   r$   r!   r;   r;   1   s!   HSM""E#s(O,,S#X
 )-37$(!$! ! &	!
  c3h0! SM! !&?h ?58AS;T ?B<<);<	x	< 2D  $1 12D 1 1"8 " "r$   r;   zCategoricalCrossentropy.v1TNrB   r=   r>   rB   r=   r>   r   c                     [        XUS9$ )Nr   r;   r   s      r!   $configure_CategoricalCrossentropy_v1r      s     # r$   zCategoricalCrossentropy.v2rB   r=   r>   rC   rC   c                     [        U UUUS9$ )Nr   r   r   s       r!   $configure_CategoricalCrossentropy_v2r      s     ##	 r$   zCategoricalCrossentropy.v3r@   rA   rD   c                     [        U UUUUS9$ NrA   r   rA   s        r!   $configure_CategoricalCrossentropy_v3r      s      ##' r$   c                       \ rS rSrSSSSSS.S\S\\\      S\\\\	4      S	\\   S
\
4
S jjrS\\   S\\   S\\\   \
4   4S jrS\\   S\\   S\\   4S jrS\\   S\\   S\
4S jrS\\   S\
4S jrSrg)SequenceCategoricalCrossentropy   TNr@   rA   rB   r=   r>   rC   rD   c                4    [        SUUUUS9U l        Xl        g )NFrA   )r;   ccrB   )r    rB   r=   r>   rC   rD   s         r!   r"   (SequenceCategoricalCrossentropy.__init__   s&     *'!+
 #r$   r%   r&   r   c                 L    U R                  X5      nU R                  U5      nX44$ r   rg   )r    r%   r&   gradslosss        r!   r,   (SequenceCategoricalCrossentropy.__call__   s*     g.''.{r$   c                    Sn[        U5      [        U5      :w  a  [        U5      e[        U5      n/ n[        X5       HF  u  pgU R                  R	                  Xg5      nU R
                  (       a  X-  nUR                  U5        MH     U$ )Nz]Cannot calculate SequenceCategoricalCrossentropy loss: guesses and truths must be same length)rQ   rV   zipr   r)   rB   rU   )	r    r%   r&   rq   nd_scoresyhyd_yhs	            r!   r)   (SequenceCategoricalCrossentropy.get_grad   sx     nw<3v;&S/!L)EB77##B*D~~	OOD!	 *
 r$   c                 B    U R                  U R                  X5      5      $ r   )rh   r)   r+   s      r!   r*   (SequenceCategoricalCrossentropy.get_loss  s     ''g(FGGr$   r   c                 V    SnU H   nX R                   R                  U5      -  nM"     U$ )Nr@   )r   rh   )r    r   r   grads       r!   rh   3SequenceCategoricalCrossentropy._get_loss_from_grad  s-    DGG//55D r$   )r   rB   )r4   r5   r6   r7   r|   r   r	   rz   r   rT   r}   r"   r   r~   r
   r   r,   r)   r*   rh   r9   r   r$   r!   r   r      s
    )-37$(!$# # &	#
  c3h0# SM# #$)3;<N3O	tH~u$	%)3;<N3O	hH)H3;<N3OH	H
();  r$   r   z"SequenceCategoricalCrossentropy.v1rB   r=   c                     [        XS9$ )Nr   r   r   s     r!   ,configure_SequenceCategoricalCrossentropy_v1r     s     +YLLr$   z"SequenceCategoricalCrossentropy.v2rB   r=   rC   c                     [        XUS9$ )Nr   r   r   s      r!   ,configure_SequenceCategoricalCrossentropy_v2r     s     +Z r$   z"SequenceCategoricalCrossentropy.v3c                     [        U UUUUS9$ r   r   rA   s        r!   ,configure_SequenceCategoricalCrossentropy_v3r   !  s      +#' r$   c                   v    \ rS rSrSS.S\4S jjrS\S\S\\\4   4S	 jr	S\S\S\4S
 jr
S\S\S\4S jrSrg)
L2Distancei3  TrB   rB   c                    Xl         g r   r   )r    rB   s     r!   r"   L2Distance.__init__4  s    "r$   r%   r&   r   c                 F    U R                  X5      U R                  X5      4$ r   r(   r+   s      r!   r,   L2Distance.__call__7  r.   r$   c                     UR                   UR                   :w  a(  SUR                    SUR                    S3n[        U5      eX-
  nU R                  (       a  XAR                   S   -  nU$ )N1Cannot calculate L2 distance: mismatched shapes: rm   rn   r   )r\   rV   rB   )r    r%   r&   rq   rr   s        r!   r)   L2Distance.get_grad:  sd    ==FLL(Egmm_TXY_YeYeXffghCS/!%
>>#mmA&66Jr$   c                     UR                   UR                   :w  a(  SUR                    SUR                    S3n[        U5      eU R                  X5      nUS-  R                  5       $ )Nr   rm   rn   rw   )r\   rV   r)   rx   )r    r%   r&   rq   rj   s        r!   r*   L2Distance.get_lossC  s_    ==FLL(Egmm_TXY_YeYeXffghCS/!--0
!!r$   N)r4   r5   r6   r7   r|   r"   r   r
   r}   r,   r)   r*   r9   r   r$   r!   r   r   3  sl    ,0 #T #N N( NuXu_?U N ( x " "( "u "r$   r   zL2Distance.v1r   c                     [        U S9$ )Nr   )r   r   s    r!   configure_L2Distancer   L  s    	**r$   c                       \ rS rSrSSS.S\S\4S jjrS\S	\S
\\\4   4S jr	S\S	\S
\4S jr
S\S	\S
\4S jrS\S	\S
\4S jrSrg)CosineDistanceiQ  TFrB   ignore_zerosrB   r   c                    Xl         X l        g r   r   )r    rB   r   s      r!   r"   CosineDistance.__init__R  s    "(r$   r%   r&   r   c                 F    U R                  X5      U R                  X5      4$ r   r(   r+   s      r!   r,   CosineDistance.__call__V  r.   r$   c                 R   UR                   UR                   :w  a(  SUR                    SUR                    S3n[        U5      e[        U5      nUS-   nUS-   nUR                  R	                  USSS9nUR                  R	                  USSS9nXx-  n	XV-  R                  SSS9U	-  n
U
$ )N7Cannot calculate cosine similarity: mismatched shapes: rm   rn   :0yE>r   Taxiskeepdims)r\   rV   r   linalgnormrx   )r    r%   r&   rq   r]   r   r   norm_yhnorm_y	mul_normscosines              r!   get_similarityCosineDistance.get_similarityY  s    ==FLL(KGMM?Z^_e_k_k^llmnCS/!g&t^TM))..!d.;D9$	&1t4y@r$   c                 F   UR                   UR                   :w  a(  SUR                    SUR                    S3n[        U5      e[        U5      nU R                  (       a!  UR	                  U5      R                  SS9S:H  nUS-   nUS-   nUR                  R                  USSS	9nUR                  R                  USSS	9n	X-  n
Xg-  R                  SSS	9U
-  nXz-  XUS
-  -  -  -
  nU R                  (       a  SUW'   U R                  (       a  XR                   S   -  nU* $ )Nr   rm   rn   r   r   r   r   Tr   rw   )	r\   rV   r   r   absrx   r   r   rB   )r    r%   r&   rq   r]   zero_indicesr   r   r   r   r   r   r   s                r!   r)   CosineDistance.get_gradh  s   ==FLL(KGMM?Z^_e_k_k^llmnCS/! g&66&>--1-5:Lt^TM))..!d.;D9$	&1t4y@&!O"<=!"D>>--**Dur$   c                    UR                   UR                   :w  a(  SUR                    SUR                    S3n[        U5      e[        U5      nU R                  X5      nUR	                  US-
  5      nU R
                  (       a%  UR	                  U5      R                  SS9S:H  nSXg'   U R                  (       a  XaR                   S   -  nUR                  5       nU$ )Nr   rm   rn   r   r   r   )r\   rV   r   r   r   r   rx   rB   )	r    r%   r&   rq   r]   r   lossesr   r   s	            r!   r*   CosineDistance.get_loss  s    ==FLL(KGMM?Z^_e_k_k^llmnCS/!g&$$W5
#66&>--1-5:L#$F >>mmA..Fzz|r$   )r   rB   N)r4   r5   r6   r7   r|   r"   r   r
   r}   r,   r   r)   r*   r9   r   r$   r!   r   r   Q  s    ,0u )T ) )N N( NuXu_?U Nh  U  ( x 4 ( u r$   r   zCosineDistance.v1Fr   r   c                     [        XS9$ )Nr   )r   r   s     r!   configure_CosineDistancer     s     IIIr$   c                 X    [        U 5      nUR                  U R                  SS9nSX1'   U$ )NrK   rL   r   )r   rP   r\   )r%   r^   r]   rc   s       r!   rY   rY     s.    	'	"B777==7,DDMKr$   c                     [        U5      nUR                  UR                  SS9nUb0  U R                  S:X  a	  SX@U:H  '   U$ UR	                  U SS9nSXEU:H  '   U$ )NrK   rL   r   r@   rN   r   )r   rP   r\   r[   argmax)r&   r%   r>   r]   rc   labelss         r!   rZ   rZ     sn    	'	"B777==7,D ;;!,/D=() K YYvBY/F,/D=()Kr$   )r   r;   r   r   )0abcr   typingr   r   r   r   r   r	   r
   r   r   r   configr   typesr   r   utilr   r   r   r   r   r   IntsOrFloatsrT   rz   r~   r   r;   r   r|   r   r   r}   r   r   r   r   r   r   r   r   r   rY   rZ   __all__r   r$   r!   <module>r      s       # 2				VX%&68Xc]HSMIJ 7665%/0 ,w"d w"t 
-. %)/3	 HSM" E#s(O,	
  / 
-. %)/3 $ HSM" E#s(O,	
   / 
-. %)/3 $  HSM" E#s(O,	
    /"2d 2j 
56MM&.x}&=M$M 7M 
56 %) $	 HSM" 	
 % 7 
56 %)/3 $  HSM" E#s(O,	
   % 7"" "2 
!.2 +t +z + "+@T @F 
$%EJJ-1JJ &JH 8  r$   