o
    6t¾iµ@  ã                
   @   st  d dl Z d dlmZ d dlmZ e  d¡Zdd„ Zdd„ Zd	d
„ Z	dZ
dd„ Zdd„ Zdd„ Zdd„ ZdZdZdZdd ddddedddiZd ddddœZd ed!efd"d#„Zd ed!efd$d%„Zd&d'„ Zd(d)„ Zd*d+„ Zd,ed!ee fd-d.„Zd,ed!ee fd/d0„Zd,ed!ee fd1d2„Zd3ZdZ d4Z!d5d6„ Z"d ed!efd7d8„Z#d9d:„ Z$d;d<„ Z%d=d>„ Z&d?Z'd@dA„ Z(dBdC„ Z)dDdE„ Z*dS )Fé    N)Úload)Úf1_scorezlm-evalc                 C   óf   | d }| d }|dkr|dkrdS |dkr|dkrdS |dkr&|dkr&dS |dkr0|dkr0dS 	 d	S ©
zP
    Function to extract the target from the dataset for sentiment analysis
    ÚoposÚonegÚ1Ú0r   é   é   é   N© ©Úxr   r   r   r   úS/home/ubuntu/.local/lib/python3.10/site-packages/lm_eval/tasks/evalita_llm/utils.pyÚsa_doc_to_target   ó   r   c                 C   r   r   r   r   r   r   r   Úsa_doc_to_target_v2   r   r   c                 C   s   g d¢S )zP
    Function to return the choices from the dataset for sentiment analysis
    )ÚPositivoÚNegativoÚNeutraleÚMistor   ©r   r   r   r   Úsa_doc_to_choice1   s   r   z	&&NOSYN&&c                 C   sn   t | d ƒdkr
tS d}| d D ]}||d d t|d ƒ d 7 }qt |ƒdkr5|d d	kr5|d
d… }|S )ú=
    Generate the target for the lexical similarity task
    Úanswersr   Ú Úwordú$$Úcountú::éþÿÿÿú:N)ÚlenÚNO_SYN_STRINGÚstr©r   Úans_strÚir   r   r   Ú_ls_gold_to_target<   s   "r)   c                 C   s^   t | d ƒdkr
tS d}| d D ]
}||d d 7 }qt |ƒdkr-|d dkr-|dd… }|S )	r   r   r   r   r   z, r!   ú,N)r#   r$   r&   r   r   r   Úls_doc_to_targetM   s   r+   c              	   C   s   | t krg g fS |  d¡}g }g }t|ƒdkrD|D ]*}d|v rC| d¡\}}| |¡ z	| t|ƒ¡ W q tyB   | d¡ Y qw q||fS )z5
    Split the gold string into a list of tuples
    r    r   r   )r$   Úsplitr#   ÚappendÚintÚ
ValueError)r   r   ÚwordsÚfreqsÚar   r   r   r   r   Ú_ls_split_gold[   s"   

ÿûr3   c                 C   sð   t | ƒ}t|ƒ\}}d}t|d tƒ}|dtdt|ƒƒ… }|r.d dd„ |d D ƒ¡|d< t|ƒdkr6dnd}t|ƒdkr@dnd}g }|D ]}	|	|v r^| |	 ¡ ¡}
|||
 7 }| 	|	¡ qFt
|ƒ}|dkrl|| }nt d	¡ d
|||fiS )zv
    Process the results of the evaluation for the lexical substitution task
    look at coqa for another example
    r   Né
   r   c                 s   s    | ]	}|  ¡ r|V  qd S )N)Úisalpha)Ú.0Úcharr   r   r   Ú	<genexpr>   s   € z%ls_process_results.<locals>.<genexpr>éÿÿÿÿr
   zH_i is 0Úf1)r)   r3   Úsplit_text_with_regexÚLS_SPLIT_REGEXÚminr#   ÚjoinÚindexÚstripr-   ÚsumÚeval_loggerÚdebug)ÚdocÚresultsÚgold_to_targetr0   r1   ÚprecÚhas_answÚhas_annotationÚmatching_resÚrÚidxÚHir   r   r   Úls_process_resultsp   s*   
€

rN   z	&&NOENT&&r*   ú$ÚPERÚLOCr
   ÚORGr   r   ÚOé   )rP   rQ   rR   rS   r   Úreturnc                 C   ó   dd„ | D ƒ}|S )úU
    Convert the gold entities to the target format according to the NER_MAPPING
    c                 S   ó   g | ]}t |d   ‘qS ©Útype©ÚNER_MAPPING©r6   Úer   r   r   Ú
