# Path: fedjax/datasets/stackoverflow.py
# SPLITS = ('train', 'held_out', 'test')
#   PAD = 0
#   BOS = 1
#   EOS = 2
# def cite():
# def load_split(split: str,
#                mode: str = 'sqlite',
#                cache_dir: Optional[str] = None) -> federated_data.FederatedData:
# def load_data(
#     mode: str = 'sqlite',
#     cache_dir: Optional[str] = None
# ) -> Tuple[federated_data.FederatedData, federated_data.FederatedData,
# def preprocess_client(
#     client_id: federated_data.ClientId,
#     examples: client_datasets.Examples) -> client_datasets.Examples:
# def default_vocab(default_vocab_size) -> List[str]:
#   def __init__(self,
#                vocab: Optional[List[str]] = None,
#                default_vocab_size: Optional[int] = 10000,
#                num_oov_buckets: int = 1):
#   def create_token_to_ids_fn(self, max_length: int):
#     def token_to_ids(tokens):
#   def as_preprocess_batch(
#       self, max_length: int
#   ) -> Callable[[client_datasets.Examples], client_datasets.Examples]:
#     def preprocess_batch(
#         examples: client_datasets.Examples) -> client_datasets.Examples:
# class StackoverflowTokenizer:


from absl.testing import absltest
from fedjax.datasets import stackoverflow
import numpy as np
import numpy.testing as npt
      npt.assert_equal(
          tokenizer.as_preprocess_batch(max_length=2)(examples), {
              'x': [[1, 3], [1, 5]],
              'y': [[3, 4], [5, 3]]
          })
      # Padding and truncation.
      npt.assert_equal(
          tokenizer.as_preprocess_batch(max_length=5)(examples), {
              'x': [[1, 3, 4, 5, 0], [1, 5, 3, 4, 5]],
              'y': [[3, 4, 5, 2, 0], [5, 3, 4, 5, 6]]
          })
      # Padding only.
      npt.assert_equal(
          tokenizer.as_preprocess_batch(max_length=7)(examples), {
              'x': [[1, 3, 4, 5, 0, 0, 0], [1, 5, 3, 4, 5, 6, 0]],
              'y': [[3, 4, 5, 2, 0, 0, 0], [5, 3, 4, 5, 6, 2, 0]]
          })

    with self.subTest('domain_id'):
      npt.assert_equal(
          tokenizer.as_preprocess_batch(max_length=2)({
              **examples, 'domain_id': np.array([1, 0])
          }), {
              'x': [[1, 3], [1, 5]],
              'y': [[3, 4], [5, 3]],
              'domain_id': [1, 0]
          })

    with self.subTest('num_oov_buckets'):
      npt.assert_equal(

# Path: fedjax/datasets/stackoverflow.py
# SPLITS = ('train', 'held_out', 'test')
#   PAD = 0
#   BOS = 1
#   EOS = 2
# def cite():
# def load_split(split: str,
#                mode: str = 'sqlite',
#                cache_dir: Optional[str] = None) -> federated_data.FederatedData:
# def load_data(
#     mode: str = 'sqlite',
#     cache_dir: Optional[str] = None
# ) -> Tuple[federated_data.FederatedData, federated_data.FederatedData,
# def preprocess_client(
#     client_id: federated_data.ClientId,
#     examples: client_datasets.Examples) -> client_datasets.Examples:
# def default_vocab(default_vocab_size) -> List[str]:
#   def __init__(self,
#                vocab: Optional[List[str]] = None,
#                default_vocab_size: Optional[int] = 10000,
#                num_oov_buckets: int = 1):
#   def create_token_to_ids_fn(self, max_length: int):
#     def token_to_ids(tokens):
#   def as_preprocess_batch(
#       self, max_length: int
#   ) -> Callable[[client_datasets.Examples], client_datasets.Examples]:
#     def preprocess_batch(
#         examples: client_datasets.Examples) -> client_datasets.Examples:
# class StackoverflowTokenizer:


# Path: fedjax/datasets/stackoverflow_test.py
from absl.testing import absltest
from fedjax.datasets import stackoverflow
import numpy as np
import numpy.testing as npt

      npt.assert_equal(
          tokenizer.as_preprocess_batch(max_length=2)(examples), {
              'x': [[1, 3], [1, 5]],
              'y': [[3, 4], [5, 3]]
          })
      # Padding and truncation.
      npt.assert_equal(
          tokenizer.as_preprocess_batch(max_length=5)(examples), {
              'x': [[1, 3, 4, 5, 0], [1, 5, 3, 4, 5]],
              'y': [[3, 4, 5, 2, 0], [5, 3, 4, 5, 6]]
          })
      # Padding only.
      npt.assert_equal(
          tokenizer.as_preprocess_batch(max_length=7)(examples), {
              'x': [[1, 3, 4, 5, 0, 0, 0], [1, 5, 3, 4, 5, 6, 0]],
              'y': [[3, 4, 5, 2, 0, 0, 0], [5, 3, 4, 5, 6, 2, 0]]
          })

    with self.subTest('domain_id'):
      npt.assert_equal(
          tokenizer.as_preprocess_batch(max_length=2)({
              **examples, 'domain_id': np.array([1, 0])
          }), {
              'x': [[1, 3], [1, 5]],
              'y': [[3, 4], [5, 3]],
              'domain_id': [1, 0]
          })

    with self.subTest('num_oov_buckets'):
      npt.assert_equal(

          stackoverflow.StackoverflowTokenizer(