wip.py 5.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139
  1. # Copyright The Lightning team.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from collections.abc import Sequence
  15. from typing import Any, Optional, Union
  16. from torch import Tensor, tensor
  17. from torchmetrics.functional.text.wip import _wip_compute, _wip_update
  18. from torchmetrics.metric import Metric
  19. from torchmetrics.utilities.imports import _MATPLOTLIB_AVAILABLE
  20. from torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE
  21. if not _MATPLOTLIB_AVAILABLE:
  22. __doctest_skip__ = ["WordInfoPreserved.plot"]
  23. class WordInfoPreserved(Metric):
  24. r"""Word Information Preserved (`WIP`_) is a metric of the performance of an automatic speech recognition system.
  25. This value indicates the percentage of words that were correctly predicted between a set of ground-
  26. truth sentences and a set of hypothesis sentences. The higher the value, the better the performance of the ASR
  27. system with a WordInfoPreserved of 1 being a perfect score. Word Information Preserved rate can then be
  28. computed as:
  29. .. math::
  30. wip = \frac{C}{N} * \frac{C}{P}
  31. where:
  32. - :math:`C` is the number of correct words,
  33. - :math:`N` is the number of words in the reference
  34. - :math:`P` is the number of words in the prediction
  35. As input to ``forward`` and ``update`` the metric accepts the following input:
  36. - ``preds`` (:class:`~List`): Transcription(s) to score as a string or list of strings
  37. - ``target`` (:class:`~List`): Reference(s) for each speech input as a string or list of strings
  38. As output of ``forward`` and ``compute`` the metric returns the following output:
  39. - ``wip`` (:class:`~torch.Tensor`): A tensor with the Word Information Preserved score
  40. Args:
  41. kwargs: Additional keyword arguments, see :ref:`Metric kwargs` for more info.
  42. Examples:
  43. >>> from torchmetrics.text import WordInfoPreserved
  44. >>> preds = ["this is the prediction", "there is an other sample"]
  45. >>> target = ["this is the reference", "there is another one"]
  46. >>> wip = WordInfoPreserved()
  47. >>> wip(preds, target)
  48. tensor(0.3472)
  49. """
  50. is_differentiable: bool = False
  51. higher_is_better: bool = False
  52. full_state_update: bool = False
  53. plot_lower_bound: float = 0.0
  54. plot_upper_bound: float = 1.0
  55. errors: Tensor
  56. preds_total: Tensor
  57. target_total: Tensor
  58. def __init__(
  59. self,
  60. **kwargs: Any,
  61. ) -> None:
  62. super().__init__(**kwargs)
  63. self.add_state("errors", tensor(0.0), dist_reduce_fx="sum")
  64. self.add_state("target_total", tensor(0.0), dist_reduce_fx="sum")
  65. self.add_state("preds_total", tensor(0.0), dist_reduce_fx="sum")
  66. def update(self, preds: Union[str, list[str]], target: Union[str, list[str]]) -> None:
  67. """Update state with predictions and targets."""
  68. errors, target_total, preds_total = _wip_update(preds, target)
  69. self.errors += errors
  70. self.target_total += target_total
  71. self.preds_total += preds_total
  72. def compute(self) -> Tensor:
  73. """Calculate the Word Information Preserved."""
  74. return _wip_compute(self.errors, self.target_total, self.preds_total)
  75. def plot(
  76. self, val: Optional[Union[Tensor, Sequence[Tensor]]] = None, ax: Optional[_AX_TYPE] = None
  77. ) -> _PLOT_OUT_TYPE:
  78. """Plot a single or multiple values from the metric.
  79. Args:
  80. val: Either a single result from calling `metric.forward` or `metric.compute` or a list of these results.
  81. If no value is provided, will automatically call `metric.compute` and plot that result.
  82. ax: An matplotlib axis object. If provided will add plot to that axis
  83. Returns:
  84. Figure and Axes object
  85. Raises:
  86. ModuleNotFoundError:
  87. If `matplotlib` is not installed
  88. .. plot::
  89. :scale: 75
  90. >>> # Example plotting a single value
  91. >>> from torchmetrics.text import WordInfoPreserved
  92. >>> metric = WordInfoPreserved()
  93. >>> preds = ["this is the prediction", "there is an other sample"]
  94. >>> target = ["this is the reference", "there is another one"]
  95. >>> metric.update(preds, target)
  96. >>> fig_, ax_ = metric.plot()
  97. .. plot::
  98. :scale: 75
  99. >>> # Example plotting multiple values
  100. >>> from torchmetrics.text import WordInfoPreserved
  101. >>> metric = WordInfoPreserved()
  102. >>> preds = ["this is the prediction", "there is an other sample"]
  103. >>> target = ["this is the reference", "there is another one"]
  104. >>> values = [ ]
  105. >>> for _ in range(10):
  106. ... values.append(metric(preds, target))
  107. >>> fig_, ax_ = metric.plot(values)
  108. """
  109. return self._plot(val, ax)