o .wi @sddlmZddlmZmZmZddlZddlmZddlm Z  ddee ee fdee ee fd e d efd d Z ddedeee fdeedd efddZ  ddee ee fdee ee fd e deedd ef ddZdS))Sequence)LiteralOptionalUnionN)Tensor)_LevenshteinEditDistancepredstargetsubstitution_costreturncst|tr|g}t|tr|g}tdd|Ds td|tdd|Ds0td|t|t|krFtdt|dt|fdd t||D}tj|tjd S) Ncs|]}t|tVqdSN isinstancestr.0xr^/home/ubuntu/sommelier/.venv/lib/python3.10/site-packages/torchmetrics/functional/text/edit.py z(_edit_distance_update..zCExpected all values in argument `preds` to be string type, but got csr rrrrrrr"rzDExpected all values in argument `target` to be string type, but got zDExpected argument `preds` and `target` to have same length, but got z and cs$g|]\}}t|d|dqS)) op_substituter) _LE_distance)rptr rr )sz)_edit_distance_update..dtype) rrall ValueErrorlenziptorchtensorint)r r r distancerrr_edit_distance_updates   r)mean edit_scores num_elements reduction)r*sumnonecCs\|dkrtjdtjdS|dkr||S|dkr |S|dus(|dkr*|Std)z3Compute final edit distance reduced over the batch.rrr*r.Nr/zHExpected argument `reduction` to either be 'sum', 'mean', 'none' or None)numelr%r&int32r.r")r+r,r-rrr_edit_distance_compute0s  r2cCst|||}t|||dS)uCalculates the Levenshtein edit distance between two sequences. The edit distance is the number of characters that need to be substituted, inserted, or deleted, to transform the predicted text into the reference text. The lower the distance, the more accurate the model is considered to be. Implementation is similar to `nltk.edit_distance `_. Args: preds: An iterable of predicted texts (strings). target: An iterable of reference texts (strings). substitution_cost: The cost of substituting one character for another. reduction: a method to reduce metric score over samples. - ``'mean'``: takes the mean over samples - ``'sum'``: takes the sum over samples - ``None`` or ``'none'``: return the score per sample Raises: ValueError: If ``preds`` and ``target`` do not have the same length. ValueError: If ``preds`` or ``target`` contain non-string values. Example:: Basic example with two strings. Going from “rain” -> “sain” -> “shin” -> “shine” takes 3 edits: >>> from torchmetrics.functional.text import edit_distance >>> edit_distance(["rain"], ["shine"]) tensor(3.) Example:: Basic example with two strings and substitution cost of 2. Going from “rain” -> “sain” -> “shin” -> “shine” takes 3 edits, where two of them are substitutions: >>> from torchmetrics.functional.text import edit_distance >>> edit_distance(["rain"], ["shine"], substitution_cost=2) tensor(5.) Example:: Multiple strings example: >>> from torchmetrics.functional.text import edit_distance >>> edit_distance(["rain", "lnaguaeg"], ["shine", "language"], reduction=None) tensor([3, 4], dtype=torch.int32) >>> edit_distance(["rain", "lnaguaeg"], ["shine", "language"], reduction="mean") tensor(3.5000) )r,r-)r)r2r0)r r r r-r(rrr edit_distanceAs 6r3)r)r*)rr*)collections.abcrtypingrrrr%r#torchmetrics.functional.text.helperrrrr'r)r2r3rrrrsL