<listcomp>§   ó    z'_ner_gold_to_target.<locals>.<listcomp>r   ©r   Úresr   r   r   Ú_ner_gold_to_target£   ó   rc   c                 C   rV   )rW   c                 S   rX   rY   r[   r]   r   r   r   r_   ¯   r`   z*_ner_gold_to_target_v2.<locals>.<listcomp>r   ra   r   r   r   Ú_ner_gold_to_target_v2«   rd   re   c                 C   óJ   | d }d}|g krt S |D ]}||d t |d  t 7 }q|d d… S )NÚentitiesr   Úentity_textrZ   r9   )ÚNO_ENT_STRINGÚNER_TYPE_SEPARATORÚNER_ENTITY_SEPARATOR©rD   ÚentsÚtarg_strr^   r   r   r   Úner_doc_to_target³   s   ÿro   c                 C   sÊ  | d }|d }t |ƒ}t|ƒ}dgt|ƒ }g }t|ƒt|ƒkrj|D ].}|d }|d }	tt|ƒƒD ]}
|||
 d krP|	||
 d krPt|	 ||
< | |
¡ q3q#tt|ƒt|ƒ ƒD ]}
| d¡ | d¡ q\nkt|ƒdkr|t|ƒdkr|dg}|}nY|D ].}|d }|d }	tt|ƒƒD ]}
|||
 d kr«|	||
 d kr«t|	 ||
< | |
¡ qŽq~tt|ƒƒD ]!}
|
|v rºq³||
 dkrÅd||
< q³||
 dkrÐd||
< q³d||
< q³t|ƒt|ƒksßJ ‚d||fiS )	zB
    Process the results of the Named Entity Recognition task
    rg   r   r
   rh   rZ   r   r   r:   )Ú_ner_process_raw_outputrc   r#   Úranger\   r-   ©rD   rE   ÚgoldÚraw_resultsÚgold_labelsÚ
res_labelsÚmatched_gold_idxrK   Úr_textÚr_typer(   r   r   r   Úner_process_resultsÁ   sT    
€ý
þ 
€ý


rz   c                 C   sî  | d }|d }t |ƒ}t|ƒ}dgt|ƒ }g }t|ƒt|ƒkrj|D ].}|d }|d }	tt|ƒƒD ]}
|||
 d krP|	||
 d krPt|	 ||
< | |
¡ q3q#tt|ƒt|ƒ ƒD ]}
| d¡ | d¡ q\n}t|ƒdkr|t|ƒdkr|dg}|}nkt|ƒdkrŽ|d tkrŽdg}|}nY|D ].}|d }|d }	tt|ƒƒD ]}
|||
 d kr½|	||
 d kr½t|	 ||
< | |
¡ q qtt|ƒƒD ]!}
|
|v rÌqÅ||
 dkr×d||
< qÅ||
 dkrâd||
< qÅd||
< qÅt|ƒt|ƒksñJ ‚d||fiS )	z¥
    Process the results of the Named Entity Recognition task
    This version considers and score explicitly when the model responds that there are no entities
    rg   r   r
   rh   rZ   rT   r   r:   )Ú_ner_process_raw_output_v2re   r#   rq   r\   r-   ri   rr   r   r   r   Úner_process_results_v2ö   sZ    
€ý
ü 
€ý


r|   Ú
llm_resultc                 C   s”   t | v rg S | dkrdgS |  t¡}g }|D ]1}| ¡ }d}d}| t¡}t|ƒdk r2|d }d}n|d }|d }|dkrG| || ¡ f¡ q|S ©Nr   ÚWRONGr   r   r
   ©ri   r,   rk   r@   rj   r#   r-   Úupper©r}   Útmp_resultsrE   rb   rK   rx   ry   Ú
r_splittedr   r   r   rp   9  s(   

€rp   c                 C   s–   t | v rt gS | dkrdgS |  t¡}g }|D ]1}| ¡ }d}d}| t¡}t|ƒdk r3|d }d}n|d }|d }|dkrH| || ¡ f¡ q|S r~   r€   r‚   r   r   r   r{   Q  s(   

€r{   c                 C   s    t | v rg S | dkrdgS |  t¡}g }|D ]-}d}d}| t¡}t|ƒdk r0|d  ¡ }d}n|d  ¡ }|d  ¡ }| ||f¡ qt|ƒt|ƒksNJ ‚|S r~   )ÚNO_REL_STRINGr,   ÚINTER_REL_SEPARATORÚINTRA_REL_SEPARATORr#   r@   r-   )r}   rƒ   Ú	relationsrb   Úr_text1Úr_text2r„   r   r   r   Ú_rel_process_raw_outputl  s$   

