
    uj"                         d dl Z d dlmZmZ d dlZd dlZd dlmc m	Z
 d dlmZ d dlmZ d dlmZ ddlmZ ddlmZmZ  ej        eej        	           G d
 de                      ZdS )    N)DictUnion)Metrics)
OutputKeys)default_group   )Metric)METRICS
MetricKeys)	group_keymodule_namec                        e Zd ZdZ fdZdedefdZedej	        dej	        de
fd	            Zed
eej        ej	        f         defd            Zde
fdZdeee
f         fdZddZd Zd Z xZS )	PplMetricz|The metric computation class for any classes.

    This metric class calculates perplexity for the whole input batches.
    c                 V     t                      j        |i | d| _        d| _        d S )Ng        r   )super__init__avg_loss	batch_num)selfargskwargs	__class__s      m/lsinfo/ai/hellotax_ai/data_center/backend/venv/lib/python3.11/site-packages/modelscope/metrics/ppl_metric.pyr   zPplMetric.__init__   s1    $)&)))!    outputsinputsc                    |t           j                 }|t           j                 }|                     ||          }|                     |t           j                           }|                     ||          | _        | xj        |z  c_        d S N)r   LOGITSLABELS	_get_loss_get_batch_num_average_lossr   r   )r   r   r   logitslabelsin_lossin_batch_nums          r   addzPplMetric.add   su    *+
)*..00**6*2C+DEE**7LAA,&r   r$   r%   returnc                     |                     d          }|                      |j        d         d          } t          j        | |                                          S )Nr   )viewshapeFcross_entropyitem)r$   r%   s     r   r!   zPplMetric._get_loss'   sH    RV\!_b11vv..33555r   matrixc                     | j         d         S )Nr   )r-   )r1   s    r   r"   zPplMetric._get_batch_num-   s    |Ar   r&   c                 B    | j         | j        z  ||z  z   | j        |z   z  S r   r   r   )r   r&   r'   s      r   r#   zPplMetric._average_loss1   s,    .<1GG~,. 	.r   c                 L    t           j        t          j        | j                  iS r   )r   PPLmathexpr   r   s    r   evaluatezPplMetric.evaluate5   s     7 788r   otherc                 z    |                      |j        |j                  | _        | xj        |j        z  c_        d S r   )r#   r   r   )r   r;   s     r   mergezPplMetric.merge8   s3    **5>5?KK%/)r   c                     | j         | j        fS r   r4   r9   s    r   __getstate__zPplMetric.__getstate__<   s    }dn,,r   c                 L    |                                   |\  | _        | _        d S r   )r   r   r   )r   states     r   __setstate__zPplMetric.__setstate__?   s!    (-%t~~~r   )r;   r   )__name__
__module____qualname____doc__r   r   r(   staticmethodtorchTensorfloatr!   r   npndarrayintr"   r#   strr:   r=   r?   rB   __classcell__)r   s   @r   r   r      sJ        
         
'4 ' ' ' ' ' 6%, 6 6 6 6 6 \6
 uRZ%=> 3    \.U . . . .9$sEz* 9 9 9 9* * * *- - -. . . . . . .r   r   )r7   typingr   r   numpyrK   rH   torch.nn.functionalnn
functionalr.   modelscope.metainfor   modelscope.outputsr   modelscope.utils.registryr   baser	   builderr
   r   register_moduler6   r    r   r   <module>r\      s                          ' ' ' ' ' ' ) ) ) ) ) ) 3 3 3 3 3 3       ( ( ( ( ( ( ( ( =gkJJJ/. /. /. /. /. /. /. KJ/. /. /.r   