<commit_before># Copyright 2020 The Flax 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.

"""SST-2 input pipeline."""

import collections
from typing import Dict, List, Set, Text, Tuple
import os

from absl import logging

import numpy as np
import tensorflow.compat.v2 as tf
import tensorflow_datasets as tfds


# pylint: disable=no-member


# pylint: disable=abstract-method
class OrderedCounter(collections.Counter, collections.OrderedDict):
  """A counter that remembers the order of added keys."""
  pass


def get_tokens(datasets: List[tf.data.Dataset]) -> Set[Text]:
  """Returns a set with all unique tokens in the given datasets."""
  counter = OrderedCounter()
  for ds in datasets:  # Data set already tokenized here.
    for example in tfds.as_numpy(ds):
      tokens = whitespace_tokenize(example['sentence'].strip())
      counter.update(tokens)
  sst_tokens = set(counter.keys())
  logging.info('Number of unique tokens: %d', len(sst_tokens))
  return sst_tokens


def save_vocab(vocab, path):
  """Saves a vocabulary to disk."""
  with open(path, mode='wb') as f:
    for token in vocab:
      f.write(token + b'\n')
  logging.info('Saved vocab to: %s', path)


def load_vocab(path):
  """Loads a vocabulary from disk."""
  vocab = collections.OrderedDict()
  with open(path, mode='rb') as f:
    for i, token in enumerate(f):
      vocab[token.rstrip(b'\n')] = i
  return vocab


def save_embed(embed, path):
  """Save word embeddings to disk."""
  assert path.endswith('.npz'), 'Path must end with .npz'
  np.savez(path, embed=embed)
  logging.info('Saved embed to: %s', path)


def load_embed(path):
  """Load word embeddings from disk."""
  return np.load(path)['embed']


# pylint: disable=dangerous-default-value
def get_glove_embeddings(
    datasets: List[tf.data.Dataset],
    glove_path: Text,
    glove_dim: int,
    specials: List[Text] = [b'<pad>', b'<unk>', b'<s>', b'</s>'],
    seed: int = 42,
    cache_dir: Text = None,
) -> Tuple[Dict[Text, int], np.ndarray]:
  """Filter Glove vectors to only use tokens that are in SST-2.

  The Glove file contains lines with the following format:

  ```
    cat 0.01 0.02 ... -0.02
    dog 0.02 0.03 ... -0.04
  ```
  That is, first the token, and then space-separated the value of each
  dimension.

  Args:
    datasets: A list of datasets to extract tokens from. We will only use
        embeddings for tokens in these datasets.
    glove_path: Path to GloVe word vectors.
    glove_dim: The size of the word vectors.
    specials: A list of special tokens (starting the vocabulary).
    seed: A seed to use to initialize the special token vectors.
    cache_dir: Directory to save/load the vocab and embeddings to/from.
      If it is empty, the vocab and embeddings will be saved.
      If both vocab and embeddings are found, they will be loaded.

  Returns:
    A tuple with the vocabulary (token -> ID) and the word embedding matrix.
  """
  sst_tokens = get_tokens(datasets)
  vectors = []
  vocab = collections.OrderedDict()
  np.random.seed(seed)

  # Special tokens.
  for token in specials:
    vectors.append(np.random.uniform(-0.05, 0.05, size=glove_dim))
    vocab[token] = len(vocab)
  vectors[0] = np.zeros_like(vectors[0])  # Zeros for padding.

  # Extract the embeddings for our set of tokens (discard the rest).
  with open(glove_path, mode='rb') as f:
    for line in f:
      parts = line.strip().rsplit(maxsplit=glove_dim)
      token = parts[0]
      if token in sst_tokens:
        vector = np.array(parts[1:], dtype=np.float32)
        vectors.append(vector)
        vocab[token] = len(vocab)

  embed = np.array(vectors)

  # Save vocab and embeddings to cache_dir.
  if cache_dir is not None:
    if not os.path.exists(cache_dir):
      os.makedir(cache_dir)
    save_vocab(vocab, os.path.join(cache_dir, 'vocab.txt'))
    save_embed(embed, os.path.join(cache_dir, 'embed.npz'))

  return vocab, embed


def whitespace_tokenize(s: Text):
  """Splits ab input into tokens by whitspace."""
  return s.split()


