U
    4Af                     @   s^   d Z ddlZddlmZ ddlmZmZ eeZG dd deZ	d	e	iZ
d	d
iZd	diZdS )z+XNLI utils (dataset loading and evaluation)    N   )logging   )DataProcessorInputExamplec                   @   s2   e Zd ZdZdddZdd Zdd Zd	d
 ZdS )XnliProcessorz
    Processor for the XNLI dataset. Adapted from
    https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/run_classifier.py#L207
    Nc                 C   s   || _ || _d S )N)languagetrain_language)selfr   r	    r   E/tmp/pip-unpacked-wheel-zw5xktn0/transformers/data/processors/xnli.py__init__!   s    zXnliProcessor.__init__c              	   C   s   | j dkr| jn| j }| tj|d| d}g }t|D ]\}}|dkrPq>d| }|d }|d }	|d dkrzd	n|d }
t|tst	d
| dt|	tst	d
|	 dt|
tst	d|
 d|
t|||	|
d q>|S )See base class.Nz$XNLI-MT-1.0/multinli/multinli.train.z.tsvr   ztrain-r      ZcontradictorycontradictionTraining input  is not a stringTraining label guidtext_atext_blabel)r	   r   	_read_tsvospathjoin	enumerate
isinstancestr	TypeErrorappendr   )r
   data_dirZlglinesexamplesiliner   r   r   r   r   r   r   get_train_examples%   s$    



z XnliProcessor.get_train_examplesc              	   C   s   |  tj|d}g }t|D ]\}}|dkr2q |d }|| jkrFq d| }|d }|d }	|d }
t|tstd| dt|	tstd|	 dt|
tstd	|
 d|	t
|||	|
d
 q |S )r   zXNLI-1.0/xnli.test.tsvr   ztest-      r   r   r   r   r   )r   r   r   r   r   r   r   r   r    r!   r   )r
   r"   r#   r$   r%   r&   r   r   r   r   r   r   r   r   get_test_examples:   s(    




zXnliProcessor.get_test_examplesc                 C   s
   dddgS )r   r   Z
entailmentZneutralr   )r
   r   r   r   
get_labelsQ   s    zXnliProcessor.get_labels)N)__name__
__module____qualname____doc__r   r'   r*   r+   r   r   r   r   r      s
   
r   ZxnliZclassification)r/   r   utilsr   r   r   Z
get_loggerr,   loggerr   Zxnli_processorsZxnli_output_modesZxnli_tasks_num_labelsr   r   r   r   <module>   s   
<   