text_encoder.py

# coding=utf-8
# Copyright 2023 The TensorFlow Datasets Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# coding=utf-8
"""TextEncoders convert between text and integers."""
from __future__ import unicode_literals

import abc
import hashlib
import json
import re

from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf


def _re_compile(pattern):
  return re.compile(pattern, flags=re.UNICODE)  # pytype: disable=wrong-keyword-args


NUM_BYTES = 2**8
ALPHANUM_REGEX = _re_compile(r"\W+")
ALL_REGEX = _re_compile(r"(\W+)")



class TextEncoderConfig:
  """Configuration for `tfds.features.Text`."""

  def __init__(
      self, encoder=None, encoder_cls=None, vocab_size=None, name=None
  ):
    if encoder:
      if encoder_cls or vocab_size:
        raise ValueError(
            "If encoder is provided, encoder_cls and vocab_size must be None"
        )
      encoder_cls = type(encoder)
      vocab_size = encoder.vocab_size
    else:
      if encoder_cls is ByteTextEncoder:
        encoder = encoder_cls()

    self.encoder = encoder
    self.encoder_cls = encoder_cls
    self.vocab_size = vocab_size
    self.name = name


class TextEncoder(metaclass=abc.ABCMeta):
  """Abstract base class for converting between text and integers.

  **A note on padding**:

    Because text data is typically variable length and nearly always requires
    padding during training, ID 0 is always reserved for padding. To accommodate
    this, all `TextEncoder`s behave in certain ways:

    * `encode`: never returns id 0 (all ids are 1+)
    * `decode`: drops 0 in the input ids
    * `vocab_size`: includes ID 0

    New subclasses should be careful to match this behavior.
  """

  @abc.abstractmethod
  def encode(self, s):
    """Encodes text into a list of integers."""
    raise NotImplementedError

  @abc.abstractmethod
  def decode(self, ids):
    """Decodes a list of integers into text."""
    raise NotImplementedError

  @abc.abstractproperty
  def vocab_size(self):
    """Size of the vocabulary. Decode produces ints [1, vocab_size)."""
    raise NotImplementedError

  @abc.abstractmethod
  def save_to_file(self, filename_prefix):
    """Store to file. Inverse of load_from_file."""
    raise NotImplementedError

  @classmethod
  @abc.abstractmethod
  def load_from_file(cls, filename_prefix):  # pylint: disable=no-self-argument
    """Load from file. Inverse of save_to_file."""
    raise NotImplementedError

  @classmethod
  def _write_lines_to_file(cls, filename, lines, metadata_dict=None):
    """Writes lines to file prepended by header and metadata."""
    write_lines_to_file(cls.__name__, filename, lines, metadata_dict)

  @classmethod
  def _read_lines_from_file(cls, filename):
    return read_lines_from_file(cls.__name__, filename)

  def __repr__(self):
    return "<%s vocab_size=%d>" % (type(self).__name__, self.vocab_size)


