summaryrefslogtreecommitdiff
path: root/text_recognizer/data/transforms/word_piece.py
blob: a15615d4dee1b70c2c8ae0c297b0dd1c46e76ed4 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
"""Target transform for word pieces."""
from pathlib import Path
from typing import Optional, Union, Set

import torch
from torch import Tensor

from text_recognizer.data.mappings.word_piece_mapping import WordPieceMapping


class WordPiece:
    """Converts EMNIST indices to Word Piece indices."""

    def __init__(
        self,
        num_features: int = 1000,
        tokens: str = "iamdb_1kwp_tokens_1000.txt",
        lexicon: str = "iamdb_1kwp_lex_1000.txt",
        use_words: bool = False,
        prepend_wordsep: bool = False,
        special_tokens: Set[str] = {"<s>", "<e>", "<p>"},
        extra_symbols: Optional[Set[str]] = {
            "\n",
        },
        max_len: int = 451,
    ) -> None:
        self.mapping = WordPieceMapping(
            num_features=num_features,
            tokens=tokens,
            lexicon=lexicon,
            use_words=use_words,
            prepend_wordsep=prepend_wordsep,
            special_tokens=special_tokens,
            extra_symbols=extra_symbols,
        )
        self.max_len = max_len

    def __call__(self, x: Tensor) -> Tensor:
        """Converts Emnist target tensor to Word piece target tensor."""
        y = self.mapping.emnist_to_wordpiece_indices(x)
        if len(y) < self.max_len:
            pad_len = self.max_len - len(y)
            y = torch.cat(
                (y, torch.LongTensor([self.mapping.get_index("<p>")] * pad_len))
            )
        else:
            y = y[: self.max_len]
        return y