| import os |
| import json |
| import torch |
| import numpy as np |
| from transformers import BatchEncoding |
| from typing import Dict, List, Tuple, Union |
| from sentencepiece import SentencePieceProcessor |
|
|
| _PATH = os.path.dirname(os.path.realpath(__file__)) |
|
|
|
|
| class IndicTransTokenizer: |
| def __init__( |
| self, |
| src_vocab_fp=None, |
| tgt_vocab_fp=None, |
| src_spm_fp=None, |
| tgt_spm_fp=None, |
| unk_token="<unk>", |
| bos_token="<s>", |
| eos_token="</s>", |
| pad_token="<pad>", |
| direction="indic-en", |
| model_max_length=256, |
| ): |
| self.model_max_length = model_max_length |
|
|
| self.supported_langs = [ |
| "asm_Beng", |
| "ben_Beng", |
| "brx_Deva", |
| "doi_Deva", |
| "eng_Latn", |
| "gom_Deva", |
| "guj_Gujr", |
| "hin_Deva", |
| "kan_Knda", |
| "kas_Arab", |
| "kas_Deva", |
| "mai_Deva", |
| "mal_Mlym", |
| "mar_Deva", |
| "mni_Beng", |
| "mni_Mtei", |
| "npi_Deva", |
| "ory_Orya", |
| "pan_Guru", |
| "san_Deva", |
| "sat_Olck", |
| "snd_Arab", |
| "snd_Deva", |
| "tam_Taml", |
| "tel_Telu", |
| "urd_Arab", |
| ] |
|
|
| self.src_vocab_fp = ( |
| src_vocab_fp |
| if (src_vocab_fp is not None) |
| else os.path.join(_PATH, direction, "dict.SRC.json") |
| ) |
| self.tgt_vocab_fp = ( |
| tgt_vocab_fp |
| if (tgt_vocab_fp is not None) |
| else os.path.join(_PATH, direction, "dict.TGT.json") |
| ) |
| self.src_spm_fp = ( |
| src_spm_fp |
| if (src_spm_fp is not None) |
| else os.path.join(_PATH, direction, "model.SRC") |
| ) |
| self.tgt_spm_fp = ( |
| tgt_spm_fp |
| if (tgt_spm_fp is not None) |
| else os.path.join(_PATH, direction, "model.TGT") |
| ) |
|
|
| self.unk_token = unk_token |
| self.pad_token = pad_token |
| self.eos_token = eos_token |
| self.bos_token = bos_token |
|
|
| self.encoder = self._load_json(self.src_vocab_fp) |
| if self.unk_token not in self.encoder: |
| raise KeyError("<unk> token must be in vocab") |
| assert self.pad_token in self.encoder |
| self.encoder_rev = {v: k for k, v in self.encoder.items()} |
|
|
| self.decoder = self._load_json(self.tgt_vocab_fp) |
| if self.unk_token not in self.encoder: |
| raise KeyError("<unk> token must be in vocab") |
| assert self.pad_token in self.encoder |
| self.decoder_rev = {v: k for k, v in self.decoder.items()} |
|
|
| |
| self.src_spm = self._load_spm(self.src_spm_fp) |
| self.tgt_spm = self._load_spm(self.tgt_spm_fp) |
|
|
| def is_special_token(self, x: str): |
| return (x == self.pad_token) or (x == self.bos_token) or (x == self.eos_token) |
|
|
| def get_vocab_size(self, src: bool) -> int: |
| """Returns the size of the vocabulary""" |
| return len(self.encoder) if src else len(self.decoder) |
|
|
| def _load_spm(self, path: str) -> SentencePieceProcessor: |
| return SentencePieceProcessor(model_file=path) |
|
|
| def _save_json(self, data, path: str) -> None: |
| with open(path, "w", encoding="utf-8") as f: |
| json.dump(data, f, indent=2) |
|
|
| def _load_json(self, path: str) -> Union[Dict, List]: |
| with open(path, "r", encoding="utf-8") as f: |
| return json.load(f) |
|
|
| def _convert_token_to_id(self, token: str, src: bool) -> int: |
| """Converts an token (str) into an index (integer) using the source/target vocabulary map.""" |
| return ( |
| self.encoder.get(token, self.encoder[self.unk_token]) |
| if src |
| else self.decoder.get(token, self.encoder[self.unk_token]) |
| ) |
|
|
| def _convert_id_to_token(self, index: int, src: bool) -> str: |
| """Converts an index (integer) into a token (str) using the source/target vocabulary map.""" |
| return ( |
| self.encoder_rev.get(index, self.unk_token) |
| if src |
| else self.decoder_rev.get(index, self.unk_token) |
| ) |
|
|
| def _convert_tokens_to_string(self, tokens: List[str], src: bool) -> str: |
| """Uses sentencepiece model for detokenization""" |
| if src: |
| if tokens[0] in self.supported_langs and tokens[1] in self.supported_langs: |
| tokens = tokens[2:] |
| return " ".join(tokens) |
| else: |
| return " ".join(tokens) |
|
|
| def _remove_translation_tags(self, text: str) -> Tuple[List, str]: |
| """Removes the translation tags before text normalization and tokenization.""" |
| tokens = text.split(" ") |
| return tokens[:2], " ".join(tokens[2:]) |
|
|
| def _tokenize_src_line(self, line: str) -> List[str]: |
| """Tokenizes a source line.""" |
| tags, text = self._remove_translation_tags(line) |
| tokens = self.src_spm.encode(text, out_type=str) |
| return tags + tokens |
|
|
| def _tokenize_tgt_line(self, line: str) -> List[str]: |
| """Tokenizes a target line.""" |
| return self.tgt_spm.encode(line, out_type=str) |
|
|
| def tokenize(self, text: str, src: bool) -> List[str]: |
| """Tokenizes a string into tokens using the source/target vocabulary.""" |
| return self._tokenize_src_line(text) if src else self._tokenize_tgt_line(text) |
|
|
| def batch_tokenize(self, batch: List[str], src: bool) -> List[List[str]]: |
| """Tokenizes a list of strings into tokens using the source/target vocabulary.""" |
| return [self.tokenize(line, src) for line in batch] |
|
|
| def _create_attention_mask(self, ids: List[int], max_seq_len: int) -> List[int]: |
| """Creates a attention mask for the input sequence.""" |
| return ([0] * (max_seq_len - len(ids))) + ([1] * (len(ids) + 1)) |
|
|
| def _pad_batch(self, tokens: List[str], max_seq_len: int) -> List[str]: |
| """Pads a batch of tokens and adds BOS/EOS tokens.""" |
| return ( |
| ([self.pad_token] * (max_seq_len - len(tokens))) + tokens + [self.eos_token] |
| ) |
|
|
| def _decode_line(self, ids: List[int], src: bool) -> List[str]: |
| return [self._convert_id_to_token(_id, src) for _id in ids] |
|
|
| def _encode_line(self, tokens: List[str], src: bool) -> List[int]: |
| return [self._convert_token_to_id(token, src) for token in tokens] |
|
|
| def _strip_special_tokens(self, tokens: List[str]) -> List[str]: |
| return [token for token in tokens if not self.is_special_token(token)] |
|
|
| def _single_input_preprocessing( |
| self, tokens: List[str], src: bool, max_seq_len: int |
| ) -> Tuple[List[int], List[int], int]: |
| """Tokenizes a string into tokens and also converts them into integers using source/target vocabulary map.""" |
| attention_mask = self._create_attention_mask(tokens, max_seq_len) |
| padded_tokens = self._pad_batch(tokens, max_seq_len) |
| input_ids = self._encode_line(padded_tokens, src) |
| return input_ids, attention_mask |
|
|
| def _single_output_postprocessing(self, ids: List[int], src: bool) -> str: |
| """Detokenizes a list of integer ids into a string using the source/target vocabulary.""" |
| tokens = self._decode_line(ids, src) |
| tokens = self._strip_special_tokens(tokens) |
| return self._convert_tokens_to_string(tokens, src) |
|
|
| def __call__( |
| self, |
| batch: Union[list, str], |
| src: bool, |
| truncation: bool = False, |
| padding: str = "longest", |
| max_length: int = None, |
| return_tensors: str = "pt", |
| return_attention_mask: bool = True, |
| return_length: bool = False, |
| ) -> BatchEncoding: |
| """Tokenizes a string into tokens and also converts them into integers using source/target vocabulary map.""" |
| assert padding in [ |
| "longest", |
| "max_length", |
| ], "padding should be either 'longest' or 'max_length'" |
|
|
| if not isinstance(batch, list): |
| raise TypeError( |
| f"batch must be a list, but current batch is of type {type(batch)}" |
| ) |
|
|
| |
| batch = self.batch_tokenize(batch, src) |
|
|
| |
| if truncation and max_length is not None: |
| batch = [ids[:max_length] for ids in batch] |
|
|
| lengths = [len(ids) for ids in batch] |
|
|
| max_seq_len = max(lengths) if padding == "longest" else max_length |
|
|
| input_ids, attention_mask = zip( |
| *[ |
| self._single_input_preprocessing( |
| tokens=tokens, src=src, max_seq_len=max_seq_len |
| ) |
| for tokens in batch |
| ] |
| ) |
|
|
| _data = {"input_ids": input_ids} |
|
|
| if return_attention_mask: |
| _data["attention_mask"] = attention_mask |
| |
| if return_length: |
| _data["lengths"] = lengths |
|
|
| return BatchEncoding(_data, tensor_type=return_tensors) |
|
|
| def batch_decode( |
| self, batch: Union[list, torch.Tensor], src: bool |
| ) -> List[List[str]]: |
| """Detokenizes a list of integer ids or a tensor into a list of strings using the source/target vocabulary.""" |
|
|
| if isinstance(batch, torch.Tensor): |
| batch = batch.detach().cpu().tolist() |
|
|
| return [self._single_output_postprocessing(ids=ids, src=src) for ids in batch] |
|
|