class ByteTextEncoder(TextEncoder):
  """Byte-encodes text."""

  def __init__(self, additional_tokens=None):
    """Constructs ByteTextEncoder.

    Args:
      additional_tokens: `list<str>`, list of additional tokens. These will be
        assigned vocab ids `[1, 1+len(additional_tokens)]`. Useful for things
        like "end-of-string" tokens (e.g. "<EOS>").
    """
    self._additional_tokens, self._additional_tokens_re = (
        _prepare_reserved_tokens(additional_tokens)
    )
    # Note that internally everything is 0-indexed. Padding is dealt with at the
    # end of encode and the beginning of decode.
    self._additional_token_to_id = dict(
        zip(self._additional_tokens, range(len(self._additional_tokens)))
    )

  def encode(self, s):
    if not self.additional_tokens:
      return pad_incr(list(bytearray(tf.compat.as_bytes(s))))

    # Handle additional tokens
    s = tf.compat.as_text(s)
    ids = []
    for substr in self._additional_tokens_re.split(s):
      if not substr:
        continue
      tok_id = self._additional_token_to_id.get(substr)
      if tok_id is None:
        offset = len(self.additional_tokens)
        tok_ids = [
            i + offset for i in list(bytearray(tf.compat.as_bytes(substr)))
        ]
      else:
        tok_ids = [tok_id]
      ids.extend(tok_ids)

    return pad_incr(ids)

  def decode(self, ids):
    ids = pad_decr(ids)
    if not self.additional_tokens:
      return tf.compat.as_text(bytes(bytearray(ids)))

    # Handle additional tokens
    # First pass picks out the additional tokens
    tmp_decoded = []
    for byte_id in ids:
      is_additional_token = byte_id < len(self.additional_tokens)
      if is_additional_token:
        tmp_decoded.append(self.additional_tokens[byte_id])
      else:
        # Leave these as ints so that we can contiguously decode bytes
        # afterwards
        tmp_decoded.append(byte_id - len(self.additional_tokens))

    # Second pass to decode contiguous bytes
    strs = []
    i = 0
    while i < len(tmp_decoded):
      el = tmp_decoded[i]
      if isinstance(el, str):
        strs.append(el)
        i += 1
      else:
        # Decode contiguous bytes
        byte_ids = []
        while i < len(tmp_decoded):
          b = tmp_decoded[i]
          if isinstance(b, int):
            byte_ids.append(b)
            i += 1
          else:
            break
        strs.append(bytes(bytearray(byte_ids)).decode("utf-8", "replace"))
    return "".join(strs)

  @property
  def vocab_size(self):
    # Plus 1 for pad
    return len(self.additional_tokens) + NUM_BYTES + 1

  @property
  def additional_tokens(self):
    return self._additional_tokens

  @classmethod
  def _filename(cls, filename_prefix):
    return filename_prefix + ".bytes"

  def save_to_file(self, filename_prefix):
    self._write_lines_to_file(
        self._filename(filename_prefix), self.additional_tokens
    )

  @classmethod
  def load_from_file(cls, filename_prefix):
    lines, _ = cls._read_lines_from_file(cls._filename(filename_prefix))
    return cls(additional_tokens=lines)