def get_batches(
    train_ds: tf.data.Dataset,
    valid_ds: tf.data.Dataset,
    test_ds: tf.data.Dataset,
    seed: int = 0,
    batch_size: int = 64):
  """Returns batched versions of each dataset."""
  autotune = tf.data.experimental.AUTOTUNE

  # For shuffling we need to know how many training examples we have.
  num_train_examples = train_ds.reduce(np.int64(0), lambda x, _: x + 1).numpy()

  # We need to pad our batches since sentences have different lengths!
  # Sentences that are shorter in a batch will get 0s added at the end, until
  # all sentences in the batch have the same length.
  # padded_shapes says what kind of shapes to expect: [] means a scalar, [-1]
  # means a vector of variable length, and [1] means a vector of size 1.
  train_batches = train_ds.shuffle(
      num_train_examples,
      seed=seed,<commit_msg>fix mkdir<commit_after># Copyright 2020 The Flax 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.

"""SST-2 input pipeline."""

import collections
from typing import Dict, List, Set, Text, Tuple
import os

from absl import logging

import numpy as np
import tensorflow.compat.v2 as tf
import tensorflow_datasets as tfds


# pylint: disable=no-member


# pylint: disable=abstract-method
class OrderedCounter(collections.Counter, collections.OrderedDict):
  """A counter that remembers the order of added keys."""
  pass


def get_tokens(datasets: List[tf.data.Dataset]) -> Set[Text]:
  """Returns a set with all unique tokens in the given datasets."""
  counter = OrderedCounter()
  for ds in datasets:  # Data set already tokenized here.
    for example in tfds.as_numpy(ds):
      tokens = whitespace_tokenize(example['sentence'].strip())
      counter.update(tokens)
  sst_tokens = set(counter.keys())
  logging.info('Number of unique tokens: %d', len(sst_tokens))
  return sst_tokens


def save_vocab(vocab, path):
  """Saves a vocabulary to disk."""
  with open(path, mode='wb') as f:
    for token in vocab:
      f.write(token + b'\n')
  logging.info('Saved vocab to: %s', path)


def load_vocab(path):
  """Loads a vocabulary from disk."""
  vocab = collections.OrderedDict()
  with open(path, mode='rb') as f:
    for i, token in enumerate(f):
      vocab[token.rstrip(b'\n')] = i
  return vocab


def save_embed(embed, path):
  """Save word embeddings to disk."""
  assert path.endswith('.npz'), 'Path must end with .npz'
  np.savez(path, embed=embed)
  logging.info('Saved embed to: %s', path)


def load_embed(path):
  """Load word embeddings from disk."""
  return np.load(path)['embed']


# pylint: disable=dangerous-default-value
def get_glove_embeddings(
    datasets: List[tf.data.Dataset],
    glove_path: Text,
    glove_dim: int,
    specials: List[Text] = [b'<pad>', b'<unk>', b'<s>', b'</s>'],
    seed: int = 42,
    cache_dir: Text = None,
) -> Tuple[Dict[Text, int], np.ndarray]:
  """Filter Glove vectors to only use tokens that are in SST-2.

  The Glove file contains lines with the following format:

  ```
    cat 0.01 0.02 ... -0.02
    dog 0.02 0.03 ... -0.04
  ```
  That is, first the token, and then space-separated the value of each
  dimension.

  Args:
    datasets: A list of datasets to extract tokens from. We will only use
        embeddings for tokens in these datasets.
    glove_path: Path to GloVe word vectors.
    glove_dim: The size of the word vectors.
    specials: A list of special tokens (starting the vocabulary).
    seed: A seed to use to initialize the special token vectors.
    cache_dir: Directory to save/load the vocab and embeddings to/from.
      If it is empty, the vocab and embeddings will be saved.
      If both vocab and embeddings are found, they will be loaded.

  Returns:
    A tuple with the vocabulary (token -> ID) and the word embedding matrix.
  """
  sst_tokens = get_tokens(datasets)
  vectors = []
  vocab = collections.OrderedDict()
  np.random.seed(seed)

  # Special tokens.
  for token in specials:
    vectors.append(np.random.uniform(-0.05, 0.05, size=glove_dim))
    vocab[token] = len(vocab)
  vectors[0] = np.zeros_like(vectors[0])  # Zeros for padding.

  # Extract the embeddings for our set of tokens (discard the rest).
  with open(glove_path, mode='rb') as f:
    for line in f:
      parts = line.strip().rsplit(maxsplit=glove_dim)
      token = parts[0]
      if token in sst_tokens:
        vector = np.array(parts[1:], dtype=np.float32)
        vectors.append(vector)
        vocab[token] = len(vocab)

  embed = np.array(vectors)

  # Save vocab and embeddings to cache_dir.
  if cache_dir is not None:
    if not os.path.exists(cache_dir):
      os.mkdir(cache_dir)
    save_vocab(vocab, os.path.join(cache_dir, 'vocab.txt'))
    save_embed(embed, os.path.join(cache_dir, 'embed.npz'))

  return vocab, embed


def whitespace_tokenize(s: Text):
  """Splits ab input into tokens by whitspace."""
  return s.split()


def get_batches(
    train_ds: tf.data.Dataset,
    valid_ds: tf.data.Dataset,
    test_ds: tf.data.Dataset,
    seed: int = 0,
    batch_size: int = 64):
  """Returns batched versions of each dataset."""
  autotune = tf.data.experimental.AUTOTUNE

  # For shuffling we need to know how many training examples we have.
  num_train_examples = train_ds.reduce(np.int64(0), lambda x, _: x + 1).numpy()

  # We need to pad our batches since sentences have different lengths!
  # Sentences that are shorter in a batch will get 0s added at the end, until
  # all sentences in the batch have the same length.
  # padded_shapes says what kind of shapes to expect: [] means a scalar, [-1]
  # means a vector of variable length, and [1] means a vector of size 1.
  train_batches = train_ds.shuffle(
      num_train_examples,
      seed=seed,
{'fraction_non_alphanumeric': 0.10916215711886545, 'fraction_numerical': 0.010076506810972196, 'mean_word_length': 3.5227848101265824, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4868613', 'n_tokens_mistral': 3548, 'n_tokens_neox': 3418, 'n_words': 1230}