o .wilB @s&ddlZddlmZddlmZmZddlmZdZdZ e dZ eGdd d e eZ Gd d d Zd eee eee fd ee ee fdeeee ee ffddZdee dee de fddZdee dfdee dffddZdee dfdeee e fee ee ffddZdS)N)Sequence)Enumunique)Unioni'g7yACc@s$eZdZdZdZdZdZdZdZdS)_EditOperationsz1Enumerations for the Levenhstein edit operations.insertdelete substitutenothing undefinedN) __name__ __module__ __qualname____doc__ OP_INSERT OP_DELETE OP_SUBSTITUTE OP_NOTHING OP_UNDEFINEDrr`/home/ubuntu/sommelier/.venv/lib/python3.10/site-packages/torchmetrics/functional/text/helper.pyr,src@sneZdZdZ d deededededdf d d Zd eedeeee d fffd dZ d eededeeeee fdeeeeeee fee d fffddZ dedeeeee fdee d ffddZ d eedeeeee fddfddZ d eedeeeeeee fffddZdedeeee ffddZdedeeee ffddZdS)!_LevenshteinEditDistancea.A convenience class for calculating the Levenshtein edit distance. Class will cache some intermediate values to hasten the calculation. The implementation follows the implementation from https://github.com/mjpost/sacrebleu/blob/master/sacrebleu/metrics/lib_ter.py, where the most of this implementation is adapted and copied from. Args: reference_tokens: list of reference tokens op_insert: cost of insertion operation op_delete: cost of deletion operation op_substitute: cost of substitution operation reference_tokens op_insert op_delete op_substitutereturnNcCs>||_t||_i|_d|_||_||_||_d|_t |_ dS)Nr) rlen reference_lencache cache_sizerrr op_nothing _INT_INFINITY op_undefined)selfrrrrrrr__init__Fs  z!_LevenshteinEditDistance.__init__prediction_tokens.cCs6||\}}||||\}}}|||||fS)a(Calculate edit distance between self._words_ref and the hypothesis. Uses cache to skip some computations. Args: prediction_tokens: A tokenized predicted sentence. Return: A tuple of a calculated edit distance and a trace of executed operations. ) _find_cache_levenshtein_edit_distance _add_cache)r&r(start_positioncached_edit_distanceedit_distance_int edit_distancetracerrr__call__Us   z!_LevenshteinEditDistance.__call__prediction_startr!cst|}fddt||D}||}|rj|nd}|dtkr-t|dtnt}t|d|dD]} t| |} td| |} | |krQjdn tjd| |} t| | D]}} | dkr{|| d| dj t j f|| | <q`|| dj | dkrj }t j}nj}t j}|| d| dd||f|| d| dj t j f|| | ddjt jff}|D]\}}|| | d|kr||f|| | <qq`q8||}|ddd|t|d|fS) aDynamic programming algorithm to compute the Levenhstein edit distance. Args: prediction_tokens: A tokenized predicted sentence. prediction_start: An index where a predicted sentence to be considered from. cache: A cached Levenshtein edit distance. Returns: Edit distance between the predicted sentence and the reference sentence csg|] }tjqSr)list_get_empty_rowr .0_r&rr }szG_LevenshteinEditDistance._levenshtein_edit_distance..g?rrN)rranger _BEAM_WIDTHmathceilfloormaxminrrrrr#rrrrr _get_trace)r&r(r2r!prediction_len empty_rowsr/ length_ratio beam_widthi pseudo_diagmin_jmax_jjcost_substituteoperation_substitute operationsoperation_costoperation_namer0rr8rr*jsD  "$  "z3_LevenshteinEditDistance._levenshtein_edit_distancerDr/cCsd}|}|j}|dks|dkrQ|||d}|g|R}|tjtjfvr.|d8}|d8}n|tjkr8|d8}n|tjkrB|d8}ntd||dks|dks|S)a@Get a trace of executed operations from the edit distance matrix. Args: prediction_len: A length of a tokenized predicted sentence. edit_distance: A matrix of the Levenshtedin edit distance. The element part of the matrix is a tuple of an edit operation cost and an edit operation itself. Return: A trace of executed operations returned as a tuple of `_EDIT_OPERATIONS` enumerates. Raises: ValueError: If an unknown operation has been applied. rrrUnknown operation )r rrrrr ValueError)r&rDr/r0rHrL operationrrrrCs        z#_LevenshteinEditDistance._get_tracec Cs|jtkrdS|j}t|t|}t|D] }|||d}qt||d|D]\}}||vrAit|f||<|jd7_||}|d}q*dS)a,Add newly computed rows to cache. Since edit distance is only calculated on the hypothesis suffix that was not in cache, the number of rows in `edit_distance` matrx may be shorter than hypothesis length. In that case we skip over these initial words. Args: prediction_tokens: A tokenized predicted sentence. edit_distance: A matrix of the Levenshtedin edit distance. The element part of the matrix is a tuple of an edit operation cost and an edit operation itself. Nrr)r"_MAX_CACHE_SIZEr!rr<ziptuple) r&r(r/nodeskip_numrHwordrowvaluerrrr+s   z#_LevenshteinEditDistance._add_cachecCsZ|j}d}||jg}|D]}||vr$|d7}||\}}||q||fS||fS)a+Find the already calculated rows of the Levenshtein edit distance metric. Args: prediction_tokens: A tokenized predicted sentence. Return: A tuple of a start hypothesis position and `edit_distance` matrix. prediction_start: An index where a predicted sentence to be considered from. edit_distance: A matrix of the cached Levenshtedin edit distance. The element part of the matrix is a tuple of an edit operation cost and an edit operation itself. rr)r!_get_initial_rowr append)r&r(rXr,r/rZr[rrrr)s  z$_LevenshteinEditDistance._find_cachelengthcCst|jtjfg|dS)aPrecomputed empty matrix row for Levenhstein edit distance. Args: length: A length of a tokenized sentence. Return: A list of tuples containing infinite edit operation costs and yet undefined edit operations. r)intr%rrr&r_rrrr4 z'_LevenshteinEditDistance._get_empty_rowcsfddt|dDS)a*First row corresponds to insertion operations of the reference, so 1 edit operation per reference word. Args: length: A length of a tokenized sentence. Return: A list of tuples containing edit operation costs of insert and insert edit operations. csg|] }|jtjfqSr)rrr)r6rHr8rrr9'sz=_LevenshteinEditDistance._get_initial_row..r)r<rarr8rr]rbz)_LevenshteinEditDistance._get_initial_row)rrr)r rrrr3strr`r'rWrr1r*rCr+r)r4r]rrrrr7sF &$ D  *&.!" r ref_corpushypothesis_corpusrcCst|tr|g}tdd|Dr!t|dkr|gndd|D}|rBtdd|DrBt|t|krBtdt|dt|||fS) aCheck and update (if needed) the format of reference and hypothesis corpora for various text evaluation metrics. Args: ref_corpus: An iterable of iterables of reference corpus. hypothesis_corpus: An iterable of hypothesis corpus. Return: ref_corpus: An iterable of iterables of reference corpus. hypothesis_corpus: An iterable of hypothesis corpus. Raises: ValueError: If length of `ref_corpus` and `hypothesis_corpus` differs. css|]}t|tVqdSN) isinstancercr6refrrr Az#_validate_inputs..rcSsg|]}|gqSrrrhrrrr9Bsz$_validate_inputs..css|]}|VqdSrfrrhrrrrjDszCorpus has different size z != )rgrcallrrS)rdrerrr_validate_inputs*s  &rmr(rcsfddtt|dD}tt|dD]}|||d<qttdD]}||d|<q(tdt|dD]I}tdtdD]=}||d|dkrb||d|d|||<qEt||d||||d||d|dd|||<qEq:|ddS)aDynamic programming algorithm to compute the edit distance. Args: prediction_tokens: A tokenized predicted sentence reference_tokens: A tokenized reference sentence Returns: Edit distance between the predicted sentence and the reference sentence csg|] }dgtdqS)rr)rr5rrrr9Tsz"_edit_distance..rrr;)r<rrB)r(rdprHrLrrnr_edit_distanceJs B rpr0.csJtjtjtjtjidtdtttfdtfddtfdd|DS)aFlip the trace of edit operations. Instead of rewriting a->b, get a recipe for rewriting b->a. Simply flips insertions and deletions. Args: trace: A tuple of edit operations. Return: inverted_trace: A tuple of inverted edit operations. rT_flip_operationsrcSs||vr ||S|Srf)get)rTrqrrr_replace_operation_or_retaints z1_flip_trace.._replace_operation_or_retainc3s|]}|VqdSrfr)r6rTrqrsrrrj{rkz_flip_trace..)rrrdictrW)r0rrtr _flip_tracebs  rvcCsd}}g}g}i}|D]c}|tjkr*|d7}|d7}|||<|d|dq |tjkrF|d7}|d7}|||<|d|dq |tjkrU|d7}|dq |tjkrh|d7}|||<|dq td|d|||fS)aTransform trace of edit operations into an alignment of the sequences. Args: trace: A trace of edit operations as a tuple of `_EDIT_OPERATIONS` enumerates. Return: alignments: A dictionary mapping aligned positions between a reference and a hypothesis. reference_errors: A list of error positions in a reference. hypothesis_errors: A list of error positions in a hypothesis. Raises: ValueError: If an unknown operation is r;rrrR.)rrr^rrrrS)r0reference_positionhypothesis_positionreference_errorshypothesis_errors alignmentsrTrrr_trace_to_alignment~s4           r})r>collections.abcrenumrrtypingrr=rUr`r$rcrrrWrmr3rprvrur}rrrrs*   t  "8