class TokenTextEncoder(TextEncoder):
  r"""TextEncoder backed by a list of tokens.

  Tokenization splits on (and drops) non-alphanumeric characters with
  regex "\W+".
  """

  def __init__(
      self,
      vocab_list,
      oov_buckets=1,
      oov_token="UNK",
      lowercase=False,
      tokenizer=None,
      strip_vocab=True,
      decode_token_separator=" ",
  ):
    """Constructs a TokenTextEncoder.

    To load from a file saved with `TokenTextEncoder.save_to_file`, use
    `TokenTextEncoder.load_from_file`.

    Args:
      vocab_list: `list<str>`, list of tokens.
      oov_buckets: `int`, the number of `int`s to reserve for OOV hash buckets.
        Tokens that are OOV will be hash-modded into a OOV bucket in `encode`.
      oov_token: `str`, the string to use for OOV ids in `decode`.
      lowercase: `bool`, whether to make all text and tokens lowercase.
      tokenizer: `Tokenizer`, responsible for converting incoming text into a
        list of tokens.
      strip_vocab: `bool`, whether to strip whitespace from the beginning and
        end of elements of `vocab_list`.
      decode_token_separator: `str`, the string used to separate tokens when
        decoding.
    """
    self._vocab_list = [tf.compat.as_text(el) for el in vocab_list]
    if strip_vocab:
      self._vocab_list = [el.strip() for el in self._vocab_list]
    self._lowercase = lowercase
    if self._lowercase:
      self._vocab_list = [t.lower() for t in self._vocab_list]
    # Note that internally everything is 0-indexed. Padding is dealt with at the
    # end of encode and the beginning of decode.
    self._token_to_id = dict(
        zip(self._vocab_list, range(len(self._vocab_list)))
    )
    self._oov_buckets = oov_buckets
    self._oov_token = tf.compat.as_text(oov_token)

    # Reserved tokens are all tokens that are mixed alphanum and non-alphanum.
    reserved_tokens = [t for t in self._vocab_list if is_mixed_alphanum(t)]
    self._tokenizer = tokenizer or Tokenizer(reserved_tokens=reserved_tokens)
    self._user_defined_tokenizer = tokenizer

    self._decode_token_separator = decode_token_separator

  def encode(self, s):
    s = tf.compat.as_text(s)
    if self.lowercase:
      s = s.lower()
    ids = []
    for token in self._tokenizer.tokenize(s):
      int_id = self._token_to_id.get(token, -1)
      if int_id < 0:
        int_id = self._oov_bucket(token)
        if int_id is None:
          raise ValueError("Out of vocabulary token %s" % token)
      ids.append(int_id)

    # Increment for pad id 0
    return pad_incr(ids)

  def decode(self, ids):
    ids = pad_decr(ids)

    tokens = []
    for int_id in ids:
      if int_id < len(self._vocab_list):
        tokens.append(self._vocab_list[int_id])
      else:
        tokens.append(self._oov_token)
    return self._decode_token_separator.join(tokens)

  @property
  def vocab_size(self):
    # Plus 1 for pad
    return len(self._vocab_list) + self._oov_buckets + 1

  @property
  def tokens(self):
    return list(self._vocab_list)

  @property
  def oov_token(self):
    return self._oov_token

  @property
  def lowercase(self):
    return self._lowercase

  @property
  def tokenizer(self):
    return self._tokenizer

  def _oov_bucket(self, token):
    if self._oov_buckets <= 0:
      return None
    if self._oov_buckets == 1:
      return len(self._vocab_list)
    hash_val = int(hashlib.md5(tf.compat.as_bytes(token)).hexdigest(), 16)
    return len(self._vocab_list) + hash_val % self._oov_buckets

  @classmethod
  def _filename(cls, filename_prefix):
    return filename_prefix + ".tokens"

  def save_to_file(self, filename_prefix):
    filename = self._filename(filename_prefix)
    kwargs = {
        "oov_buckets": self._oov_buckets,
        "lowercase": self._lowercase,
        "oov_token": self._oov_token,
    }
    if self._user_defined_tokenizer is not None:
      self._tokenizer.save_to_file(filename)
      kwargs["has_tokenizer"] = True
    self._write_lines_to_file(filename, self._vocab_list, kwargs)

  @classmethod
  def load_from_file(cls, filename_prefix):
    filename = cls._filename(filename_prefix)
    vocab_lines, kwargs = cls._read_lines_from_file(filename)
    has_tokenizer = kwargs.pop("has_tokenizer", False)
    if has_tokenizer:
      kwargs["tokenizer"] = Tokenizer.load_from_file(filename)
    return cls(vocab_list=vocab_lines, **kwargs)


class Tokenizer(object):
  """Splits a string into tokens, and joins them back."""

  def __init__(self, alphanum_only=True, reserved_tokens=None):
    """Constructs a Tokenizer.

    Note that the Tokenizer is invertible if `alphanum_only=False`.
    i.e. `s == t.join(t.tokenize(s))`.

    Args:
      alphanum_only: `bool`, if `True`, only parse out alphanumeric tokens
        (non-alphanumeric characters are dropped); otherwise, keep all
        characters (individual tokens will still be either all alphanumeric or
        all non-alphanumeric).
      reserved_tokens: `list<str>`, a list of strings that, if any are in `s`,
        will be preserved as whole tokens, even if they contain mixed
        alphanumeric/non-alphanumeric characters.
    """
    self._alphanum_only = alphanum_only
    reserved_tokens, self._reserved_tokens_re = _prepare_reserved_tokens(
        reserved_tokens
    )
    self._reserved_tokens = set(reserved_tokens)

  @property
  def alphanum_only(self):
    return self._alphanum_only

  @property
  def reserved_tokens(self):
    return self._reserved_tokens

  def tokenize(self, s):
    """Splits a string into tokens."""
    s = tf.compat.as_text(s)

    if self.reserved_tokens:
      # First split out the reserved tokens
      substrs = self._reserved_tokens_re.split(s)
    else:
      substrs = [s]

    toks = []
    for substr in substrs:
      if substr in self.reserved_tokens:
        toks.append(substr)
      elif self._alphanum_only:
        toks.extend(ALPHANUM_REGEX.split(substr))
      else:
        toks.extend(ALL_REGEX.split(substr))

    # Filter out empty strings
    toks = [t for t in toks if t]
    return toks

  def join(self, tokens):
    """Joins tokens into a string."""
    if self._alphanum_only:
      return " ".join(tokens)
    else:
      # Fully invertible
      return "".join(tokens)

  @classmethod
  def _filename(cls, filename_prefix):
    return filename_prefix + ".tokenizer"

  def save_to_file(self, filename_prefix):
    filename = self._filename(filename_prefix)
    kwargs = {
        "reserved_tokens": list(self._reserved_tokens),
        "alphanum_only": self._alphanum_only,
    }
    write_lines_to_file(type(self).__name__, filename, [], kwargs)

  @classmethod
  def load_from_file(cls, filename_prefix):
    filename = cls._filename(filename_prefix)
    _, kwargs = read_lines_from_file(cls.__name__, filename)
    return cls(**kwargs)


