summaryrefslogtreecommitdiff
path: root/text_recognizer/models/metrics/cer.py
diff options
context:
space:
mode:
Diffstat (limited to 'text_recognizer/models/metrics/cer.py')
-rw-r--r--text_recognizer/models/metrics/cer.py23
1 files changed, 23 insertions, 0 deletions
diff --git a/text_recognizer/models/metrics/cer.py b/text_recognizer/models/metrics/cer.py
new file mode 100644
index 0000000..238ecc3
--- /dev/null
+++ b/text_recognizer/models/metrics/cer.py
@@ -0,0 +1,23 @@
+"""Character Error Rate (CER)."""
+from typing import Sequence
+
+import torch
+import torchmetrics
+
+
+class CharacterErrorRate(torchmetrics.CharErrorRate):
+ """Character error rate metric, allowing for tokens to be ignored."""
+
+ def __init__(self, ignore_tokens: Sequence[int], *args):
+ super().__init__(*args)
+ self.ignore_tokens = set(ignore_tokens)
+
+ def update(self, preds: torch.Tensor, targets: torch.Tensor) -> None:
+ preds_l = [
+ [t for t in pred if t not in self.ignore_tokens] for pred in preds.tolist()
+ ]
+ targets_l = [
+ [t for t in target if t not in self.ignore_tokens]
+ for target in targets.tolist()
+ ]
+ super().update(preds_l, targets_l)