r‹   ú%z	&&NOREL&&c                 C   rf   )Nrˆ   r   r   r
   r9   )ri   r‡   r†   rl   r   r   r   Úre_doc_to_target‡  ó   r   c                 C   s   | g krdgS dgt | ƒ S )Nr   r
   )r#   r   r   r   r   Ú_rel_gold_to_target“  s   r   c                 C   sJ   | d }d}|g krt S |D ]}||d d |d  d 7 }q|d d… S )Nrˆ   r   r   rO   r
   rŒ   r9   )r…   )rD   Úrelrn   rK   r   r   r   Úrel_doc_to_targetš  rŽ   r‘   c                 C   sp   g }| D ]'}d}d}|  t¡}t|ƒdk r|d }d}n|d }|d }| ||f¡ qt|ƒt| ƒks6J ‚|S )Nr   r   r   r
   )r,   r‡   r#   r-   )rE   rˆ   rK   r‰   rŠ   r„   r   r   r   Ú_extract_relations¦  s   
r’   c                 C   s¾  | d }|d }t |v rdnd}|g krdnd}g }g }|dkr.|r.t|ƒ}dgt|ƒ }n¡|dkr=|dkr=t|ƒ}|}n’|rO|dkrOt|ƒ}dgt|ƒ }n€t|ƒ}t|ƒ}dgt|ƒ }t|ƒdksfJ ‚tt|ƒƒD ]:}tt|ƒƒD ]1}	||	 d }
||	 d }|
|| d kr›||| d kr›d||< d||	< qt|
dkr¤|dkr¤qt	 qtqlt|ƒt|ƒ dkrÏtt|ƒt|ƒ ƒD ]}|| dkrÄq»| d¡ | d¡ q»t|ƒt|ƒksÙJ ‚d||fiS )zq
    Process the results of the Relation extraction task not considering the order of the relation extracted
    rˆ   r   r
   )ÚDELETEDr“   r“   r:   )r…   r   r#   r‹   rq   r-   )rD   rE   rs   rt   Úhas_resultsÚhas_goldrv   ru   r(   Újr‰   rŠ   r   r   r   Úrel_process_results_v3·  sL    
ö
r—   ú[^,]+c                 C   s€   ddl }|  d¡d } | || ¡}g }|D ]!}| d¡r1| d¡r1| d|dd… ¡}| |¡ q| |¡ qdd	„ |D ƒ}|S )
z\
    pattern: str - a regex pattern to match the text
    text: str - the text to split
    r   NÚ
ú"r˜   r
   r9   c                 S   s   g | ]
}|  ¡  d d¡‘qS )rš   r   )r@   Úreplace)r6   Úelementr   r   r   r_   	  s    z)split_text_with_regex.<locals>.<listcomp>)Úrer,   ÚfindallÚ
startswithÚendswithÚextendr-   )ÚtextÚpatternr   ÚmatchesÚresultÚmatchÚinner_matchesr   r   r   r;   ò  s   r;   c                 C   sN   | d dkrdS | d dkrdS | d dkrdS | d dkr d	S t  d
¡ d S )NÚcorrect_answerÚAr   ÚBr
   ÚCr   ÚDr   z@WARNING: correct answer not found or not in ["A", "B", "C", "D"]©rB   Úwarningr   r   r   r   Úfaq_doc_to_target  s   ÿr¯   c                 C   s.   | d dkrdS | d dkrdS t  d¡ d S )NÚsourceÚ
ilgiornaler   Ú
repubblicar
   z@WARNING: source not found or not in ["ilgiornale", "repubblica"]r­   r   r   r   r   Úht_doc_to_target  s   ÿr³   )+ÚloggingÚevaluater   Úsklearn.metricsr   Ú	getLoggerrB   r   r   r   r$   r)   r+   r3   rN   ri   rk   rj   ÚNER_MAPPING_V2r\   Úlistrc   re   ro   rz   r|   r%   Útuplerp   r{   r‹   r†   r‡   r…   r   r   r‘   r’   r—   r<   r;   r¯   r³   r   r   r   r   Ú<module>   sJ    
,5C8