def pad_decr(ids):
  """Strip ID 0 and decrement ids by 1."""
  if len(ids) < 1:
    return list(ids)
  if not any(ids):
    return []  # all padding.
  idx = -1
  while not ids[idx]:
    idx -= 1
  if idx == -1:
    ids = ids  # pylint: disable=self-assigning-variable
  else:
    ids = ids[: idx + 1]
  return [i - 1 for i in ids]


def pad_incr(ids):
  """Add 1 to ids to account for pad."""
  return [i + 1 for i in ids]


def _prepare_reserved_tokens(reserved_tokens):
  """Prepare reserved tokens and a regex for splitting them out of strings."""
  reserved_tokens = [tf.compat.as_text(tok) for tok in reserved_tokens or []]
  dups = _find_duplicates(reserved_tokens)
  if dups:
    raise ValueError("Duplicates found in tokens: %s" % dups)
  reserved_tokens_re = _make_reserved_tokens_re(reserved_tokens)
  return reserved_tokens, reserved_tokens_re


def _re_escape(s):
  """Escape regex control characters."""
  escaped = re.sub(r"[(){}\[\].*?|^$\\+-]", r"\\\g<0>", s)
  return escaped


def _make_reserved_tokens_re(reserved_tokens):
  """Constructs compiled regex to parse out reserved tokens."""
  if not reserved_tokens:
    return None
  escaped_tokens = [_re_escape(rt) for rt in reserved_tokens]
  pattern = "(%s)" % "|".join(escaped_tokens)
  reserved_tokens_re = _re_compile(pattern)
  return reserved_tokens_re


def _find_duplicates(els):
  seen = set()
  dups = []
  for x in els:
    if x in seen:
      dups.append(x)
    else:
      seen.add(x)
  return dups


def is_mixed_alphanum(token):
  return len([s for s in ALL_REGEX.split(token) if s]) > 1


_HEADER_PREFIX = "### "
_METADATA_PREFIX = "### Metadata: "


def write_lines_to_file(cls_name, filename, lines, metadata_dict):
  """Writes lines to file prepended by header and metadata."""
  metadata_dict = metadata_dict or {}
  header_line = "%s%s" % (_HEADER_PREFIX, cls_name)
  metadata_line = "%s%s" % (
      _METADATA_PREFIX,
      json.dumps(metadata_dict, sort_keys=True),
  )
  with tf.io.gfile.GFile(filename, "wb") as f:
    for line in [header_line, metadata_line]:
      f.write(tf.compat.as_bytes(line))
      f.write(tf.compat.as_bytes("\n"))
    if lines:
      f.write(tf.compat.as_bytes("\n".join(lines)))
      f.write(tf.compat.as_bytes("\n"))


def read_lines_from_file(cls_name, filename):
  """Read lines from file, parsing out header and metadata."""
  with tf.io.gfile.GFile(filename, "rb") as f:
    lines = [tf.compat.as_text(line)[:-1] for line in f]
  header_line = "%s%s" % (_HEADER_PREFIX, cls_name)
  if lines[0] != header_line:
    raise ValueError(
        "File {fname} does not seem to have been created from "
        "{name}.save_to_file.".format(fname=filename, name=cls_name)
    )
  metadata_dict = json.loads(lines[1][len(_METADATA_PREFIX) :])
  return lines[2:], metadata_dict