o
    *jw*                     @   s   d Z ddlZddlZddlZddlZddlmZ ddlmZ ddl	Z	ddl
mZ ddlmZ ddlmZ ddlmZmZ dd	lmZmZmZ d
d Zdd Zdd Zda					dddZdZdd ZdS )zEvaluation utilities.    N)OrderedDict)List)process_batch)mpuf1_score)InputExamplebuild_data_loader)debug_finetune_dataget_spare_portprint_rank_0c                 C   sR   d}t t| d}t| t|ksJ t| |D ]
\}}|||k7 }q|d | S )Nr      g      Y@)maxlenzip)predictionslabelsexamplescountZnum_predictions
predictionlabel r   l/var/www/html/Deteccion_Ine/venv/lib/python3.10/site-packages/modelscope/models/nlp/mglm/tasks/eval_utils.pyaccuracy_metric   s   r   c                 C   s
   t || S )Nr   r   r   r   r   r   r   	f1_metric(   s   
r   c                 C   s   t || ddS )Nmacro)Zaverager   r   r   r   r   f1_macro_metric,   s   r   FTc              	      s   |a |rtj rtj dkrdS r# js# jdur jndg}n jdur+ jndg}du r4tg  j	r< j	n j
}	|D ]}
| |
}t||	 jdd|d}|j|f qA		d	 fdd	}|S )
z,Provide function that calculates accuracies.r   NtestdevF)num_workersZ	drop_lastshuffle
only_rank0c                    s   t d ttrtdd D ndi}dD ]\}}d }t|jdr*|jj}t }| ||\ t | }	|rWtj	
 dkrWtjj|d }
	 |
 t} fdd	 D }d
|||}| D ])\}}|d||7 }|d ur|dkrstdkr|d| d| || qs|d|	7 }tdkrt | |D ]}||  || | 7  < q|7 qfdd	| D }d|}| D ] \}}|d||7 }|d ur|dkrs|d| || qt | |S )Nzcalculating metrics ...c                 S   s   g | ]}|d fqS )        r   ).0keyr   r   r   
<listcomp>Z   s    z@accuracy_func_provider.<locals>.metrics_func.<locals>.<listcomp>r#   r   r   z.jsonlc                    s   i | ]\}}|| qS r   r   )r$   r%   Zmetric)r   r   r   r   r   
<dictcomp>k       z@accuracy_func_provider.<locals>.metrics_func.<locals>.<dictcomp>z' > |epoch: {}| metrics for {}: total {}z {} = {:.4f} %r   zTrain/valid__z elapsed time (sec): {:.3f}c                    s   i | ]\}}||t   qS r   )float)r$   r%   score)totalr   r   r'   }   r(   z# >> |epoch: {}| overall: total = {}z {} = {:.4f})r   
isinstancedictr   hasattrdatasetr   timetorchdistributedget_rankospathjoinZlog_dirr   itemsformatZ
add_scalar)modelepochZoutput_predictionsZsummary_writerZ
score_dictname
dataloaderexample_dict
start_timeZelapsed_timefilenameZtotal_countZsingle_dictZ
output_strr%   valuer+   argsZdataloaders	eval_funcis_testmetric_dictoutput_func)r   r   r   r,   r   metrics_funcU   sl   


z,accuracy_func_provider.<locals>.metrics_func)FN)global_tokenizerr2   r3   Zis_initializedr4   Z
eval_validZ	test_dataZ
valid_datamultichoice_evaluateeval_batch_sizeZ
batch_sizer	   r    appendZdataset_name)Zsingle_dataset_providerrF   rC   rE   rD   rG   r"   	tokenizerZ	datapathsrK   datapathr0   r=   rH   r   rB   r   accuracy_func_provider3   s4   
4rO   
   c           '   
      sd  |    t|}td|  tj|j|tj tj dkt	j
dd}t  t|D ]\}}t||}|jrU|d |d |d |d f\}	}
}}|	|
|g}no|jr|d |d |d	 }	}}|d
 |d |d }}}|js|	||||g}|jr|d }|| n9|d |d |d }}}|d |d }}|	|||||||g}n|d |d |d	 |d
 f\}	}}}|	||g}t|d jdkr|d dtkrg }t|d dd t d D ]  fdd|D }|jr| | }n| | ^}}|| qtj|dd}nT|jrU|jrUg }|d d}t|d t d D ]! |dd  fdd|dd D  }| | ^}}|| q+tj|dd}n|jr^| | }n| | ^}}d|v rddlm} d|v rz||d  }|||d dd}nd|v r|d }|| dd|   }|d }t|tjr|   }tj |d d }| } |j!" d!krd"d |D }t#$ dkrt%||| D ]\}!}"}#|&|!t'|"|#f qq0W d   n	1 sw   Y  | (  tj)  g g g }$} }%|* D ]\}!}&t |+|!\}"}#|$|" | |# |%|& q
tj)  |$| |%fS )#zcCalculate correct over total answers and return prediction if the
    `output_predictions` is true.zUsing port r      )secondstexttypesr   Zpadding_maskpositionmasktarget
logit_mask
prompt_posZdec_textZdec_positionZdec_maskZ
dec_targetdec_logit_mask   r   c                    ,   g | ]}|d d  t   d t  f qS Nr   segment_lengthr$   argir   r   r&          z(multichoice_evaluate.<locals>.<listcomp>)dimNc                    r\   r]   r^   r`   rb   r   r   r&      rd   Z
segment_id)scatter_sum	loss_maskg     @g      ?uidZwscc                 S   s   g | ]
}|d kr
dnd qS )r   r   r   )r$   predr   r   r   r&      s    ),evalr   r   r2   r3   ZTCPStoreZ	master_ipZget_world_sizer4   datetime	timedeltaZno_grad	enumerater   Zpretrained_bertZ
cloze_evalZfast_decodeZcontinuous_promptrL   r   shapesizer_   rangecatZtorch_scatterrf   r-   ZTensorcpunumpytolistZargmaxtasklowerr   Zget_model_parallel_rankr   setstrtrainZbarrierr8   get)'r:   r=   r>   rC   portstorer)   batchdatatokensrT   Zlabels_Zattention_maskZinputsZposition_idsZ
target_idsrX   rY   Zdec_input_idsZdec_position_idsZdec_attention_maskZdec_target_idsrZ   Z
logit_listZinput_batchZlogitsZmemsZnum_choicesrf   rg   Zuid_listZ	predictedr   rh   r   r   r   r   Zexampler   rb   r   rJ      s   






"






T




rJ   )FNNTN)__doc__rl   r5   randomr1   collectionsr   typingr   r2   Zfinetune_glmr   Zmegatron_utilr   Zsklearn.metricsr   Ztasks.data_utilsr   r	   utilsr
   r   r   r   r   r   rI   rO   r_   rJ   r   r   r   r   <module>   s2   	
Y