murtazadahmardeh commited on
Commit
6b913da
1 Parent(s): 9f091f5
Files changed (2) hide show
  1. requirements.txt +5 -0
  2. tokenizer_base.py +132 -0
requirements.txt ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ torch==1.11.0
2
+ torchvision==0.12.0
3
+ onnx==1.14.0
4
+ onnxruntime==1.15.1
5
+ Pillow==10.0.0
tokenizer_base.py ADDED
@@ -0,0 +1,132 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import re
2
+ from abc import ABC, abstractmethod
3
+ from itertools import groupby
4
+ from typing import List, Optional, Tuple
5
+
6
+ import torch
7
+ from torch import Tensor
8
+ from torch.nn.utils.rnn import pad_sequence
9
+
10
+
11
+ class CharsetAdapter:
12
+ """Transforms labels according to the target charset."""
13
+
14
+ def __init__(self, target_charset) -> None:
15
+ super().__init__()
16
+ self.charset = target_charset ###
17
+ self.lowercase_only = target_charset == target_charset.lower()
18
+ self.uppercase_only = target_charset == target_charset.upper()
19
+ # self.unsupported = f'[^{re.escape(target_charset)}]'
20
+
21
+ def __call__(self, label):
22
+ if self.lowercase_only:
23
+ label = label.lower()
24
+ elif self.uppercase_only:
25
+ label = label.upper()
26
+ return label
27
+
28
+
29
+ class BaseTokenizer(ABC):
30
+
31
+ def __init__(self, charset: str, specials_first: tuple = (), specials_last: tuple = ()) -> None:
32
+ self._itos = specials_first + tuple(charset+'[UNK]') + specials_last
33
+ self._stoi = {s: i for i, s in enumerate(self._itos)}
34
+
35
+ def __len__(self):
36
+ return len(self._itos)
37
+
38
+ def _tok2ids(self, tokens: str) -> List[int]:
39
+ return [self._stoi[s] for s in tokens]
40
+
41
+ def _ids2tok(self, token_ids: List[int], join: bool = True) -> str:
42
+ tokens = [self._itos[i] for i in token_ids]
43
+ return ''.join(tokens) if join else tokens
44
+
45
+ @abstractmethod
46
+ def encode(self, labels: List[str], device: Optional[torch.device] = None) -> Tensor:
47
+ """Encode a batch of labels to a representation suitable for the model.
48
+
49
+ Args:
50
+ labels: List of labels. Each can be of arbitrary length.
51
+ device: Create tensor on this device.
52
+
53
+ Returns:
54
+ Batched tensor representation padded to the max label length. Shape: N, L
55
+ """
56
+ raise NotImplementedError
57
+
58
+ @abstractmethod
59
+ def _filter(self, probs: Tensor, ids: Tensor) -> Tuple[Tensor, List[int]]:
60
+ """Internal method which performs the necessary filtering prior to decoding."""
61
+ raise NotImplementedError
62
+
63
+ def decode(self, token_dists: Tensor, raw: bool = False) -> Tuple[List[str], List[Tensor]]:
64
+ """Decode a batch of token distributions.
65
+
66
+ Args:
67
+ token_dists: softmax probabilities over the token distribution. Shape: N, L, C
68
+ raw: return unprocessed labels (will return list of list of strings)
69
+
70
+ Returns:
71
+ list of string labels (arbitrary length) and
72
+ their corresponding sequence probabilities as a list of Tensors
73
+ """
74
+ batch_tokens = []
75
+ batch_probs = []
76
+ for dist in token_dists:
77
+ probs, ids = dist.max(-1) # greedy selection
78
+ if not raw:
79
+ probs, ids = self._filter(probs, ids)
80
+ tokens = self._ids2tok(ids, not raw)
81
+ batch_tokens.append(tokens)
82
+ batch_probs.append(probs)
83
+ return batch_tokens, batch_probs
84
+
85
+
86
+ class Tokenizer(BaseTokenizer):
87
+ BOS = '[B]'
88
+ EOS = '[E]'
89
+ PAD = '[P]'
90
+
91
+ def __init__(self, charset: str) -> None:
92
+ specials_first = (self.EOS,)
93
+ specials_last = (self.BOS, self.PAD)
94
+ super().__init__(charset, specials_first, specials_last)
95
+ self.eos_id, self.bos_id, self.pad_id = [self._stoi[s] for s in specials_first + specials_last]
96
+
97
+ def encode(self, labels: List[str], device: Optional[torch.device] = None) -> Tensor:
98
+ batch = [torch.as_tensor([self.bos_id] + self._tok2ids(y) + [self.eos_id], dtype=torch.long, device=device)
99
+ for y in labels]
100
+ return pad_sequence(batch, batch_first=True, padding_value=self.pad_id)
101
+
102
+ def _filter(self, probs: Tensor, ids: Tensor) -> Tuple[Tensor, List[int]]:
103
+ ids = ids.tolist()
104
+ try:
105
+ eos_idx = ids.index(self.eos_id)
106
+ except ValueError:
107
+ eos_idx = len(ids) # Nothing to truncate.
108
+ # Truncate after EOS
109
+ ids = ids[:eos_idx]
110
+ probs = probs[:eos_idx + 1] # but include prob. for EOS (if it exists)
111
+ return probs, ids
112
+
113
+
114
+ class CTCTokenizer(BaseTokenizer):
115
+ BLANK = '[B]'
116
+
117
+ def __init__(self, charset: str) -> None:
118
+ # BLANK uses index == 0 by default
119
+ super().__init__(charset, specials_first=(self.BLANK,))
120
+ self.blank_id = self._stoi[self.BLANK]
121
+
122
+ def encode(self, labels: List[str], device: Optional[torch.device] = None) -> Tensor:
123
+ # We use a padded representation since we don't want to use CUDNN's CTC implementation
124
+ batch = [torch.as_tensor(self._tok2ids(y), dtype=torch.long, device=device) for y in labels]
125
+ return pad_sequence(batch, batch_first=True, padding_value=self.blank_id)
126
+
127
+ def _filter(self, probs: Tensor, ids: Tensor) -> Tuple[Tensor, List[int]]:
128
+ # Best path decoding:
129
+ ids = list(zip(*groupby(ids.tolist())))[0] # Remove duplicate tokens
130
+ ids = [x for x in ids if x != self.blank_id] # Remove BLANKs
131
+ # `probs` is just pass-through since all positions are considered part of the path
132
+ return probs, ids