, lines.
if line.startswith("") or line.startswith("
"):
return ""
# Strip tags.
line = line.strip()
if line.startswith(""):
i = line.index(">")
return line[i + 1:-6] # Strip first and last .
def _clean_sentences(sentence_pairs):
res_pairs = []
for cleaned in cleaner_en_xx.clean_en_xx_pairs(sentence_pairs):
res_pairs.append(cleaned)
return res_pairs
def _tmx_to_source_target(tmx_file, source_resfile, target_resfile,
do_cleaning=False):
source_target_pairs = cleaner_en_xx.paracrawl_v3_pairs(tmx_file)
if do_cleaning:
source_target_pairs = cleaner_en_xx.clean_en_xx_pairs(source_target_pairs)
for source, target in source_target_pairs:
source_resfile.write(source)
source_resfile.write("\n")
target_resfile.write(target)
target_resfile.write("\n")
def compile_data(tmp_dir, datasets, filename, datatypes_to_clean=None):
"""Concatenates all `datasets` and saves to `filename`."""
datatypes_to_clean = datatypes_to_clean or []
filename = os.path.join(tmp_dir, filename)
lang1_fname = filename + ".lang1"
lang2_fname = filename + ".lang2"
if tf.gfile.Exists(lang1_fname) and tf.gfile.Exists(lang2_fname):
tf.logging.info("Skipping compile data, found files:\n%s\n%s", lang1_fname,
lang2_fname)
return filename
with tf.gfile.GFile(lang1_fname, mode="w") as lang1_resfile:
with tf.gfile.GFile(lang2_fname, mode="w") as lang2_resfile:
for dataset in datasets:
url = dataset[0]
compressed_filename = os.path.basename(url)
compressed_filepath = os.path.join(tmp_dir, compressed_filename)
if url.startswith("http"):
generator_utils.maybe_download(tmp_dir, compressed_filename, url)
if compressed_filename.endswith(".zip"):
zipfile.ZipFile(os.path.join(compressed_filepath),
"r").extractall(tmp_dir)
if dataset[1][0] == "tmx":
cleaning_requested = "tmx" in datatypes_to_clean
tmx_filename = os.path.join(tmp_dir, dataset[1][1])
if tmx_filename.endswith(".gz"):
with gzip.open(tmx_filename, "rb") as tmx_file:
_tmx_to_source_target(tmx_file, lang1_resfile, lang2_resfile,
do_cleaning=cleaning_requested)
else:
with tf.gfile.Open(tmx_filename) as tmx_file:
_tmx_to_source_target(tmx_file, lang1_resfile, lang2_resfile,
do_cleaning=cleaning_requested)
elif dataset[1][0] == "tsv":
_, src_column, trg_column, glob_pattern = dataset[1]
filenames = tf.gfile.Glob(os.path.join(tmp_dir, glob_pattern))
if not filenames:
# Capture *.tgz and *.tar.gz too.
mode = "r:gz" if compressed_filepath.endswith("gz") else "r"
with tarfile.open(compressed_filepath, mode) as corpus_tar:
corpus_tar.extractall(tmp_dir)
filenames = tf.gfile.Glob(os.path.join(tmp_dir, glob_pattern))
for tsv_filename in filenames:
if tsv_filename.endswith(".gz"):
new_filename = tsv_filename.strip(".gz")
generator_utils.gunzip_file(tsv_filename, new_filename)
tsv_filename = new_filename
with tf.gfile.Open(tsv_filename) as tsv_file:
for line in tsv_file:
if line and "\t" in line:
parts = line.split("\t")
source, target = parts[src_column], parts[trg_column]
source, target = source.strip(), target.strip()
clean_pairs = [(source, target)]
if "tsv" in datatypes_to_clean:
clean_pairs = cleaner_en_xx.clean_en_xx_pairs(clean_pairs)
for source, target in clean_pairs:
if source and target:
lang1_resfile.write(source)
lang1_resfile.write("\n")
lang2_resfile.write(target)
lang2_resfile.write("\n")
else:
lang1_filename, lang2_filename = dataset[1]
lang1_filepath = os.path.join(tmp_dir, lang1_filename)
lang2_filepath = os.path.join(tmp_dir, lang2_filename)
is_sgm = (
lang1_filename.endswith("sgm") and lang2_filename.endswith("sgm"))
if not (tf.gfile.Exists(lang1_filepath) and
tf.gfile.Exists(lang2_filepath)):
# For .tar.gz and .tgz files, we read compressed.
mode = "r:gz" if compressed_filepath.endswith("gz") else "r"
with tarfile.open(compressed_filepath, mode) as corpus_tar:
corpus_tar.extractall(tmp_dir)
if lang1_filepath.endswith(".gz"):
new_filepath = lang1_filepath.strip(".gz")
generator_utils.gunzip_file(lang1_filepath, new_filepath)
lang1_filepath = new_filepath
if lang2_filepath.endswith(".gz"):
new_filepath = lang2_filepath.strip(".gz")
generator_utils.gunzip_file(lang2_filepath, new_filepath)
lang2_filepath = new_filepath
for example in text_problems.text2text_txt_iterator(
lang1_filepath, lang2_filepath):
line1res = _preprocess_sgm(example["inputs"], is_sgm)
line2res = _preprocess_sgm(example["targets"], is_sgm)
clean_pairs = [(line1res, line2res)]
if "txt" in datatypes_to_clean:
clean_pairs = cleaner_en_xx.clean_en_xx_pairs(clean_pairs)
for line1res, line2res in clean_pairs:
if line1res and line2res:
lang1_resfile.write(line1res)
lang1_resfile.write("\n")
lang2_resfile.write(line2res)
lang2_resfile.write("\n")
return filename
class TranslateDistillProblem(TranslateProblem):
"""Base class for translation problems."""
@property
def is_generate_per_split(self):
return True
def example_reading_spec(self):
data_fields = {"dist_targets": tf.VarLenFeature(tf.int64)}
if self.has_inputs:
data_fields["inputs"] = tf.VarLenFeature(tf.int64)
# hack: ignoring true targets and putting dist_targets in targets
data_items_to_decoders = {
"inputs": contrib.slim().tfexample_decoder.Tensor("inputs"),
"targets": contrib.slim().tfexample_decoder.Tensor("dist_targets"),
}
return (data_fields, data_items_to_decoders)
def get_or_create_vocab(self, data_dir, tmp_dir, force_get=False):
"""Get vocab for distill problems."""
# We assume that vocab file is present in data_dir directory where the
# data generated will be stored.
vocab_filepath = os.path.join(data_dir, self.vocab_filename)
encoder = text_encoder.SubwordTextEncoder(vocab_filepath)
return encoder
def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
generator = self.generate_samples(data_dir, tmp_dir, dataset_split)
vocab = self.get_or_create_vocab(data_dir, tmp_dir)
# For each example, encode the text and append EOS ID.
for sample in generator:
if self.has_inputs:
sample["inputs"] = vocab.encode(sample["inputs"])
sample["inputs"].append(text_encoder.EOS_ID)
sample["targets"] = vocab.encode(sample["targets"])
sample["targets"].append(text_encoder.EOS_ID)
sample["dist_targets"] = vocab.encode(sample["dist_targets"])
sample["dist_targets"].append(text_encoder.EOS_ID)
yield sample
def generate_samples(self, data_dir, tmp_dir, dataset_split):
data_path = self.source_data_files(dataset_split)
assert tf.gfile.Exists(data_path)
return text_problems.text2text_distill_iterator(data_path + "inputs",
data_path + "gold",
data_path + "prediction")
class TranslateWmt20Problem(TranslateProblem):
"""Base class for WMT20 Datasets."""
@property
def is_generate_per_split(self):
return True
def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
generator = self.generate_samples(data_dir, tmp_dir, dataset_split)
vocab = self.get_or_create_vocab(data_dir, tmp_dir)
# For each example, encode the text and append EOS ID.
for sample in generator:
if self.has_inputs:
sample["inputs"] = vocab.encode(sample["inputs"])
sample["inputs"].append(text_encoder.EOS_ID)
sample["targets"] = vocab.encode(sample["targets"])
sample["targets"].append(text_encoder.EOS_ID)
yield sample
def generate_text_for_vocab(self, data_dir, tmp_dir):
for i, sample in enumerate(
self.generate_samples(data_dir, tmp_dir, problem.DatasetSplit.TRAIN)):
if self.has_inputs:
yield sample["inputs"]
yield sample["targets"]
if self.max_samples_for_vocab and (i + 1) >= self.max_samples_for_vocab:
break
def generate_samples(self, data_dir, tmp_dir, dataset_split):
data_path = self.source_data_files(dataset_split)[0]
return text_problems.text2text_txt_tab_iterator(data_path)
class TranslateSamanantarProblem(TranslateWmt20Problem):
"""Base class for Samanantar Datasets."""
def generate_samples(self, data_dir, tmp_dir, dataset_split):
src_data_path = self.source_data_files(dataset_split)[0]
tgt_data_path = self.source_data_files(dataset_split)[1]
return text_problems.text2text_txt_iterator(src_data_path, tgt_data_path)
================================================
FILE: tensor2tensor/data_generators/translate_encs.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Data generators for translation data-sets."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from tensor2tensor.data_generators import problem
from tensor2tensor.data_generators import text_encoder
from tensor2tensor.data_generators import text_problems
from tensor2tensor.data_generators import translate
from tensor2tensor.utils import registry
# End-of-sentence marker.
EOS = text_encoder.EOS_ID
_ENCS_TRAIN_DATASETS = [
[("https://lindat.mff.cuni.cz/repository/xmlui/bitstream/handle/"
"11234/1-1458/data-plaintext-format.tar"),
("tsv", 3, 2, "data.plaintext-format/*train.gz")],
[
"http://data.statmt.org/wmt18/translation-task/training-parallel-nc-v13.tgz", # pylint: disable=line-too-long
("training-parallel-nc-v13/news-commentary-v13.cs-en.en",
"training-parallel-nc-v13/news-commentary-v13.cs-en.cs")
],
[
"http://www.statmt.org/wmt13/training-parallel-commoncrawl.tgz",
("commoncrawl.cs-en.en", "commoncrawl.cs-en.cs")
],
[
"http://www.statmt.org/wmt13/training-parallel-europarl-v7.tgz",
("training/europarl-v7.cs-en.en", "training/europarl-v7.cs-en.cs")
],
]
_ENCS_TEST_DATASETS = [
[
"http://data.statmt.org/wmt17/translation-task/dev.tgz",
("dev/newstest2013.en", "dev/newstest2013.cs")
],
]
@registry.register_problem
class TranslateEncsWmt32k(translate.TranslateProblem):
"""Problem spec for WMT English-Czech translation."""
@property
def approx_vocab_size(self):
return 2**15 # 32768
def source_data_files(self, dataset_split):
train = dataset_split == problem.DatasetSplit.TRAIN
return _ENCS_TRAIN_DATASETS if train else _ENCS_TEST_DATASETS
def vocab_data_files(self):
datasets = self.source_data_files(problem.DatasetSplit.TRAIN)
vocab_datasets = []
if datasets[0][0].endswith("data-plaintext-format.tar"):
vocab_datasets.append([
datasets[0][0], [
"%s-compiled-train.lang1" % self.name,
"%s-compiled-train.lang2" % self.name
]
])
datasets = datasets[1:]
vocab_datasets += [[item[0], [item[1][0], item[1][1]]] for item in datasets]
return vocab_datasets
@registry.register_problem
class TranslateEncsWmtCharacters(translate.TranslateProblem):
"""Problem spec for WMT En-Cs character-based translation."""
@property
def vocab_type(self):
return text_problems.VocabType.CHARACTER
def generate_samples(self, data_dir, tmp_dir, dataset_split):
train = dataset_split == problem.DatasetSplit.TRAIN
datasets = _ENCS_TRAIN_DATASETS if train else _ENCS_TEST_DATASETS
tag = "train" if train else "dev"
data_path = translate.compile_data(tmp_dir, datasets,
"wmt_encs_chr_%s" % tag)
return text_problems.text2text_txt_iterator(data_path + ".lang1",
data_path + ".lang2")
================================================
FILE: tensor2tensor/data_generators/translate_encs_cubbitt.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Data generators for English-Czech backtranslation NMT data-sets.
To use this problem you need to provide backtranslated (synthetic) data to
tmp_dir (cs_mono_{en,cs}.txt{0,1,2} - each file of a similar size to the
authentic training data).
You can either translate the monolingual data yourself or you can download
"csmono" data from CzEng2.0 (http://ufal.mff.cuni.cz/czeng, registration needed)
which comes with synthetic translations into English using a
backtranslation-trained model, thus the final model will be using
"iterated" backtranslation.
To get the best results out of the Block-Backtranslation
(where blocks of synthetic and authentic training data are concatenated
without shuffling), you should use checkpoint averaging (see t2t-avg-all).
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
from tensor2tensor.data_generators import problem
from tensor2tensor.data_generators import text_problems
from tensor2tensor.data_generators import translate
from tensor2tensor.data_generators import translate_encs
from tensor2tensor.utils import registry
@registry.register_problem
class TranslateEncsCubbitt(translate_encs.TranslateEncsWmt32k):
"""Problem spec for English-Czech CUBBITT (CUni Block-Backtranslation-Improved Transformer Translation)."""
@property
def use_vocab_from_other_problem(self):
return translate_encs.TranslateEncsWmt32k()
@property
def already_shuffled(self):
return True
@property
def skip_random_fraction_when_training(self):
return False
@property
def backtranslate_data_filenames(self):
"""List of pairs of files with matched back-translated data."""
# Files must be placed in tmp_dir, each similar size to authentic data.
return [("cs_mono_en.txt%d" % i, "cs_mono_cs.txt%d" % i) for i in [0, 1, 2]]
@property
def dataset_splits(self):
"""Splits of data to produce and number of output shards for each."""
return [{
"split": problem.DatasetSplit.TRAIN,
"shards": 1, # Use just 1 shard so as to not mix data.
}, {
"split": problem.DatasetSplit.EVAL,
"shards": 1,
}]
def generate_samples(self, data_dir, tmp_dir, dataset_split):
datasets = self.source_data_files(dataset_split)
tag = "train" if dataset_split == problem.DatasetSplit.TRAIN else "dev"
data_path = translate.compile_data(
tmp_dir, datasets, "%s-compiled-%s" % (self.name, tag))
# For eval, use authentic data.
if dataset_split != problem.DatasetSplit.TRAIN:
for example in text_problems.text2text_txt_iterator(
data_path + ".lang1", data_path + ".lang2"):
yield example
else: # For training, mix synthetic and authentic data as follows.
for (file1, file2) in self.backtranslate_data_filenames:
path1 = os.path.join(tmp_dir, file1)
path2 = os.path.join(tmp_dir, file2)
# Synthetic data first.
for example in text_problems.text2text_txt_iterator(path1, path2):
yield example
# Now authentic data.
for example in text_problems.text2text_txt_iterator(
data_path + ".lang1", data_path + ".lang2"):
yield example
================================================
FILE: tensor2tensor/data_generators/translate_ende.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Data generators for translation data-sets."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from tensor2tensor.data_generators import problem
from tensor2tensor.data_generators import text_problems
from tensor2tensor.data_generators import translate
from tensor2tensor.data_generators import wiki_lm
from tensor2tensor.utils import registry
_ENDE_TRAIN_DATASETS = [
[
"http://data.statmt.org/wmt18/translation-task/training-parallel-nc-v13.tgz", # pylint: disable=line-too-long
("training-parallel-nc-v13/news-commentary-v13.de-en.en",
"training-parallel-nc-v13/news-commentary-v13.de-en.de")
],
[
"http://www.statmt.org/wmt13/training-parallel-commoncrawl.tgz",
("commoncrawl.de-en.en", "commoncrawl.de-en.de")
],
[
"http://www.statmt.org/wmt13/training-parallel-europarl-v7.tgz",
("training/europarl-v7.de-en.en", "training/europarl-v7.de-en.de")
],
]
_ENDE_EVAL_DATASETS = [
[
"http://data.statmt.org/wmt17/translation-task/dev.tgz",
("dev/newstest2013.en", "dev/newstest2013.de")
],
]
_ENDE_RAPID_TRAIN_DATASET = [
# additional training data available for WMT 18 news task training data
# as defined by http://www.statmt.org/wmt18/translation-task.html
[
"http://data.statmt.org/wmt18/translation-task/rapid2016.tgz",
("rapid2016.de-en.en", "rapid2016.de-en.de"),
],
]
_ENDE_PARACRAWL_DATASETS = [
[
"https://s3.amazonaws.com/web-language-models/paracrawl/release4/en-de.bicleaner07.tmx.gz", # pylint: disable=line-too-long
("tmx", "en-de.bicleaner07.tmx.gz")
]
]
@registry.register_problem
class TranslateEndeWmt32k(translate.TranslateProblem):
"""En-de translation trained on WMT corpus."""
@property
def additional_training_datasets(self):
"""Allow subclasses to add training datasets."""
return []
def source_data_files(self, dataset_split):
train = dataset_split == problem.DatasetSplit.TRAIN
train_datasets = _ENDE_TRAIN_DATASETS + self.additional_training_datasets
return train_datasets if train else _ENDE_EVAL_DATASETS
@registry.register_problem
class TranslateEnde2018Wmt32k(translate.TranslateProblem):
"""En-de translation trained on WMT18 corpus."""
@property
def use_vocab_from_other_problem(self):
return TranslateEndeWmt32k()
@property
def additional_training_datasets(self):
"""WMT18 adds rapid data."""
return _ENDE_RAPID_TRAIN_DATASET
@registry.register_problem
class TranslateEndeWmtClean32k(TranslateEndeWmt32k):
"""En-de translation trained on WMT with further cleaning."""
@property
def use_vocab_from_other_problem(self):
return TranslateEndeWmt32k()
@property
def datatypes_to_clean(self):
return ["txt"]
@registry.register_problem
class TranslateEndePc32k(translate.TranslateProblem):
"""En-de translation trained on Paracrawl (bicleaner corpus)."""
@property
def use_vocab_from_other_problem(self):
return TranslateEndeWmt32k()
@property
def additional_training_datasets(self):
"""Allow subclasses to add training datasets."""
return []
def source_data_files(self, dataset_split):
train = dataset_split == problem.DatasetSplit.TRAIN
train_datasets = (
_ENDE_PARACRAWL_DATASETS + self.additional_training_datasets)
return train_datasets if train else _ENDE_EVAL_DATASETS
@registry.register_problem
class TranslateEndePcClean32k(TranslateEndePc32k):
"""En-de translation trained on Paracrawl with further cleaning."""
@property
def datatypes_to_clean(self):
return ["tmx"]
@registry.register_problem
class TranslateEndeWmtPc32k(TranslateEndeWmt32k):
"""En-de translation trained on WMT plus Paracrawl."""
@property
def use_vocab_from_other_problem(self):
return TranslateEndeWmt32k()
@property
def additional_training_datasets(self):
return _ENDE_PARACRAWL_DATASETS
@registry.register_problem
class TranslateEndeWmtCleanPc32k(TranslateEndeWmtPc32k):
"""En-de translation trained on cleaned WMT plus Paracrawl."""
@property
def datatypes_to_clean(self):
return ["txt"]
@registry.register_problem
class TranslateEndeWmtPcClean32k(TranslateEndeWmtPc32k):
"""En-de translation trained on WMT plus cleaned Paracrawl."""
@property
def datatypes_to_clean(self):
return ["tmx"]
@registry.register_problem
class TranslateEndeWmtCleanPcClean32k(TranslateEndeWmtPcClean32k):
"""En-de translation trained on cleaned WMT plus cleaned Paracrawl."""
@property
def datatypes_to_clean(self):
return ["txt", "tmx"]
@registry.register_problem
class TranslateEndeWmt32kPacked(TranslateEndeWmt32k):
@property
def packed_length(self):
return 256
@property
def use_vocab_from_other_problem(self):
return TranslateEndeWmt32k()
@registry.register_problem
class TranslateEndeWmt8k(TranslateEndeWmt32k):
"""Problem spec for WMT En-De translation."""
@property
def approx_vocab_size(self):
return 2**13 # 8192
@registry.register_problem
class TranslateEndeWmt8kPacked(TranslateEndeWmt8k):
@property
def packed_length(self):
return 256
@property
def use_vocab_from_other_problem(self):
return TranslateEndeWmt8k()
@registry.register_problem
class TranslateEndeWmtCharacters(TranslateEndeWmt8k):
"""Problem spec for WMT En-De translation."""
@property
def vocab_type(self):
return text_problems.VocabType.CHARACTER
@registry.register_problem
class TranslateEndeWmtMulti64k(TranslateEndeWmt8k):
"""Translation with muli-lingual vocabulary."""
@property
def use_vocab_from_other_problem(self):
return wiki_lm.LanguagemodelDeEnFrRoWiki64k()
@registry.register_problem
class TranslateEndeWmtMulti64kPacked1k(TranslateEndeWmtMulti64k):
"""Translation with muli-lingual vocabulary."""
@property
def packed_length(self):
return 1024
@property
def num_training_examples(self):
return 173800
@property
def inputs_prefix(self):
return "translate English German "
@property
def targets_prefix(self):
return "translate German English "
================================================
FILE: tensor2tensor/data_generators/translate_ende_test.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Tests for tensor2tensor.data_generators.translate_ende."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from tensor2tensor.data_generators import problem
from tensor2tensor.data_generators import translate_ende
import tensorflow.compat.v1 as tf
class TranslateEndeTest(tf.test.TestCase):
"""Tests that some TranslateEnde subclasses inherit information correctly."""
def test_vocab_size(self):
wmt_8k = translate_ende.TranslateEndeWmt8k()
wmt_32k = translate_ende.TranslateEndeWmt32k()
self.assertEqual(wmt_8k.approx_vocab_size, 8192)
self.assertEqual(wmt_32k.approx_vocab_size, 32768)
def test_additional_datasets(self):
wmt_8k = translate_ende.TranslateEndeWmt8k()
wmt_32k = translate_ende.TranslateEndeWmt32k()
self.assertListEqual(wmt_8k.additional_training_datasets, [])
self.assertListEqual(wmt_32k.additional_training_datasets, [])
def test_source_data_files(self):
wmt_8k = translate_ende.TranslateEndeWmt8k()
wmt_32k = translate_ende.TranslateEndeWmt32k()
eval_split = problem.DatasetSplit.EVAL
train_split = problem.DatasetSplit.TRAIN
wmt_8k_eval_files = wmt_8k.source_data_files(eval_split)
wmt_32k_eval_files = wmt_32k.source_data_files(eval_split)
self.assertListEqual(wmt_8k_eval_files, wmt_32k_eval_files)
self.assertGreater(len(wmt_8k_eval_files), 0)
wmt_8k_train_files = wmt_8k.source_data_files(train_split)
wmt_32k_train_files = wmt_32k.source_data_files(train_split)
self.assertListEqual(wmt_8k_train_files, wmt_32k_train_files)
self.assertGreater(len(wmt_8k_train_files), 0)
if __name__ == '__main__':
tf.test.main()
================================================
FILE: tensor2tensor/data_generators/translate_enes.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Data generators for translation data-sets."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from tensor2tensor.data_generators import problem
from tensor2tensor.data_generators import text_encoder
from tensor2tensor.data_generators import text_problems
from tensor2tensor.data_generators import translate
from tensor2tensor.utils import registry
# End-of-sentence marker.
EOS = text_encoder.EOS_ID
_ENES_TRAIN_DATASETS = [
[
"http://www.statmt.org/wmt13/training-parallel-commoncrawl.tgz",
("commoncrawl.es-en.en", "commoncrawl.es-en.es")
],
[
"http://www.statmt.org/wmt13/training-parallel-europarl-v7.tgz",
("training/europarl-v7.es-en.en", "training/europarl-v7.es-en.es")
],
[
"http://www.statmt.org/wmt13/training-parallel-un.tgz",
("un/undoc.2000.es-en.en", "un/undoc.2000.es-en.es")
],
[
"https://s3.amazonaws.com/web-language-models/paracrawl/release1/paracrawl-release1.en-es.zipporah0-dedup-clean.tgz",
("paracrawl-release1.en-es.zipporah0-dedup-clean.en",
"paracrawl-release1.en-es.zipporah0-dedup-clean.es")
]
]
_ENES_TEST_DATASETS = [
[
"http://data.statmt.org/wmt17/translation-task/dev.tgz",
("dev/newstest2013.en", "dev/newstest2013.es")
],
]
@registry.register_problem
class TranslateEnesWmt32k(translate.TranslateProblem):
"""En-es translation trained on WMT corpus."""
@property
def additional_training_datasets(self):
"""Allow subclasses to add training datasets."""
return []
def source_data_files(self, dataset_split):
train = dataset_split == problem.DatasetSplit.TRAIN
train_datasets = _ENES_TRAIN_DATASETS + self.additional_training_datasets
return train_datasets if train else _ENES_TEST_DATASETS
def vocab_data_files(self):
return _ENES_TRAIN_DATASETS
@registry.register_problem
class TranslateEnesWmtClean32k(TranslateEnesWmt32k):
"""En-es translation trained on WMT with further cleaning."""
@property
def use_vocab_from_other_problem(self):
return TranslateEnesWmt32k()
@property
def datatypes_to_clean(self):
return ["txt"]
@registry.register_problem
class TranslateEnesWmt32kPacked(TranslateEnesWmt32k):
@property
def packed_length(self):
return 256
@property
def use_vocab_from_other_problem(self):
return TranslateEnesWmt32k()
@registry.register_problem
class TranslateEnesWmt8k(TranslateEnesWmt32k):
"""Problem spec for WMT En-Es translation."""
@property
def approx_vocab_size(self):
return 2**13 # 8192
@registry.register_problem
class TranslateEnesWmt8kPacked(TranslateEnesWmt8k):
@property
def packed_length(self):
return 256
@property
def use_vocab_from_other_problem(self):
return TranslateEnesWmt8k()
@registry.register_problem
class TranslateEnesWmtCharacters(TranslateEnesWmt8k):
"""Problem spec for WMT En-Es translation."""
@property
def vocab_type(self):
return text_problems.VocabType.CHARACTER
================================================
FILE: tensor2tensor/data_generators/translate_enet.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Data generators for En-Et translation."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from tensor2tensor.data_generators import problem
from tensor2tensor.data_generators import text_encoder
from tensor2tensor.data_generators import text_problems
from tensor2tensor.data_generators import translate
from tensor2tensor.utils import registry
# End-of-sentence marker.
EOS = text_encoder.EOS_ID
# For English-Estonian the WMT18 data is used
# The complete corpus has ~ 2,18M sentences
_ENET_TRAIN_DATASETS = [
[
"http://data.statmt.org/wmt18/translation-task/training-parallel-ep-v8.tgz", # pylint: disable=line-too-long
("training/europarl-v8.et-en.en", "training/europarl-v8.et-en.et")
],
[
"https://s3.amazonaws.com/web-language-models/paracrawl/release1/paracrawl-release1.en-et.zipporah0-dedup-clean.tgz", # pylint: disable=line-too-long
("paracrawl-release1.en-et.zipporah0-dedup-clean.en",
"paracrawl-release1.en-et.zipporah0-dedup-clean.et")
],
[
"http://data.statmt.org/wmt18/translation-task/rapid2016.tgz",
("rapid2016.en-et.en", "rapid2016.en-et.et")
],
]
# For development 2,000 parallel sentences are used
_ENET_TEST_DATASETS = [[
"https://github.com/stefan-it/nmt-en-et/raw/master/data/newsdev2018-enet.tar.gz", # pylint: disable=line-too-long
("newsdev2018-enet-src.en", "newsdev2018-enet-ref.et")
]]
@registry.register_problem
class TranslateEnetWmt32k(translate.TranslateProblem):
"""Problem spec for WMT18 En-Et translation."""
@property
def approx_vocab_size(self):
return 2**15 # 32768
def source_data_files(self, dataset_split):
train = dataset_split == problem.DatasetSplit.TRAIN
return _ENET_TRAIN_DATASETS if train else _ENET_TEST_DATASETS
@registry.register_problem
class TranslateEnetWmtCharacters(translate.TranslateProblem):
"""Problem spec for WMT18 En-Et translation."""
@property
def vocab_type(self):
return text_problems.VocabType.CHARACTER
def source_data_files(self, dataset_split):
train = dataset_split == problem.DatasetSplit.TRAIN
return _ENET_TRAIN_DATASETS if train else _ENET_TEST_DATASETS
================================================
FILE: tensor2tensor/data_generators/translate_enfr.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Data generators for translation data-sets."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
from tensor2tensor.data_generators import problem
from tensor2tensor.data_generators import text_encoder
from tensor2tensor.data_generators import text_problems
from tensor2tensor.data_generators import translate
from tensor2tensor.data_generators import wiki_lm
from tensor2tensor.utils import registry
# End-of-sentence marker.
EOS = text_encoder.EOS_ID
_ENFR_TRAIN_SMALL_DATA = [
[
"https://s3.amazonaws.com/opennmt-trainingdata/baseline-1M-enfr.tgz",
("baseline-1M-enfr/baseline-1M_train.en",
"baseline-1M-enfr/baseline-1M_train.fr")
],
]
_ENFR_TEST_SMALL_DATA = [
[
"https://s3.amazonaws.com/opennmt-trainingdata/baseline-1M-enfr.tgz",
("baseline-1M-enfr/baseline-1M_valid.en",
"baseline-1M-enfr/baseline-1M_valid.fr")
],
]
_ENFR_TRAIN_LARGE_DATA = [
[
"http://www.statmt.org/wmt13/training-parallel-commoncrawl.tgz",
("commoncrawl.fr-en.en", "commoncrawl.fr-en.fr")
],
[
"http://www.statmt.org/wmt13/training-parallel-europarl-v7.tgz",
("training/europarl-v7.fr-en.en", "training/europarl-v7.fr-en.fr")
],
[
"http://www.statmt.org/wmt14/training-parallel-nc-v9.tgz",
("training/news-commentary-v9.fr-en.en",
"training/news-commentary-v9.fr-en.fr")
],
[
"http://www.statmt.org/wmt10/training-giga-fren.tar",
("giga-fren.release2.fixed.en.gz",
"giga-fren.release2.fixed.fr.gz")
],
[
"http://www.statmt.org/wmt13/training-parallel-un.tgz",
("un/undoc.2000.fr-en.en", "un/undoc.2000.fr-en.fr")
],
]
_ENFR_TEST_LARGE_DATA = [
[
"http://data.statmt.org/wmt17/translation-task/dev.tgz",
("dev/newstest2013.en", "dev/newstest2013.fr")
],
]
@registry.register_problem
class TranslateEnfrWmtSmall8k(translate.TranslateProblem):
"""Problem spec for WMT En-Fr translation."""
@property
def approx_vocab_size(self):
return 2**13 # 8192
@property
def use_small_dataset(self):
return True
def source_data_files(self, dataset_split):
train = dataset_split == problem.DatasetSplit.TRAIN
if self.use_small_dataset:
datasets = _ENFR_TRAIN_SMALL_DATA if train else _ENFR_TEST_SMALL_DATA
else:
datasets = _ENFR_TRAIN_LARGE_DATA if train else _ENFR_TEST_LARGE_DATA
return datasets
def vocab_data_files(self):
return (_ENFR_TRAIN_SMALL_DATA if self.use_small_dataset
else _ENFR_TRAIN_LARGE_DATA)
@registry.register_problem
class TranslateEnfrWmtSmall32k(TranslateEnfrWmtSmall8k):
@property
def approx_vocab_size(self):
return 2**15 # 32768
@registry.register_problem
class TranslateEnfrWmt8k(TranslateEnfrWmtSmall8k):
@property
def use_small_dataset(self):
return False
@registry.register_problem
class TranslateEnfrWmt32k(TranslateEnfrWmtSmall32k):
@property
def use_small_dataset(self):
return False
@registry.register_problem
class TranslateEnfrWmt32kPacked(TranslateEnfrWmt32k):
@property
def packed_length(self):
return 256
@property
def use_vocab_from_other_problem(self):
return TranslateEnfrWmt32k()
@registry.register_problem
class TranslateEnfrWmt32kWithBacktranslateFr(TranslateEnfrWmt32k):
"""En-Fr translation with added French data, back-translated."""
@property
def use_vocab_from_other_problem(self):
return TranslateEnfrWmt32k()
@property
def already_shuffled(self):
return True
@property
def skip_random_fraction_when_training(self):
return False
@property
def backtranslate_data_filenames(self):
"""List of pairs of files with matched back-translated data."""
# Files must be placed in tmp_dir, each similar size to authentic data.
return [("fr_mono_en.txt", "fr_mono_fr.txt")]
@property
def dataset_splits(self):
"""Splits of data to produce and number of output shards for each."""
return [{
"split": problem.DatasetSplit.TRAIN,
"shards": 1, # Use just 1 shard so as to not mix data.
}, {
"split": problem.DatasetSplit.EVAL,
"shards": 1,
}]
def generate_samples(self, data_dir, tmp_dir, dataset_split):
datasets = self.source_data_files(dataset_split)
tag = "train" if dataset_split == problem.DatasetSplit.TRAIN else "dev"
data_path = translate.compile_data(
tmp_dir, datasets, "%s-compiled-%s" % (self.name, tag))
# For eval, use authentic data.
if dataset_split != problem.DatasetSplit.TRAIN:
for example in text_problems.text2text_txt_iterator(
data_path + ".lang1", data_path + ".lang2"):
yield example
else: # For training, mix synthetic and authentic data as follows.
for (file1, file2) in self.backtranslate_data_filenames:
path1 = os.path.join(tmp_dir, file1)
path2 = os.path.join(tmp_dir, file2)
# Synthetic data first.
for example in text_problems.text2text_txt_iterator(path1, path2):
yield example
# Now authentic data.
for example in text_problems.text2text_txt_iterator(
data_path + ".lang1", data_path + ".lang2"):
yield example
@registry.register_problem
class TranslateEnfrWmt32kWithBacktranslateEn(
TranslateEnfrWmt32kWithBacktranslateFr):
"""En-Fr translation with added English data, back-translated."""
@property
def backtranslate_data_filenames(self):
"""List of pairs of files with matched back-translated data."""
# Files must be placed in tmp_dir, each similar size to authentic data.
return [("en_mono_en.txt%d" % i, "en_mono_fr.txt%d" % i) for i in [0, 1, 2]]
@registry.register_problem
class TranslateEnfrWmtSmallCharacters(translate.TranslateProblem):
"""Problem spec for WMT En-Fr translation."""
@property
def vocab_type(self):
return text_problems.VocabType.CHARACTER
@property
def use_small_dataset(self):
return True
def source_data_files(self, dataset_split):
train = dataset_split == problem.DatasetSplit.TRAIN
if self.use_small_dataset:
datasets = _ENFR_TRAIN_SMALL_DATA if train else _ENFR_TEST_SMALL_DATA
else:
datasets = _ENFR_TRAIN_LARGE_DATA if train else _ENFR_TEST_LARGE_DATA
return datasets
@registry.register_problem
class TranslateEnfrWmtCharacters(TranslateEnfrWmtSmallCharacters):
@property
def use_small_dataset(self):
return False
@registry.register_problem
class TranslateEnfrWmtMulti64k(TranslateEnfrWmtSmall32k):
"""Translation with muli-lingual vocabulary."""
@property
def use_small_dataset(self):
return False
@property
def use_vocab_from_other_problem(self):
return wiki_lm.LanguagemodelDeEnFrRoWiki64k()
@registry.register_problem
class TranslateEnfrWmtMulti64kPacked1k(TranslateEnfrWmtMulti64k):
"""Translation with muli-lingual vocabulary."""
@property
def packed_length(self):
return 1024
@property
def num_training_examples(self):
return 1760600
@property
def inputs_prefix(self):
return "translate English French "
@property
def targets_prefix(self):
return "translate French English "
================================================
FILE: tensor2tensor/data_generators/translate_enid.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Data generators for En-Id translation."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
# Dependency imports
from tensor2tensor.data_generators import problem
from tensor2tensor.data_generators import text_encoder
from tensor2tensor.data_generators import translate
from tensor2tensor.utils import registry
# End-of-sentence marker.
EOS = text_encoder.EOS_ID
_REPO = "https://github.com/prasastoadi/parallel-corpora-en-id/raw/master/"
# IWSLT17 :
# 109335 sentences
# https://wit3.fbk.eu/mt.php?release=2017-01-more
# PANL-BPPT :
# 24024 sentences
# http://www.panl10n.net/english/outputs/Indonesia/BPPT/0902/BPPTIndToEngCorpusHalfM.zip # pylint: disable=line-too-long
_ENID_TRAIN_DATASETS = [
[
_REPO + "IWSLT17.train.en-id.tgz",
("IWSLT17.train.en-id.en", "IWSLT17.train.en-id.id")
],
[
_REPO + "PANL-BPPT-ECO-EN-ID-150Kw.tgz",
("PANL-BPPT-ECO-EN-150Kw.txt", "PANL-BPPT-ECO-ID-150Kw.txt")
],
[
_REPO + "PANL-BPPT-INT-EN-ID-150Kw.tgz",
("PANL-BPPT-INT-EN-150Kw.txt", "PANL-BPPT-INT-ID-150Kw.txt")
],
[
_REPO + "PANL-BPPT-SCI-EN-ID-100Kw.tgz",
("PANL-BPPT-SCI-EN-100Kw.txt", "PANL-BPPT-SCI-ID-100Kw.txt")
],
[
_REPO + "PANL-BPPT-SPO-EN-ID-100Kw.tgz",
("PANL-BPPT-SPO-EN-100Kw.txt", "PANL-BPPT-SPO-ID-100Kw.txt")
],
]
# IWSLT17 :
# 1478 sentences
# https://wit3.fbk.eu/mt.php?release=2017-01-more
_ENID_TEST_DATASETS = [
[
_REPO + "IWSLT17.TED.tst2017plus.en-id.tgz",
("IWSLT17.TED.tst2017plus.en-id.en",
"IWSLT17.TED.tst2017plus.en-id.id")
]
]
@registry.register_problem
class TranslateEnidIwslt32k(translate.TranslateProblem):
"""Problem spec for IWSLT'15 En-Vi translation."""
@property
def approx_vocab_size(self):
return 2**15 # 32768
def source_data_files(self, dataset_split):
train = dataset_split == problem.DatasetSplit.TRAIN
return _ENID_TRAIN_DATASETS if train else _ENID_TEST_DATASETS
================================================
FILE: tensor2tensor/data_generators/translate_enmk.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Data generators for translation data-sets."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from tensor2tensor.data_generators import problem
from tensor2tensor.data_generators import text_encoder
from tensor2tensor.data_generators import text_problems
from tensor2tensor.data_generators import translate
from tensor2tensor.utils import registry
# End-of-sentence marker.
EOS = text_encoder.EOS_ID
# For English-Macedonian the SETimes corpus
# from http://nlp.ffzg.hr/resources/corpora/setimes/ is used.
_ENMK_TRAIN_DATASETS = [[
"http://nlp.ffzg.hr/data/corpora/setimes/setimes.en-mk.txt.tgz",
("setimes.en-mk.en.txt", "setimes.en-mk.mk.txt")
]]
# For development the MULTEXT-East "1984" corpus from
# https://www.clarin.si/repository/xmlui/handle/11356/1043 is used.
# 4,986 parallel sentences are used for evaluation.
_ENMK_DEV_DATASETS = [[
"https://github.com/stefan-it/nmt-en-mk/raw/master/data/MTE-1984-dev.enmk.tgz", # pylint: disable=line-too-long
("MTE1984-dev.en", "MTE1984-dev.mk")
]]
# See this PR on github for some results with Transformer on these Problems.
# https://github.com/tensorflow/tensor2tensor/pull/738
@registry.register_problem
class TranslateEnmkSetimes32k(translate.TranslateProblem):
"""Problem spec for SETimes En-Mk translation."""
@property
def approx_vocab_size(self):
return 2**15 # 32768
def source_data_files(self, dataset_split):
train = dataset_split == problem.DatasetSplit.TRAIN
return _ENMK_TRAIN_DATASETS if train else _ENMK_DEV_DATASETS
@registry.register_problem
class TranslateEnmkSetimesCharacters(translate.TranslateProblem):
"""Problem spec for SETimes En-Mk translation."""
@property
def vocab_type(self):
return text_problems.VocabType.CHARACTER
def source_data_files(self, dataset_split):
train = dataset_split == problem.DatasetSplit.TRAIN
return _ENMK_TRAIN_DATASETS if train else _ENMK_DEV_DATASETS
================================================
FILE: tensor2tensor/data_generators/translate_enro.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Data generators for translation data-sets."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import random
from tensor2tensor.data_generators import problem
from tensor2tensor.data_generators import text_problems
from tensor2tensor.data_generators import translate
from tensor2tensor.data_generators import wiki_lm
from tensor2tensor.utils import registry
_ENRO_TRAIN_DATASETS = [
[
"http://www.statmt.org/europarl/v7/ro-en.tgz",
("europarl-v7.ro-en.en", "europarl-v7.ro-en.ro")
],
[
"http://opus.nlpl.eu/download.php?f=SETIMES/v2/moses/en-ro.txt.zip",
("SETIMES.en-ro.en", "SETIMES.en-ro.ro")
]
]
_ENRO_TEST_DATASETS = [
[
("http://data.statmt.org/wmt16/translation-task/"
"dev-romanian-updated.tgz"),
("dev/newsdev2016-roen-ref.en.sgm", "dev/newsdev2016-roen-src.ro.sgm")
],
]
@registry.register_problem
class TranslateEnroWmt8k(translate.TranslateProblem):
"""Problem spec for WMT En-Ro translation."""
@property
def approx_vocab_size(self):
return 2**13 # 8192
def source_data_files(self, dataset_split):
train = dataset_split == problem.DatasetSplit.TRAIN
return _ENRO_TRAIN_DATASETS if train else _ENRO_TEST_DATASETS
@registry.register_problem
class TranslateEnroWmt32k(TranslateEnroWmt8k):
@property
def approx_vocab_size(self):
return 2**15 # 32768
@registry.register_problem
class TranslateEnroWmtCharacters(TranslateEnroWmt8k):
"""Problem spec for WMT En-Ro translation."""
@property
def vocab_type(self):
return text_problems.VocabType.CHARACTER
@registry.register_problem
class TranslateEnroWmtMulti64k(TranslateEnroWmt8k):
"""Translation with muli-lingual vocabulary."""
@property
def use_vocab_from_other_problem(self):
return wiki_lm.LanguagemodelDeEnFrRoWiki64k()
@registry.register_problem
class TranslateEnroWmtMultiSmall64k(TranslateEnroWmt8k):
"""Translation with muli-lingual vocabulary, small (6K) training data."""
@property
def dataset_splits(self):
"""Splits of data to produce and number of output shards for each."""
return [{
"split": problem.DatasetSplit.TRAIN,
"shards": 16, # It's a small dataset, TPUs like at least a few shards.
}, {
"split": problem.DatasetSplit.EVAL,
"shards": 1,
}]
@property
def use_vocab_from_other_problem(self):
return wiki_lm.LanguagemodelDeEnFrRoWiki64k()
@property
def how_many_examples_to_sample(self):
return 6000
def generate_samples(self, data_dir, tmp_dir, dataset_split):
"""Generate just the first 6k samples for training."""
# If not training, do the same as before.
if dataset_split != problem.DatasetSplit.TRAIN:
for x in super(TranslateEnroWmtMultiSmall64k, self).generate_samples(
data_dir, tmp_dir, dataset_split):
yield x
raise StopIteration
# Now we assume we're training.
counter = 0
# The size of this data-set in total is around 614K, we want to sample so
# that in expectation we take the requested number of samples in 1 go.
sample_prob = self.how_many_examples_to_sample / float(614000)
# Let's sample.
for x in super(TranslateEnroWmtMultiSmall64k, self).generate_samples(
data_dir, tmp_dir, dataset_split):
if random.random() > sample_prob:
continue
counter += 1
if counter > self.how_many_examples_to_sample:
raise StopIteration
yield x
# We do it again if we don't have enough samples.
if counter < self.how_many_examples_to_sample:
for x in super(TranslateEnroWmtMultiSmall64k, self).generate_samples(
data_dir, tmp_dir, dataset_split):
if random.random() > sample_prob:
continue
counter += 1
if counter > self.how_many_examples_to_sample:
raise StopIteration
yield x
@registry.register_problem
class TranslateEnroWmtMultiTiny64k(TranslateEnroWmtMultiSmall64k):
"""Translation with muli-lingual vocabulary, tiny (600) training data."""
@property
def how_many_examples_to_sample(self):
return 600
@registry.register_problem
class TranslateEnroWmtMultiTiny64kPacked1k(TranslateEnroWmtMultiTiny64k):
"""Translation with muli-lingual vocabulary."""
@property
def packed_length(self):
return 1024
@property
def num_training_examples(self):
return 32
@property
def inputs_prefix(self):
return "translate English Romanian "
@property
def targets_prefix(self):
return "translate Romanian English "
================================================
FILE: tensor2tensor/data_generators/translate_entn.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Data generators for translation data-sets."""
from tensor2tensor.data_generators import problem
from tensor2tensor.data_generators import text_encoder
from tensor2tensor.data_generators import translate
from tensor2tensor.utils import registry
EOS = text_encoder.EOS_ID
_URL = "https://github.com/LauraMartinus/ukuxhumana/blob/master/data/en_tn"
_ENTN_TRAIN_DATASETS = [[
_URL + "/eng_tswane.train.tar.gz?raw=true",
("entn_parallel.train.en", "entn_parallel.train.tn")
]]
_ENTN_TEST_DATASETS = [[
_URL + "/eng_tswane.dev.tar.gz?raw=true",
("entn_parallel.dev.en", "entn_parallel.dev.tn")
]]
@registry.register_problem
class TranslateEntnRma(translate.TranslateProblem):
"""Problem spec for English-Setswana translation.
Uses the RMA Autshumato dataset.
"""
@property
def approx_vocab_size(self):
return 2**15 # 32768
@property
def vocab_filename(self):
return "vocab.entn.%d" % self.approx_vocab_size
def source_data_files(self, dataset_split):
train = dataset_split == problem.DatasetSplit.TRAIN
return _ENTN_TRAIN_DATASETS if train else _ENTN_TEST_DATASETS
================================================
FILE: tensor2tensor/data_generators/translate_envi.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Data generators for En-Vi translation."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from tensor2tensor.data_generators import problem
from tensor2tensor.data_generators import text_encoder
from tensor2tensor.data_generators import translate
from tensor2tensor.utils import registry
# End-of-sentence marker.
EOS = text_encoder.EOS_ID
# For English-Vietnamese the IWSLT'15 corpus
# from https://nlp.stanford.edu/projects/nmt/ is used.
# The original dataset has 133K parallel sentences.
_ENVI_TRAIN_DATASETS = [[
"https://github.com/stefan-it/nmt-en-vi/raw/master/data/train-en-vi.tgz", # pylint: disable=line-too-long
("train.en", "train.vi")
]]
# For development 1,553 parallel sentences are used.
_ENVI_TEST_DATASETS = [[
"https://github.com/stefan-it/nmt-en-vi/raw/master/data/dev-2012-en-vi.tgz", # pylint: disable=line-too-long
("tst2012.en", "tst2012.vi")
]]
# See this PR on github for some results with Transformer on this Problem.
# https://github.com/tensorflow/tensor2tensor/pull/611
@registry.register_problem
class TranslateEnviIwslt32k(translate.TranslateProblem):
"""Problem spec for IWSLT'15 En-Vi translation."""
@property
def approx_vocab_size(self):
return 2**15 # 32768
def source_data_files(self, dataset_split):
train = dataset_split == problem.DatasetSplit.TRAIN
return _ENVI_TRAIN_DATASETS if train else _ENVI_TEST_DATASETS
================================================
FILE: tensor2tensor/data_generators/translate_enzh.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Data generators for translation data-sets."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
from tensor2tensor.data_generators import generator_utils
from tensor2tensor.data_generators import problem
from tensor2tensor.data_generators import text_encoder
from tensor2tensor.data_generators import text_problems
from tensor2tensor.data_generators import translate
from tensor2tensor.utils import registry
import tensorflow.compat.v1 as tf
# End-of-sentence marker.
EOS = text_encoder.EOS_ID
# This is far from being the real WMT18 task - only toyset here
# you need to register to get UN data and CWT data. Also, by convention,
# this is EN to ZH - use translate_enzh_wmt8k_rev for ZH to EN task
#
# News Commentary, around 252k lines
# This dataset is only a small fraction of full WMT18 task
_STAT_MT_URL = "http://data.statmt.org/wmt18/translation-task/"
_NC_TRAIN_DATASETS = [[
_STAT_MT_URL + "training-parallel-nc-v13.tgz", [
"training-parallel-nc-v13/news-commentary-v13.zh-en.en",
"training-parallel-nc-v13/news-commentary-v13.zh-en.zh"
]
]]
# Test set from News Commentary. 2000 lines
_NC_TEST_DATASETS = [[
_STAT_MT_URL + "dev.tgz",
("dev/newsdev2017-enzh-src.en.sgm", "dev/newsdev2017-enzh-ref.zh.sgm")
]]
# UN parallel corpus. 15,886,041 lines
# Visit source website to download manually:
# https://conferences.unite.un.org/UNCorpus
#
# NOTE: You need to register to download dataset from official source
# place into tmp directory e.g. /tmp/t2t_datagen/dataset.tgz
_UN_TRAIN_DATASETS = [[
"https://s3-us-west-2.amazonaws.com/twairball.wmt17.zh-en/UNv1.0.en-zh.tar"
".gz", ["en-zh/UNv1.0.en-zh.en", "en-zh/UNv1.0.en-zh.zh"]
]]
# CWMT corpus
# Visit source website to download manually:
# http://nlp.nju.edu.cn/cwmt-wmt/
#
# casia2015: 1,050,000 lines
# casict2015: 2,036,833 lines
# datum2015: 1,000,003 lines
# datum2017: 1,999,968 lines
# NEU2017: 2,000,000 lines
#
# NOTE: You need to register to download dataset from official source
# place into tmp directory e.g. /tmp/t2t_datagen/dataset.tgz
_CWMT_TRAIN_DATASETS = [[
"https://s3-us-west-2.amazonaws.com/twairball.wmt17.zh-en/cwmt.tgz",
["cwmt/casia2015/casia2015_en.txt", "cwmt/casia2015/casia2015_ch.txt"]
], [
"https://s3-us-west-2.amazonaws.com/twairball.wmt17.zh-en/cwmt.tgz",
["cwmt/casict2015/casict2015_en.txt", "cwmt/casict2015/casict2015_ch.txt"]
], [
"https://s3-us-west-2.amazonaws.com/twairball.wmt17.zh-en/cwmt.tgz",
["cwmt/neu2017/NEU_en.txt", "cwmt/neu2017/NEU_cn.txt"]
], [
"https://s3-us-west-2.amazonaws.com/twairball.wmt17.zh-en/cwmt.tgz",
["cwmt/datum2015/datum_en.txt", "cwmt/datum2015/datum_ch.txt"]
], [
"https://s3-us-west-2.amazonaws.com/twairball.wmt17.zh-en/cwmt.tgz",
["cwmt/datum2017/Book1_en.txt", "cwmt/datum2017/Book1_cn.txt"]
], [
"https://s3-us-west-2.amazonaws.com/twairball.wmt17.zh-en/cwmt.tgz",
["cwmt/datum2017/Book2_en.txt", "cwmt/datum2017/Book2_cn.txt"]
], [
"https://s3-us-west-2.amazonaws.com/twairball.wmt17.zh-en/cwmt.tgz",
["cwmt/datum2017/Book3_en.txt", "cwmt/datum2017/Book3_cn.txt"]
], [
"https://s3-us-west-2.amazonaws.com/twairball.wmt17.zh-en/cwmt.tgz",
["cwmt/datum2017/Book4_en.txt", "cwmt/datum2017/Book4_cn.txt"]
], [
"https://s3-us-west-2.amazonaws.com/twairball.wmt17.zh-en/cwmt.tgz",
["cwmt/datum2017/Book5_en.txt", "cwmt/datum2017/Book5_cn.txt"]
], [
"https://s3-us-west-2.amazonaws.com/twairball.wmt17.zh-en/cwmt.tgz",
["cwmt/datum2017/Book6_en.txt", "cwmt/datum2017/Book6_cn.txt"]
], [
"https://s3-us-west-2.amazonaws.com/twairball.wmt17.zh-en/cwmt.tgz",
["cwmt/datum2017/Book7_en.txt", "cwmt/datum2017/Book7_cn.txt"]
], [
"https://s3-us-west-2.amazonaws.com/twairball.wmt17.zh-en/cwmt.tgz",
["cwmt/datum2017/Book8_en.txt", "cwmt/datum2017/Book8_cn.txt"]
], [
"https://s3-us-west-2.amazonaws.com/twairball.wmt17.zh-en/cwmt.tgz",
["cwmt/datum2017/Book9_en.txt", "cwmt/datum2017/Book9_cn.txt"]
], [
"https://s3-us-west-2.amazonaws.com/twairball.wmt17.zh-en/cwmt.tgz",
["cwmt/datum2017/Book10_en.txt", "cwmt/datum2017/Book10_cn.txt"]
], [
"https://s3-us-west-2.amazonaws.com/twairball.wmt17.zh-en/cwmt.tgz",
["cwmt/datum2017/Book11_en.txt", "cwmt/datum2017/Book11_cn.txt"]
], [
"https://s3-us-west-2.amazonaws.com/twairball.wmt17.zh-en/cwmt.tgz",
["cwmt/datum2017/Book12_en.txt", "cwmt/datum2017/Book12_cn.txt"]
], [
"https://s3-us-west-2.amazonaws.com/twairball.wmt17.zh-en/cwmt.tgz",
["cwmt/datum2017/Book13_en.txt", "cwmt/datum2017/Book13_cn.txt"]
], [
"https://s3-us-west-2.amazonaws.com/twairball.wmt17.zh-en/cwmt.tgz",
["cwmt/datum2017/Book14_en.txt", "cwmt/datum2017/Book14_cn.txt"]
], [
"https://s3-us-west-2.amazonaws.com/twairball.wmt17.zh-en/cwmt.tgz",
["cwmt/datum2017/Book15_en.txt", "cwmt/datum2017/Book15_cn.txt"]
], [
"https://s3-us-west-2.amazonaws.com/twairball.wmt17.zh-en/cwmt.tgz",
["cwmt/datum2017/Book16_en.txt", "cwmt/datum2017/Book16_cn.txt"]
], [
"https://s3-us-west-2.amazonaws.com/twairball.wmt17.zh-en/cwmt.tgz",
["cwmt/datum2017/Book17_en.txt", "cwmt/datum2017/Book17_cn.txt"]
], [
"https://s3-us-west-2.amazonaws.com/twairball.wmt17.zh-en/cwmt.tgz",
["cwmt/datum2017/Book18_en.txt", "cwmt/datum2017/Book18_cn.txt"]
], [
"https://s3-us-west-2.amazonaws.com/twairball.wmt17.zh-en/cwmt.tgz",
["cwmt/datum2017/Book19_en.txt", "cwmt/datum2017/Book19_cn.txt"]
], [
"https://s3-us-west-2.amazonaws.com/twairball.wmt17.zh-en/cwmt.tgz",
["cwmt/datum2017/Book20_en.txt", "cwmt/datum2017/Book20_cn.txt"]
]]
def get_filename(dataset):
return dataset[0][0].split("/")[-1]
@registry.register_problem
class TranslateEnzhWmt32k(translate.TranslateProblem):
"""Problem spec for WMT En-Zh translation.
Attempts to use full training dataset, which needs website
registration and downloaded manually from official sources:
CWMT:
- http://nlp.nju.edu.cn/cwmt-wmt/
- Website contains instructions for FTP server access.
- You'll need to download CASIA, CASICT, DATUM2015, DATUM2017,
NEU datasets
UN Parallel Corpus:
- https://conferences.unite.un.org/UNCorpus
- You'll need to register your to download the dataset.
NOTE: place into tmp directory e.g. /tmp/t2t_datagen/dataset.tgz
"""
@property
def approx_vocab_size(self):
return 2**15 # 32k
@property
def source_vocab_name(self):
return "%s.en" % self.vocab_filename
@property
def target_vocab_name(self):
return "%s.zh" % self.vocab_filename
def get_training_dataset(self, tmp_dir):
"""UN Parallel Corpus and CWMT Corpus need to be downloaded manually.
Append to training dataset if available
Args:
tmp_dir: path to temporary dir with the data in it.
Returns:
paths
"""
full_dataset = _NC_TRAIN_DATASETS
for dataset in [_CWMT_TRAIN_DATASETS, _UN_TRAIN_DATASETS]:
filename = get_filename(dataset)
tmp_filepath = os.path.join(tmp_dir, filename)
if tf.gfile.Exists(tmp_filepath):
full_dataset += dataset
else:
tf.logging.info("[TranslateEzhWmt] dataset incomplete, you need to "
"manually download %s" % filename)
return full_dataset
def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
train = dataset_split == problem.DatasetSplit.TRAIN
train_dataset = self.get_training_dataset(tmp_dir)
datasets = train_dataset if train else _NC_TEST_DATASETS
source_datasets = [[item[0], [item[1][0]]] for item in train_dataset]
target_datasets = [[item[0], [item[1][1]]] for item in train_dataset]
source_vocab = generator_utils.get_or_generate_vocab(
data_dir,
tmp_dir,
self.source_vocab_name,
self.approx_vocab_size,
source_datasets,
file_byte_budget=1e8,
max_subtoken_length=self.max_subtoken_length)
target_vocab = generator_utils.get_or_generate_vocab(
data_dir,
tmp_dir,
self.target_vocab_name,
self.approx_vocab_size,
target_datasets,
file_byte_budget=1e8,
max_subtoken_length=self.max_subtoken_length)
tag = "train" if train else "dev"
filename_base = "wmt_enzh_%sk_tok_%s" % (self.approx_vocab_size, tag)
data_path = translate.compile_data(tmp_dir, datasets, filename_base)
return text_problems.text2text_generate_encoded(
text_problems.text2text_txt_iterator(data_path + ".lang1",
data_path + ".lang2"),
source_vocab, target_vocab)
def feature_encoders(self, data_dir):
source_vocab_filename = os.path.join(data_dir, self.source_vocab_name)
target_vocab_filename = os.path.join(data_dir, self.target_vocab_name)
source_token = text_encoder.SubwordTextEncoder(source_vocab_filename)
target_token = text_encoder.SubwordTextEncoder(target_vocab_filename)
return {
"inputs": source_token,
"targets": target_token,
}
@registry.register_problem
class TranslateEnzhWmt8k(TranslateEnzhWmt32k):
"""Problem spec for WMT En-Zh translation.
This is far from being the real WMT17 task - only toyset here
"""
@property
def approx_vocab_size(self):
return 2**13 # 8192
@property
def dataset_splits(self):
return [
{
"split": problem.DatasetSplit.TRAIN,
"shards": 10, # this is a small dataset
},
{
"split": problem.DatasetSplit.EVAL,
"shards": 1,
}
]
def get_training_dataset(self, tmp_dir):
"""Uses only News Commentary Dataset for training."""
return _NC_TRAIN_DATASETS
================================================
FILE: tensor2tensor/data_generators/translate_test.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Translate generators test."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import shutil
import tarfile
from tensor2tensor.data_generators import text_problems
from tensor2tensor.data_generators import translate
import tensorflow.compat.v1 as tf
class TranslateTest(tf.test.TestCase):
DATASETS = [
["data1.tgz", ("train1.en", "train1.de")],
["data2.tgz", ("train2.en", "train2.de")],
["data3.tgz", ("train3.en", "train3.de")],
]
@classmethod
def setUpClass(cls):
tmp_dir = tf.test.get_temp_dir()
compressed_dir = os.path.join(tmp_dir, "compressed")
shutil.rmtree(tmp_dir)
tf.gfile.MakeDirs(compressed_dir)
en_data = [str(i) for i in range(10, 40)]
de_data = [str(i) for i in range(100, 130)]
data = list(zip(en_data, de_data))
for i, dataset in enumerate(cls.DATASETS):
tar_file = dataset[0]
en_file, de_file = [
os.path.join(compressed_dir, name) for name in dataset[1]
]
with tf.gfile.Open(en_file, "w") as en_f:
with tf.gfile.Open(de_file, "w") as de_f:
start = i * 10
end = start + 10
for en_line, de_line in data[start:end]:
en_f.write(en_line)
en_f.write("\n")
de_f.write(de_line)
de_f.write("\n")
with tarfile.open(os.path.join(tmp_dir, tar_file), "w:gz") as tar_f:
tar_f.add(en_file, os.path.basename(en_file))
tar_f.add(de_file, os.path.basename(de_file))
cls.tmp_dir = tmp_dir
cls.data = data
def testCompileData(self):
filename = "out"
filepath = os.path.join(self.tmp_dir, filename)
translate.compile_data(self.tmp_dir, self.DATASETS, filename)
count = 0
for i, example in enumerate(
text_problems.text2text_txt_iterator(filepath + ".lang1",
filepath + ".lang2")):
expected = self.data[i]
self.assertEqual(list(expected), [example["inputs"], example["targets"]])
count += 1
self.assertEqual(count, len(self.data))
if __name__ == "__main__":
tf.test.main()
================================================
FILE: tensor2tensor/data_generators/video_generated.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Data generators for video problems with artificially generated frames."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import math
import numpy as np
from tensor2tensor.data_generators import video_utils
from tensor2tensor.layers import modalities
from tensor2tensor.utils import contrib
from tensor2tensor.utils import registry
import tensorflow.compat.v1 as tf
try:
import matplotlib # pylint: disable=g-import-not-at-top
matplotlib.use("agg")
import matplotlib.pyplot as plt # pylint: disable=g-import-not-at-top
except ImportError:
pass
@registry.register_problem
class VideoStochasticShapes10k(video_utils.VideoProblem):
"""Shapes moving in a stochastic way."""
@property
def is_generate_per_split(self):
"""Whether we have a train/test split or just hold out data."""
return False # Just hold out some generated data for evals.
@property
def frame_height(self):
return 64
@property
def frame_width(self):
return 64
@property
def total_number_of_frames(self):
# 10k videos
return 10000 * self.video_length
@property
def video_length(self):
return 5
@property
def random_skip(self):
return False
@property
def only_keep_videos_from_0th_frame(self):
return True
@property
def use_not_breaking_batching(self):
return True
def eval_metrics(self):
return []
@property
def extra_reading_spec(self):
"""Additional data fields to store on disk and their decoders."""
data_fields = {
"frame_number": tf.FixedLenFeature([1], tf.int64),
}
decoders = {
"frame_number":
contrib.slim().tfexample_decoder.Tensor(tensor_key="frame_number"),
}
return data_fields, decoders
def hparams(self, defaults, unused_model_hparams):
p = defaults
p.modality = {
"inputs": modalities.ModalityType.VIDEO,
"targets": modalities.ModalityType.VIDEO,
}
p.vocab_size = {
"inputs": 256,
"targets": 256,
}
@staticmethod
def get_circle(x, y, z, c, s):
"""Draws a circle with center(x, y), color c, size s and z-order of z."""
cir = plt.Circle((x, y), s, fc=c, zorder=z)
return cir
@staticmethod
def get_rectangle(x, y, z, c, s):
"""Draws a rectangle with center(x, y), color c, size s and z-order of z."""
rec = plt.Rectangle((x-s, y-s), s*2.0, s*2.0, fc=c, zorder=z)
return rec
@staticmethod
def get_triangle(x, y, z, c, s):
"""Draws a triangle with center (x, y), color c, size s and z-order of z."""
points = np.array([[0, 0], [s, s*math.sqrt(3.0)], [s*2.0, 0]])
tri = plt.Polygon(points + [x-s, y-s], fc=c, zorder=z)
return tri
def generate_stochastic_shape_instance(self):
"""Yields one video of a shape moving to a random direction.
The size and color of the shapes are random but
consistent in a single video. The speed is fixed.
Raises:
ValueError: The frame size is not square.
"""
if self.frame_height != self.frame_width or self.frame_height % 2 != 0:
raise ValueError("Generator only supports square frames with even size.")
lim = 10.0
direction = np.array([[+1.0, +1.0],
[+1.0, +0.0],
[+1.0, -1.0],
[+0.0, +1.0],
[+0.0, -1.0],
[-1.0, +1.0],
[-1.0, +0.0],
[-1.0, -1.0]
])
sp = np.array([lim/2.0, lim/2.0])
rnd = np.random.randint(len(direction))
di = direction[rnd]
colors = ["b", "g", "r", "c", "m", "y"]
color = np.random.choice(colors)
shape = np.random.choice([
VideoStochasticShapes10k.get_circle,
VideoStochasticShapes10k.get_rectangle,
VideoStochasticShapes10k.get_triangle])
speed = 1.0
size = np.random.uniform(0.5, 1.5)
back_color = str(0.0)
plt.ioff()
xy = np.array(sp)
for _ in range(self.video_length):
fig = plt.figure()
fig.set_dpi(self.frame_height//2)
fig.set_size_inches(2, 2)
ax = plt.axes(xlim=(0, lim), ylim=(0, lim))
# Background
ax.add_patch(VideoStochasticShapes10k.get_rectangle(
0.0, 0.0, -1.0, back_color, 25.0))
# Foreground
ax.add_patch(shape(xy[0], xy[1], 0.0, color, size))
plt.axis("off")
plt.tight_layout(pad=-2.0)
fig.canvas.draw()
image = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="")
image = image.reshape(fig.canvas.get_width_height()[::-1] + (3,))
image = np.copy(np.uint8(image))
plt.close()
xy += speed * di
yield image
def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
counter = 0
done = False
while not done:
for frame_number, frame in enumerate(
self.generate_stochastic_shape_instance()):
if counter >= self.total_number_of_frames:
done = True
break
yield {"frame": frame, "frame_number": [frame_number]}
counter += 1
================================================
FILE: tensor2tensor/data_generators/video_utils.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Base classes and utilities for video datasets."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import functools
import os
from absl import flags
import numpy as np
import six
from tensor2tensor.data_generators import generator_utils
from tensor2tensor.data_generators import image_utils
from tensor2tensor.data_generators import problem
from tensor2tensor.data_generators import text_encoder
from tensor2tensor.layers import common_layers
from tensor2tensor.layers import common_video
from tensor2tensor.layers import modalities
from tensor2tensor.utils import contrib
from tensor2tensor.utils import metrics
from tensor2tensor.utils import video_metrics
import tensorflow.compat.v1 as tf
from tensorflow.compat.v1 import estimator as tf_estimator
FLAGS = flags.FLAGS
flags.DEFINE_bool(
"disable_ffmpeg", False, "Disable FFMPEG when generating debug videos."
)
def resize_video_frames(images, size):
return [tf.to_int64(tf.image.resize_images(
image, [size, size], tf.image.ResizeMethod.BILINEAR)) for image in images]
def video_augmentation(features, hue=False, saturate=False, contrast=False):
"""Augments video with optional hue, saturation and constrast.
Args:
features: dict, with keys "inputs", "targets".
features["inputs"], 4-D Tensor, shape=(THWC)
features["targets"], 4-D Tensor, shape=(THWC)
hue: bool, apply hue_transform.
saturate: bool, apply saturation transform.
contrast: bool, apply constrast transform.
Returns:
augment_features: dict with transformed "inputs" and "targets".
"""
inputs, targets = features["inputs"], features["targets"]
in_steps = common_layers.shape_list(inputs)[0]
# makes sure that the same augmentation is applied to both input and targets.
# if input is 4-D, then tf.image applies the same transform across the batch.
video = tf.concat((inputs, targets), axis=0)
if hue:
video = tf.image.random_hue(video, max_delta=0.2)
if saturate:
video = tf.image.random_saturation(video, lower=0.5, upper=1.5)
if contrast:
video = tf.image.random_contrast(video, lower=0.5, upper=1.5)
features["inputs"], features["targets"] = video[:in_steps], video[in_steps:]
return features
def create_border(video, color="blue", border_percent=2):
"""Creates a border around each frame to differentiate input and target.
Args:
video: 5-D NumPy array.
color: string, "blue", "red" or "green".
border_percent: Percentarge of the frame covered by the border.
Returns:
video: 5-D NumPy array.
"""
# Do not create border if the video is not in RGB format
if video.shape[-1] != 3:
return video
color_to_axis = {"blue": 2, "red": 0, "green": 1}
axis = color_to_axis[color]
_, _, height, width, _ = video.shape
border_height = np.ceil(border_percent * height / 100.0).astype(int)
border_width = np.ceil(border_percent * width / 100.0).astype(int)
video[:, :, :border_height, :, axis] = 255
video[:, :, -border_height:, :, axis] = 255
video[:, :, :, :border_width, axis] = 255
video[:, :, :, -border_width:, axis] = 255
return video
def convert_videos_to_summaries(input_videos, output_videos, target_videos,
tag, decode_hparams,
display_ground_truth=False):
"""Converts input, output and target videos into video summaries.
Args:
input_videos: 5-D NumPy array, (NTHWC) conditioning frames.
output_videos: 5-D NumPy array, (NTHWC) model predictions.
target_videos: 5-D NumPy array, (NTHWC) target frames.
tag: tf summary tag.
decode_hparams: HParams.
display_ground_truth: Whether or not to display ground truth videos.
Returns:
summaries: a list of tf frame-by-frame and video summaries.
"""
fps = decode_hparams.frames_per_second
border_percent = decode_hparams.border_percent
max_outputs = decode_hparams.max_display_outputs
target_steps = target_videos.shape[1]
all_summaries = []
input_videos = create_border(
input_videos, color="blue", border_percent=border_percent)
target_videos = create_border(
target_videos, color="red", border_percent=border_percent)
output_videos = create_border(
output_videos, color="red", border_percent=border_percent)
all_input = np.concatenate((input_videos, target_videos), axis=1)
all_output = np.concatenate((input_videos, output_videos), axis=1)
output_summ_vals, _ = common_video.py_gif_summary(
"%s/output" % tag, all_output, max_outputs=max_outputs, fps=fps,
return_summary_value=True)
all_summaries.extend(output_summ_vals)
# Optionally display ground truth.
if display_ground_truth:
input_summ_vals, _ = common_video.py_gif_summary(
"%s/input" % tag, all_input, max_outputs=max_outputs, fps=fps,
return_summary_value=True)
all_summaries.extend(input_summ_vals)
# Frame-by-frame summaries
iterable = zip(output_videos[:max_outputs, :target_steps],
target_videos[:max_outputs])
for ind, (input_video, output_video) in enumerate(iterable):
t, h, w, c = input_video.shape
# Tile vertically
input_frames = np.reshape(input_video, (t*h, w, c))
output_frames = np.reshape(output_video, (t*h, w, c))
# Concat across width.
all_frames = np.concatenate((input_frames, output_frames), axis=1)
tag = "input/output/%s_sample_%d" % (tag, ind)
frame_by_frame_summ = image_utils.image_to_tf_summary_value(
all_frames, tag=tag)
all_summaries.append(frame_by_frame_summ)
return all_summaries
def display_video_hooks(hook_args):
"""Hooks to display videos at decode time."""
predictions = hook_args.predictions
max_outputs = hook_args.decode_hparams.max_display_outputs
max_decodes = hook_args.decode_hparams.max_display_decodes
with tf.Graph().as_default():
_, best_decodes = video_metrics.compute_video_metrics_from_predictions(
predictions, decode_hparams=hook_args.decode_hparams)
all_summaries = []
# Displays decodes corresponding to the best/worst metric,
for metric, metric_decode_inds in best_decodes.items():
curr_metric_inds = metric_decode_inds[:max_outputs]
best_inputs, best_outputs, best_targets = [], [], []
for sample_ind, decode_ind in enumerate(curr_metric_inds):
curr_decode = predictions[decode_ind][sample_ind]
best_inputs.append(curr_decode["inputs"])
best_outputs.append(curr_decode["outputs"])
best_targets.append(curr_decode["targets"])
best_inputs = np.array(best_inputs, dtype=np.uint8)
best_outputs = np.array(best_outputs, dtype=np.uint8)
best_targets = np.array(best_targets, dtype=np.uint8)
summaries = convert_videos_to_summaries(
best_inputs, best_outputs, best_targets,
tag=metric, decode_hparams=hook_args.decode_hparams)
all_summaries.extend(summaries)
# Display random decodes for ten conditioning frames.
for decode_ind, decode in enumerate(predictions[: max_decodes]):
target_videos = video_metrics.stack_data_given_key(decode, "targets")
output_videos = video_metrics.stack_data_given_key(decode, "outputs")
input_videos = video_metrics.stack_data_given_key(decode, "inputs")
target_videos = np.asarray(target_videos, dtype=np.uint8)
output_videos = np.asarray(output_videos, dtype=np.uint8)
input_videos = np.asarray(input_videos, dtype=np.uint8)
summaries = convert_videos_to_summaries(
input_videos, output_videos, target_videos,
tag="decode_%d" % decode_ind, decode_hparams=hook_args.decode_hparams,
display_ground_truth=decode_ind == 0)
all_summaries.extend(summaries)
return all_summaries
def summarize_video_metrics(hook_args):
"""Computes video metrics summaries using the decoder output."""
problem_name = hook_args.problem.name
current_problem = hook_args.problem
hparams = hook_args.hparams
output_dirs = hook_args.output_dirs
predictions = hook_args.predictions
frame_shape = [
current_problem.frame_height, current_problem.frame_width,
current_problem.num_channels
]
metrics_graph = tf.Graph()
with metrics_graph.as_default():
if predictions:
metrics_results, _ = video_metrics.compute_video_metrics_from_predictions(
predictions, decode_hparams=hook_args.decode_hparams)
else:
metrics_results, _ = video_metrics.compute_video_metrics_from_png_files(
output_dirs, problem_name, hparams.video_num_target_frames,
frame_shape)
summary_values = []
for name, array in six.iteritems(metrics_results):
for ind, val in enumerate(array):
tag = "metric_{}/{}".format(name, ind)
summary_values.append(tf.Summary.Value(tag=tag, simple_value=val))
return summary_values
def debug_video_writer_factory(output_dir):
"""Creates a VideoWriter for debug videos."""
if FLAGS.disable_ffmpeg:
return common_video.IndividualFrameWriter(output_dir)
else:
output_path = os.path.join(output_dir, "video.avi")
return common_video.WholeVideoWriter(
fps=10, output_path=output_path, file_format="avi"
)
class VideoProblem(problem.Problem):
"""Base class for problems with videos."""
def __init__(self, *args, **kwargs):
super(VideoProblem, self).__init__(*args, **kwargs)
# Path to a directory to dump generated frames as png for debugging.
# If empty, no debug frames will be generated.
self.debug_dump_frames_path = ""
# Whether to skip random inputs at the beginning or not.
self.settable_random_skip = True
self.settable_use_not_breaking_batching = True
self.shuffle = True
def max_frames_per_video(self, hparams):
"""Maximum number of frames per video as determined by the dataset.
This is used only in PREDICT mode and handles the corner case where
video_num_input_frames + video_num_target_frames is greater than the
maximum number of frames per video in the dataset. For eg, 30 in BAIR.
For this special case, setting this to return "x" limits the input pipeline
to handle "x" (input + target) frames. The corresponding video model can
then decode arbitrary number of target frames via
hparams.video_num_target_frames.
Args:
hparams: HParams.
Returns:
num_frames: int.
"""
return hparams.video_num_input_frames + hparams.video_num_target_frames
@property
def num_channels(self):
"""Number of color channels in each frame."""
return 3
@property
def frame_height(self):
"""Height of each frame."""
raise NotImplementedError
@property
def frame_width(self):
"""Width of each frame."""
raise NotImplementedError
@property
def frame_shape(self):
"""Shape of a frame: a list [height , width , channels]."""
return [self.frame_height, self.frame_width, self.num_channels]
@property
def total_number_of_frames(self):
"""The total number of frames, needed for sharding."""
# It can also be a lower number -- we will switch shards every
# total_number_of_frames // num_shards time, so for example if
# you know that every video is 30 frames long and you have 100 shards
# then it's sufficient to set this to 30 * 100 so no shard-switching
# occurs during the generation of a video. For videos of variable length,
# just make this large so switching shards mid-video is very rare.
raise NotImplementedError
@property
def random_skip(self):
"""Whether to skip random inputs at the beginning or not."""
return True
@property
def extra_reading_spec(self):
"""Additional data fields to store on disk and their decoders."""
return {}, {}
@property
def dataset_splits(self):
"""Splits of data to produce and number of output shards for each."""
return [{
"split": problem.DatasetSplit.TRAIN,
"shards": 10,
}, {
"split": problem.DatasetSplit.EVAL,
"shards": 1,
}]
@property
def only_keep_videos_from_0th_frame(self):
return True
@property
def avoid_overlapping_frames(self):
"""When True, each video has non overlapping frames with every other."""
return False
@property
def use_not_breaking_batching(self):
return True
def preprocess_example(self, example, mode, hparams):
"""Runtime preprocessing, e.g., resize example["frame"]."""
if getattr(hparams, "preprocess_resize_frames", None) is not None:
example["frame"] = tf.image.resize_images(
example["frame"], hparams.preprocess_resize_frames,
tf.image.ResizeMethod.BILINEAR)
return example
@property
def decode_hooks(self):
return [summarize_video_metrics, display_video_hooks]
@property
def is_generate_per_split(self):
"""A single call to `generate_samples` generates for all `dataset_splits`.
Set to True if you already have distinct subsets of data for each dataset
split specified in `self.dataset_splits`. `self.generate_samples` will be
called once for each split.
Set to False if you have a unified dataset that you'd like to have split out
into training and evaluation data automatically. `self.generate_samples`
will be called only once and the data will be sharded across the dataset
splits specified in `self.dataset_splits`.
Returns:
bool
"""
raise NotImplementedError()
def example_reading_spec(self):
extra_data_fields, extra_data_items_to_decoders = self.extra_reading_spec
data_fields = {
"image/encoded": tf.FixedLenFeature((), tf.string),
"image/format": tf.FixedLenFeature((), tf.string),
}
data_fields.update(extra_data_fields)
data_items_to_decoders = {
"frame":
contrib.slim().tfexample_decoder.Image(
image_key="image/encoded",
format_key="image/format",
shape=[self.frame_height, self.frame_width, self.num_channels],
channels=self.num_channels),
}
data_items_to_decoders.update(extra_data_items_to_decoders)
return data_fields, data_items_to_decoders
def serving_input_fn(self, hparams):
"""For serving/predict, assume that only video frames are provided."""
video_input_frames = tf.placeholder(
dtype=tf.float32,
shape=[
None, hparams.video_num_input_frames, self.frame_width,
self.frame_height, self.num_channels
])
# TODO(michalski): add support for passing input_action and input_reward.
return tf_estimator.export.ServingInputReceiver(
features={"inputs": video_input_frames},
receiver_tensors=video_input_frames)
def preprocess(self, dataset, mode, hparams, interleave=True):
def split_on_batch(x):
"""Split x on batch dimension into x[:size, ...] and x[size:, ...]."""
length = len(x.get_shape())
size = hparams.video_num_input_frames
if length < 1:
raise ValueError("Batched tensor of length < 1.")
if length == 1:
return x[:size], x[size:]
if length == 2:
return x[:size, :], x[size:, :]
if length == 3:
return x[:size, :, :], x[size:, :, :]
if length == 4:
return x[:size, :, :, :], x[size:, :, :, :]
# TODO(lukaszkaiser): use tf.split for the general case.
raise ValueError("Batch splitting on general dimensions not done yet.")
def features_from_batch(batched_prefeatures):
"""Construct final features from the batched inputs.
This function gets prefeatures.
Args:
batched_prefeatures: single-frame features (from disk) as batch tensors.
Returns:
Features dictionary with joint features per-frame.
"""
features = {}
for k, v in six.iteritems(batched_prefeatures):
if k == "frame": # We rename past frames to inputs and targets.
s1, s2 = split_on_batch(v)
features["inputs"] = s1
features["targets"] = s2
else:
s1, s2 = split_on_batch(v)
features["input_%s" % k] = s1
features["target_%s" % k] = s2
return features
# Batch and construct features.
def _preprocess(example):
return self.preprocess_example(example, mode, hparams)
def avoid_break_batching(dataset):
"""Smart preprocessing to avoid break between videos!
Simple batching of images into videos may result into broken videos
with two parts from two different videos. This preprocessing avoids
this using the frame number.
Args:
dataset: raw not-batched dataset.
Returns:
batched not-broken videos.
"""
def check_integrity_and_batch(*datasets):
"""Checks whether a sequence of frames are from the same video.
Args:
*datasets: datasets each skipping 1 frame from the previous one.
Returns:
batched data and the integrity flag.
"""
not_broken = tf.constant(True)
if "frame_number" in datasets[0]:
frame_numbers = [dataset["frame_number"][0] for dataset in datasets]
not_broken = tf.equal(frame_numbers[-1] - frame_numbers[0],
num_frames - 1)
if self.only_keep_videos_from_0th_frame:
not_broken = tf.logical_and(not_broken, tf.equal(
frame_numbers[0], 0))
if self.avoid_overlapping_frames:
non_overlap = tf.equal(tf.mod(frame_numbers[0], num_frames), 0)
not_broken = tf.logical_and(not_broken, non_overlap)
else:
tf.logging.warning("use_not_breaking_batching is True but "
"no frame_number is in the dataset.")
features = {}
for key in datasets[0].keys():
values = [dataset[key] for dataset in datasets]
batch = tf.stack(values)
features[key] = batch
return features, not_broken
ds = [dataset.skip(i) for i in range(num_frames)]
dataset = tf.data.Dataset.zip(tuple(ds))
dataset = dataset.map(check_integrity_and_batch)
dataset = dataset.filter(lambda _, not_broken: not_broken)
dataset = dataset.map(lambda features, _: features)
return dataset
preprocessed_dataset = dataset.map(_preprocess)
num_frames = (
hparams.video_num_input_frames + hparams.video_num_target_frames)
if mode == tf_estimator.ModeKeys.PREDICT:
num_frames = min(self.max_frames_per_video(hparams), num_frames)
# We jump by a random position at the beginning to add variety.
if (self.random_skip and self.settable_random_skip and interleave and
mode == tf_estimator.ModeKeys.TRAIN):
random_skip = tf.random_uniform([], maxval=num_frames, dtype=tf.int64)
preprocessed_dataset = preprocessed_dataset.skip(random_skip)
if (self.use_not_breaking_batching and
self.settable_use_not_breaking_batching):
batch_dataset = avoid_break_batching(preprocessed_dataset)
else:
batch_dataset = preprocessed_dataset.batch(num_frames,
drop_remainder=True)
dataset = batch_dataset.map(features_from_batch)
if self.shuffle and interleave and mode == tf_estimator.ModeKeys.TRAIN:
dataset = dataset.shuffle(hparams.get("shuffle_buffer_size", 128))
return dataset
def eval_metrics(self):
eval_metrics = [
metrics.Metrics.ACC, metrics.Metrics.ACC_PER_SEQ,
metrics.Metrics.NEG_LOG_PERPLEXITY, metrics.Metrics.IMAGE_SUMMARY
]
return eval_metrics
def validate_frame(self, frame):
height, width, channels = frame.shape
if channels != self.num_channels:
raise ValueError("Generated frame has %d channels while the class "
"assumes %d channels." % (channels, self.num_channels))
if height != self.frame_height:
raise ValueError("Generated frame has height %d while the class "
"assumes height %d." % (height, self.frame_height))
if width != self.frame_width:
raise ValueError("Generated frame has width %d while the class "
"assumes width %d." % (width, self.frame_width))
def generate_samples(self, data_dir, tmp_dir, dataset_split):
"""Generate samples of the frames with possible extra data.
Args:
data_dir: final data directory. Typically only used in this method to copy
over user-supplied vocab files if there are extra fields needing them.
tmp_dir: temporary directory that you can use for downloading and scratch.
dataset_split: problem.DatasetSplit, which data split to generate samples
for (for example, training and evaluation). You can assume it's TRAIN if
self.
Yields:
Sample: dict; we assume that there is
a "frame" feature with unencoded frame which is a numpy arrays of shape
[frame_height, frame_width, num_channels] and which will be transcoded
into an image format by generate_encodeded_samples.
"""
raise NotImplementedError()
def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
"""Generate samples of the encoded frames with possible extra data.
By default this function just encodes the numpy array returned as "frame"
from `self.generate_samples` into a PNG image. Override this function to
get other encodings on disk.
Args:
data_dir: final data directory. Typically only used in this method to copy
over user-supplied vocab files if there are extra fields needing them.
tmp_dir: temporary directory that you can use for downloading and scratch.
dataset_split: problem.DatasetSplit, which data split to generate samples
for (for example, training and evaluation).
Yields:
Sample: dict which is in disk encoding.
Raises:
ValueError: if the frame has a different number of channels than required.
"""
writer = None
with tf.Graph().as_default():
image_t = tf.placeholder(dtype=tf.uint8, shape=(None, None, None))
encoded_image_t = tf.image.encode_png(image_t)
with tf.Session() as sess:
for features in self.generate_samples(data_dir, tmp_dir, dataset_split):
unencoded_frame = features.pop("frame")
self.validate_frame(unencoded_frame)
height, width, _ = unencoded_frame.shape
encoded_frame = sess.run(
encoded_image_t, feed_dict={image_t: unencoded_frame})
features["image/encoded"] = [encoded_frame]
features["image/format"] = ["png"]
features["image/height"] = [height]
features["image/width"] = [width]
has_debug_image = "image/debug" in features
if has_debug_image:
unencoded_debug = features.pop("image/debug")
encoded_debug = sess.run(
encoded_image_t, feed_dict={image_t: unencoded_debug})
features["image/encoded_debug"] = [encoded_debug]
if self.debug_dump_frames_path:
# Defer creating debug writer until we know debug_dump_frames_path.
if writer is None:
if not tf.gfile.Exists(self.debug_dump_frames_path):
tf.gfile.MkDir(self.debug_dump_frames_path)
writer = debug_video_writer_factory(self.debug_dump_frames_path)
img = unencoded_debug if has_debug_image else unencoded_frame
encoded_img = encoded_debug if has_debug_image else encoded_frame
writer.write(img, encoded_img)
yield features
if self.debug_dump_frames_path:
writer.finish_to_disk()
def generate_data(self, data_dir, tmp_dir, task_id=-1):
"""The function generating the data."""
filepath_fns = {
problem.DatasetSplit.TRAIN: self.training_filepaths,
problem.DatasetSplit.EVAL: self.dev_filepaths,
problem.DatasetSplit.TEST: self.test_filepaths,
}
# We set shuffled=True as we don't want to shuffle on disk later.
split_paths = [(split["split"], filepath_fns[split["split"]](
data_dir, split["shards"], shuffled=True))
for split in self.dataset_splits]
all_paths = []
for _, paths in split_paths:
all_paths.extend(paths)
if self.is_generate_per_split:
for split, paths in split_paths:
generator_utils.generate_files(
self.generate_encoded_samples(data_dir, tmp_dir, split),
paths,
cycle_every_n=self.total_number_of_frames // len(paths))
else:
generator_utils.generate_files(
self.generate_encoded_samples(data_dir, tmp_dir,
problem.DatasetSplit.TRAIN),
all_paths,
cycle_every_n=self.total_number_of_frames // len(all_paths))
# TODO(lukaszkaiser): remove this version after everything is ported.
class VideoProblemOld(problem.Problem):
"""Base class for problems with videos: previous version."""
@property
def num_channels(self):
"""Number of color channels."""
return 3
def example_reading_spec(self):
data_fields = {
"image/encoded": tf.FixedLenFeature((), tf.string),
"image/format": tf.FixedLenFeature((), tf.string),
}
data_items_to_decoders = {
"inputs":
contrib.slim().tfexample_decoder.Image(
image_key="image/encoded",
format_key="image/format",
channels=self.num_channels),
}
return data_fields, data_items_to_decoders
def eval_metrics(self):
eval_metrics = [
metrics.Metrics.ACC, metrics.Metrics.ACC_TOP5,
metrics.Metrics.NEG_LOG_PERPLEXITY
]
return eval_metrics
class VideoAugmentationProblem(VideoProblem):
"""Base class for video data-augmentation.
By default applies a random hue, contrast and saturation transformation
to every video. To disable any of these transformations, inherit
this class and set the corresponding property to False.
"""
@property
def hue(self):
return True
@property
def contrast(self):
return True
@property
def saturate(self):
return True
def preprocess(self, dataset, mode, hparams, interleave=True):
dataset = super(VideoAugmentationProblem, self).preprocess(
dataset=dataset, mode=mode, hparams=hparams, interleave=interleave)
video_augment_func = functools.partial(
video_augmentation, hue=self.hue, contrast=self.contrast,
saturate=self.saturate)
if mode == tf_estimator.ModeKeys.TRAIN:
dataset = dataset.map(video_augment_func)
return dataset
class Video2ClassProblem(VideoProblemOld):
"""Base class for image classification problems."""
@property
def is_small(self):
raise NotImplementedError()
@property
def num_classes(self):
raise NotImplementedError()
@property
def train_shards(self):
raise NotImplementedError()
@property
def dev_shards(self):
return 1
@property
def class_labels(self):
return ["ID_%d" % i for i in range(self.num_classes)]
@property
def image_size(self):
raise NotImplementedError()
def feature_encoders(self, data_dir):
del data_dir
return {
"inputs": text_encoder.ImageEncoder(),
"targets": text_encoder.ClassLabelEncoder(self.class_labels)
}
def generator(self, data_dir, tmp_dir, is_training):
raise NotImplementedError()
def example_reading_spec(self):
label_key = "image/class/label"
data_fields, data_items_to_decoders = (
super(Video2ClassProblem, self).example_reading_spec())
data_fields[label_key] = tf.FixedLenFeature((1,), tf.int64)
data_items_to_decoders["targets"] = contrib.slim().tfexample_decoder.Tensor(
label_key)
return data_fields, data_items_to_decoders
def hparams(self, defaults, unused_model_hparams):
p = defaults
p.modality = {"inputs": modalities.ModalityType.IMAGE,
"targets": modalities.ModalityType.CLASS_LABEL}
p.vocab_size = {"inputs": 256,
"targets": self.num_classes}
p.input_space_id = problem.SpaceID.IMAGE
p.target_space_id = problem.SpaceID.IMAGE_LABEL
def generate_data(self, data_dir, tmp_dir, task_id=-1):
generator_utils.generate_dataset_and_shuffle(
self.generator(data_dir, tmp_dir, True),
self.training_filepaths(data_dir, self.train_shards, shuffled=False),
self.generator(data_dir, tmp_dir, False),
self.dev_filepaths(data_dir, self.dev_shards, shuffled=False))
================================================
FILE: tensor2tensor/data_generators/video_utils_test.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""video_utils test."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from absl.testing import parameterized
import numpy as np
from tensor2tensor.data_generators import video_generated # pylint: disable=unused-import
from tensor2tensor.data_generators import video_utils
from tensor2tensor.utils import decoding
from tensor2tensor.utils import registry
import tensorflow.compat.v1 as tf
class VideoUtilsTest(parameterized.TestCase, tf.test.TestCase):
def get_predictions(self, num_decodes=2):
rng = np.random.RandomState(0)
# num_samples=4
inputs = rng.randint(0, 255, (4, 2, 64, 64, 3))
outputs = rng.randint(0, 255, (4, 5, 64, 64, 3))
targets = rng.randint(0, 255, (4, 5, 64, 64, 3))
predictions = []
for input_, output, target in zip(inputs, outputs, targets):
curr_pred = {"inputs": input_, "outputs": output, "targets": target}
predictions.append(curr_pred)
# num_decodes=2
predictions = [predictions] * num_decodes
problem = registry.problem("video_stochastic_shapes10k")
return predictions, problem
def testVideoAugmentation(self):
# smoke-test, test for shapes.
with tf.Graph().as_default():
inputs = tf.random_uniform(shape=(3, 64, 64, 3))
targets = tf.random_uniform(shape=(10, 64, 64, 3))
features = {"inputs": inputs, "targets": targets}
augment = video_utils.video_augmentation(
features, hue=True, saturate=True, contrast=True)
with tf.Session() as sess:
augment_dict = sess.run(augment)
self.assertEqual(augment_dict["inputs"].shape, (3, 64, 64, 3))
self.assertEqual(augment_dict["targets"].shape, (10, 64, 64, 3))
def testDecodeInMemoryTrue(self):
predictions, problem = self.get_predictions()
decode_hparams = decoding.decode_hparams()
decode_hparams.decode_in_memory = True
decode_hooks = decoding.DecodeHookArgs(
estimator=None, problem=problem, output_dirs=None,
hparams=decode_hparams, decode_hparams=decode_hparams,
predictions=predictions)
metrics = video_utils.summarize_video_metrics(decode_hooks)
@parameterized.named_parameters(
("d5_o6", 5, 6))
# ("d5", 5), ("d10", 10), ("d5_o6", 5, 6))
def testConvertPredictionsToVideoSummaries(self, num_decodes=5,
max_output_steps=5):
# Initialize predictions.
rng = np.random.RandomState(0)
inputs = rng.randint(0, 255, (2, 32, 32, 3))
outputs = rng.randint(0, 255, (max_output_steps, 32, 32, 3))
targets = rng.randint(0, 255, (5, 32, 32, 3))
# batch it up.
prediction = [{"outputs": outputs, "inputs": inputs, "targets": targets}]*5
predictions = [prediction] * num_decodes
decode_hparams = decoding.decode_hparams(
overrides="max_display_decodes=5")
decode_hooks = decoding.DecodeHookArgs(
estimator=None, problem=None, output_dirs=None,
hparams=decode_hparams, decode_hparams=decode_hparams,
predictions=predictions)
summaries = video_utils.display_video_hooks(decode_hooks)
for summary in summaries:
self.assertIsInstance(summary, tf.Summary.Value)
if __name__ == "__main__":
tf.test.main()
================================================
FILE: tensor2tensor/data_generators/vqa.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Data generators for VQA data sets."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import base64
import csv
import json
import os
import random
import sys
import tarfile
import zipfile
import numpy as np
from tensor2tensor.data_generators import generator_utils
from tensor2tensor.data_generators import image_utils
from tensor2tensor.data_generators import problem
from tensor2tensor.data_generators import text_encoder
from tensor2tensor.data_generators import vqa_utils
from tensor2tensor.layers import modalities
from tensor2tensor.utils import contrib
from tensor2tensor.utils import metrics
from tensor2tensor.utils import registry
import tensorflow.compat.v1 as tf
def _get_vqa_v2_annotations(directory,
annotation_url,
annotation_filename="vqa_v2.tar.gz"):
"""Extract the VQA V2 annotation files to directory unless it's there."""
annotation_file = generator_utils.maybe_download_from_drive(
directory, annotation_filename, annotation_url)
with tarfile.open(annotation_file, "r:gz") as annotation_tar:
annotation_tar.extractall(directory)
def _get_vqa_v2_image_raw_dataset(directory, image_root_url, image_urls):
"""Extract the VQA V2 image data set to directory unless it's there."""
for url in image_urls:
filename = os.path.basename(url)
download_url = os.path.join(image_root_url, url)
path = generator_utils.maybe_download(directory, filename, download_url)
unzip_dir = os.path.join(directory, filename.strip(".zip"))
if not tf.gfile.Exists(unzip_dir):
zipfile.ZipFile(path, "r").extractall(directory)
def _get_vqa_v2_image_feature_dataset(
directory, feature_url, feature_filename="mscoco_feat.tar.gz"):
"""Extract the VQA V2 feature data set to directory unless it's there."""
feature_file = generator_utils.maybe_download_from_drive(
directory, feature_filename, feature_url)
with tarfile.open(feature_file, "r:gz") as feature_tar:
feature_tar.extractall(directory)
class ImageQuestion2MultilabelProblem(image_utils.ImageProblem):
"""Base class for image question answer problem."""
@property
def target_space_id(self):
raise NotImplementedError()
@property
def vocab_size(self):
raise NotImplementedError
@property
def num_classes(self):
raise NotImplementedError()
@property
def vocab_filename(self):
raise NotImplementedError()
@property
def label_filename(self):
raise NotImplementedError()
@property
def train_shards(self):
raise NotImplementedError()
@property
def dev_shards(self):
raise NotImplementedError()
def source_data_files(self, dataset_split):
raise NotImplementedError()
def generator(self, data_dir, tmp_dir, dataset_split):
raise NotImplementedError()
def eval_metrics(self):
return [
metrics.Metrics.ACC_MULTILABEL_MATCH3,
]
def feature_encoders(self, data_dir):
input_encoder = text_encoder.ImageEncoder(channels=self.num_channels)
vocab_file = os.path.join(data_dir, self.vocab_filename)
question_encoder = text_encoder.TokenTextEncoder(
vocab_file, replace_oov="UNK")
label_file = os.path.join(data_dir, self.label_filename)
target_encoder = text_encoder.ClassLabelEncoder(
class_labels_fname=label_file)
return {"inputs": input_encoder,
"question": question_encoder,
"targets": target_encoder}
def hparams(self, defaults, unused_model_hparams):
p = defaults
question_encoder = self._encoders["question"]
targets_encoder = self._encoders["targets"]
p.modality = {
"inputs": modalities.ModalityType.IDENTITY,
"question": modalities.ModalityType.SYMBOL,
"targets": modalities.ModalityType.MULTI_LABEL,
}
p.vocab_size = {
"inputs": None,
"question": question_encoder.vocab_size,
"targets": targets_encoder.vocab_size,
}
p.input_space_id = problem.SpaceID.IMAGE # multiple input features?
p.target_space_id = self.target_space_id
def generate_data(self, data_dir, tmp_dir, task_id=-1):
generator_utils.generate_dataset_and_shuffle(
self.generator(data_dir, tmp_dir, problem.DatasetSplit.TRAIN),
self.training_filepaths(data_dir, self.train_shards, shuffled=False),
self.generator(data_dir, tmp_dir, problem.DatasetSplit.EVAL),
self.dev_filepaths(data_dir, self.dev_shards, shuffled=False))
@registry.register_problem
class ImageVqav2Tokens10kLabels3k(ImageQuestion2MultilabelProblem):
"""VQA V2, raw images, 10k question vocab, 3k answer label."""
_MSCOCO_ROOT_URL = "http://msvocds.blob.core.windows.net/"
_MSCOCO_IMAGE_URLS = [
"coco2014/train2014.zip", "coco2014/val2014.zip", "coco2014/test2014.zip",
]
_VQA_V2_ANNOTATION_URL = ("https://drive.google.com/uc?export=download&id="
"1xfMU54ObCLvMRAekT3cfcIg-AgY39fWB")
_VQA_V2_TRAIN_DATASETS = [
("trainval_resnet101_faster_rcnn_genome_36.tsv",
"v2_train2014_annotations.json"),
]
_VQA_V2_DEV_DATASETS = [
("trainval_resnet101_faster_rcnn_genome_36.tsv",
"v2_val2014_annotations.json"),
]
_VQA_V2_TEST_DATASETS = [
("test2015_resnet101_faster_rcnn_genome_36.tsv",
"v2_test2015_annotations.json"),
]
def source_data_files(self, dataset_split):
train = dataset_split == problem.DatasetSplit.TRAIN
return self._VQA_V2_TRAIN_DATASETS if train else self._VQA_V2_DEV_DATASETS
@property
def target_space_id(self):
return problem.SpaceID.GENERIC
@property
def vocab_size(self):
return 10000
@property
def num_classes(self):
return 3000
@property
def vocab_filename(self):
return "question.vocab.%d" % self.vocab_size
@property
def label_filename(self):
return "answer.label.%d" % self.num_classes
@property
def train_shards(self):
return 128
@property
def dev_shards(self):
return 64
def example_reading_spec(self):
data_fields, data_items_to_decoders = (
super(ImageVqav2Tokens10kLabels3k, self).example_reading_spec())
data_fields["image/image_id"] = tf.FixedLenFeature((), tf.int64)
data_fields["image/question_id"] = tf.FixedLenFeature((), tf.int64)
data_fields["image/question"] = tf.FixedLenSequenceFeature(
(), tf.int64, allow_missing=True)
data_fields["image/answer"] = tf.FixedLenSequenceFeature(
(), tf.int64, allow_missing=True)
slim = contrib.slim()
data_items_to_decoders["question"] = slim.tfexample_decoder.Tensor(
"image/question")
data_items_to_decoders["targets"] = slim.tfexample_decoder.Tensor(
"image/answer")
return data_fields, data_items_to_decoders
def preprocess_example(self, example, mode, hparams):
# hparams is model_hparams
image = example["inputs"]
example["inputs"] = vqa_utils.vqa_v2_preprocess_image(
image, hparams.height, hparams.width, mode,
resize_side=hparams.resize_side, distort=hparams.distort,
image_model_fn=hparams.image_model_fn)
return example
def generator(self, data_dir, tmp_dir, dataset_split):
datasets = self.source_data_files(dataset_split)
return self.vqa_v2_generator(data_dir, tmp_dir, datasets)
def vqa_v2_generator(self, data_dir, tmp_dir, datasets):
"""VQA v2 generator using raw images."""
_get_vqa_v2_annotations(tmp_dir, self._VQA_V2_ANNOTATION_URL)
_get_vqa_v2_image_raw_dataset(tmp_dir, self._MSCOCO_ROOT_URL,
self._MSCOCO_IMAGE_URLS)
vocab_path = os.path.join(data_dir, self.vocab_filename)
if not tf.gfile.Exists(vocab_path):
vocab_tmp_path = os.path.join(tmp_dir, self.vocab_filename)
tf.gfile.Copy(vocab_tmp_path, vocab_path)
with tf.gfile.GFile(vocab_path, mode="r") as f:
vocab_data = "\n\n" + f.read() + "UNK\n"
with tf.gfile.GFile(vocab_path, mode="w") as f:
f.write(vocab_data)
label_path = os.path.join(data_dir, self.label_filename)
if not tf.gfile.Exists(label_path):
label_tmp_path = os.path.join(tmp_dir, self.label_filename)
tf.gfile.Copy(label_tmp_path, label_path)
vocab_encoder = text_encoder.TokenTextEncoder(vocab_path, replace_oov="UNK")
label_encoder = text_encoder.ClassLabelEncoder(
class_labels_fname=label_path)
prefix_annotation = []
for prefix, annotation_file in datasets:
annotation_path = os.path.join(tmp_dir, annotation_file)
with tf.gfile.Open(annotation_path) as f:
annotation_json = json.loads(f.read())
prefix_annotation += [(prefix, anno) for anno in annotation_json]
random.shuffle(prefix_annotation)
annotation_count = len(prefix_annotation)
tf.logging.info("Processing %d annotations for vqa v2" %(annotation_count))
for prefix, anno in prefix_annotation:
image_id = anno["image_id"]
question = vocab_encoder.encode(anno["question"])
answer = [label_encoder.encode(ans) for ans in anno["answer"]]
answer = answer if answer else [0] # 0 indicates padding
image_filename = "COCO_" + prefix + "_" + str(image_id).zfill(12) + ".jpg"
image_filepath = os.path.join(tmp_dir, prefix, image_filename)
with tf.gfile.Open(image_filepath, "r") as f:
encoded_image_data = f.read()
yield {
"image/encoded": [encoded_image_data],
"image/format": ["jpeg"],
"image/image_id": [image_id],
"image/question_id": [anno["question_id"]],
"image/question": question,
"image/answer": answer,
}
@registry.register_problem
class ImageVqav2RcnnFeatureTokens10kLabels3k(ImageVqav2Tokens10kLabels3k):
"""VQA V2, image feature, 10k question vocab, 3k answer label."""
_VQA_V2_FEATURE_URL = ("https://drive.google.com/uc?export=download&id="
"1yTTFUWqx1SScC-Whs2vRbF3tDsEEjrtt")
@property
def num_boxes(self):
return 36
@property
def feature_dimension(self):
return 2048
@property
def spatial_feature_dimension(self):
return 6
@property
def feature_file_field_names(self):
return ["image_id",
"image_w",
"image_h",
"num_boxes",
"boxes",
"features"]
def preprocess_example(self, example, mode, hparams):
# reshape some features
example["inputs"] = tf.reshape(
example["inputs"], [self.num_boxes, 1, self.feature_dimension])
example["spatial_feature"] = tf.reshape(
example["spatial_feature"],
[self.num_boxes, 1, self.spatial_feature_dimension])
return example
def example_reading_spec(self):
slim = contrib.slim()
data_fields, data_items_to_decoders = {}, {}
data_fields["image/feature"] = tf.FixedLenSequenceFeature(
(), tf.float32, allow_missing=True)
data_fields["image/spatial_feature"] = tf.FixedLenSequenceFeature(
(), tf.float32, allow_missing=True)
data_fields["image/image_id"] = tf.FixedLenFeature((), tf.int64)
data_fields["image/question_id"] = tf.FixedLenFeature((), tf.int64)
data_fields["image/question"] = tf.FixedLenSequenceFeature(
(), tf.int64, allow_missing=True)
data_fields["image/answer"] = tf.FixedLenSequenceFeature(
(), tf.int64, allow_missing=True)
data_items_to_decoders["inputs"] = slim.tfexample_decoder.Tensor(
"image/feature")
data_items_to_decoders["question_id"] = slim.tfexample_decoder.Tensor(
"image/question_id")
data_items_to_decoders["image_id"] = slim.tfexample_decoder.Tensor(
"image/image_id")
data_items_to_decoders["spatial_feature"] = slim.tfexample_decoder.Tensor(
"image/spatial_feature")
data_items_to_decoders["question"] = slim.tfexample_decoder.Tensor(
"image/question")
data_items_to_decoders["targets"] = slim.tfexample_decoder.Tensor(
"image/answer")
return data_fields, data_items_to_decoders
def vqa_v2_generator(self, data_dir, tmp_dir, datasets):
"""VQA v2 generator using image features."""
_get_vqa_v2_annotations(tmp_dir, self._VQA_V2_ANNOTATION_URL)
_get_vqa_v2_image_feature_dataset(tmp_dir, self._VQA_V2_FEATURE_URL)
vocab_path = os.path.join(data_dir, self.vocab_filename)
if not tf.gfile.Exists(vocab_path):
vocab_tmp_path = os.path.join(tmp_dir, self.vocab_filename)
tf.gfile.Copy(vocab_tmp_path, vocab_path)
with tf.gfile.GFile(vocab_path, mode="r") as f:
vocab_data = "\n\n" + f.read() + "UNK\n"
with tf.gfile.GFile(vocab_path, mode="w") as f:
f.write(vocab_data)
label_path = os.path.join(data_dir, self.label_filename)
if not tf.gfile.Exists(label_path):
label_tmp_path = os.path.join(tmp_dir, self.label_filename)
tf.gfile.Copy(label_tmp_path, label_path)
vocab_encoder = text_encoder.TokenTextEncoder(vocab_path, replace_oov="UNK")
label_encoder = text_encoder.ClassLabelEncoder(
class_labels_fname=label_path)
# merge annotations
annotation_json = []
for _, annotation_file in datasets:
annotation_path = os.path.join(tmp_dir, annotation_file)
with tf.gfile.Open(annotation_path) as f:
annotation_json += json.loads(f.read())
annotation_count = len(annotation_json)
tf.logging.info("Processing %d annotations for vqa v2" %(annotation_count))
imageid2annotation = {}
for anno in annotation_json:
if anno["image_id"] not in imageid2annotation:
imageid2annotation[anno["image_id"]] = [anno]
else:
imageid2annotation[anno["image_id"]].append(anno)
csv.field_size_limit(sys.maxsize)
for feature_file, _ in datasets:
feature_file_path = os.path.join(tmp_dir, feature_file)
with open(feature_file_path, "r+b") as tsv_file:
csv_reader = csv.DictReader(
tsv_file, delimiter="\t", fieldnames=self.feature_file_field_names)
for item in csv_reader:
item["num_boxes"] = int(item["num_boxes"])
image_id = int(item["image_id"])
image_w = float(item["image_w"])
image_h = float(item["image_h"])
bboxes = np.frombuffer(base64.decodestring(item["boxes"]),
dtype=np.float32).reshape(
(item["num_boxes"], -1))
box_width = bboxes[:, 2] - bboxes[:, 0]
box_height = bboxes[:, 3] - bboxes[:, 1]
scaled_width = box_width / image_w
scaled_height = box_height / image_h
scaled_x = bboxes[:, 0] / image_w
scaled_y = bboxes[:, 1] / image_h
box_width = box_width[..., np.newaxis]
box_height = box_height[..., np.newaxis]
scaled_width = scaled_width[..., np.newaxis]
scaled_height = scaled_height[..., np.newaxis]
scaled_x = scaled_x[..., np.newaxis]
scaled_y = scaled_y[..., np.newaxis]
spatial_features = np.concatenate(
(scaled_x,
scaled_y,
scaled_x + scaled_width,
scaled_y + scaled_height,
scaled_width,
scaled_height),
axis=1)
if image_id in imageid2annotation:
for anno in imageid2annotation[image_id]:
question = vocab_encoder.encode(anno["question"])
answer = [label_encoder.encode(ans) for ans in anno["answer"]]
answer = answer if answer else [0] # 0 indicates padding
yield {
"image/feature":
np.frombuffer(base64.decodestring(item["features"]),
dtype=np.float32).tolist(),
"image/spatial_feature": spatial_features.flatten().tolist(),
"image/height": [image_h],
"image/width": [image_w],
"image/bboxes": bboxes.flatten().tolist(),
"image/image_id": [image_id],
"image/question_id": [anno["question_id"]],
"image/question": question,
"image/answer": answer,
}
del imageid2annotation[image_id]
# assert all annotations are included
assert not imageid2annotation
================================================
FILE: tensor2tensor/data_generators/vqa_utils.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Utilities for VQA data sets."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import tensorflow.compat.v1 as tf
from tensorflow.compat.v1 import estimator as tf_estimator
from tensorflow.python.ops import control_flow_ops
# some functions are copied and modified from
# vgg_preprocessing and inception_preprocessing in
# models/research/slim/preprocessing/
_R_MEAN = 123.68
_G_MEAN = 116.78
_B_MEAN = 103.94
def _smallest_size_at_least(height, width, smallest_side):
"""Computes new shape with the smallest side equal to `smallest_side`.
Computes new shape with the smallest side equal to `smallest_side` while
preserving the original aspect ratio.
Args:
height: an int32 scalar tensor indicating the current height.
width: an int32 scalar tensor indicating the current width.
smallest_side: A python integer or scalar `Tensor` indicating the size of
the smallest side after resize.
Returns:
new_height: an int32 scalar tensor indicating the new height.
new_width: and int32 scalar tensor indicating the new width.
"""
smallest_side = tf.convert_to_tensor(smallest_side, dtype=tf.int32)
height = tf.to_float(height)
width = tf.to_float(width)
smallest_side = tf.to_float(smallest_side)
scale = tf.cond(
tf.greater(height, width), lambda: smallest_side / width,
lambda: smallest_side / height)
new_height = tf.to_int32(height * scale)
new_width = tf.to_int32(width * scale)
return new_height, new_width
def _aspect_preserving_resize(image, smallest_side):
"""Resize images preserving the original aspect ratio.
Args:
image: A 3-D image `Tensor`.
smallest_side: A python integer or scalar `Tensor` indicating the size of
the smallest side after resize.
Returns:
resized_image: A 3-D tensor containing the resized image.
"""
smallest_side = tf.convert_to_tensor(smallest_side, dtype=tf.int32)
shape = tf.shape(image)
height = shape[0]
width = shape[1]
new_height, new_width = _smallest_size_at_least(height, width, smallest_side)
image = tf.expand_dims(image, 0)
resized_image = tf.image.resize_images(
image, size=[new_height, new_width], method=tf.image.ResizeMethod.BICUBIC)
resized_image = tf.squeeze(resized_image)
resized_image.set_shape([None, None, 3])
return resized_image
def _flip(image):
"""Random horizontal image flip."""
image = tf.image.random_flip_left_right(image)
return image
def _distort_color(image, color_ordering=0, scope=None):
"""Distort the color of a Tensor image.
Each color distortion is non-commutative and thus ordering of the color ops
matters. Ideally we would randomly permute the ordering of the color ops.
Rather then adding that level of complication, we select a distinct ordering
of color ops for each preprocessing thread.
Args:
image: 3-D Tensor containing single image in [0, 1].
color_ordering: Python int, a type of distortion (valid values: 0-3).
scope: Optional scope for name_scope.
Returns:
3-D Tensor color-distorted image on range [0, 1]
Raises:
ValueError: if color_ordering not in [0, 3]
"""
with tf.name_scope(scope, "distort_color", [image]):
if color_ordering == 0:
image = tf.image.random_brightness(image, max_delta=32. / 255.)
image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
image = tf.image.random_hue(image, max_delta=0.2)
image = tf.image.random_contrast(image, lower=0.5, upper=1.5)
elif color_ordering == 1:
image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
image = tf.image.random_brightness(image, max_delta=32. / 255.)
image = tf.image.random_contrast(image, lower=0.5, upper=1.5)
image = tf.image.random_hue(image, max_delta=0.2)
elif color_ordering == 2:
image = tf.image.random_contrast(image, lower=0.5, upper=1.5)
image = tf.image.random_hue(image, max_delta=0.2)
image = tf.image.random_brightness(image, max_delta=32. / 255.)
image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
elif color_ordering == 3:
image = tf.image.random_hue(image, max_delta=0.2)
image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
image = tf.image.random_contrast(image, lower=0.5, upper=1.5)
image = tf.image.random_brightness(image, max_delta=32. / 255.)
else:
raise ValueError("color_ordering must be in [0, 3]")
# The random_* ops do not necessarily clamp.
return tf.clip_by_value(image, 0.0, 1.0)
def _apply_with_random_selector(x, func, num_cases):
"""Computes func(x, sel), with sel sampled from [0...num_cases-1].
Args:
x: input Tensor.
func: Python function to apply.
num_cases: Python int32, number of cases to sample sel from.
Returns:
The result of func(x, sel), where func receives the value of the
selector as a python integer, but sel is sampled dynamically.
"""
sel = tf.random_uniform([], maxval=num_cases, dtype=tf.int32)
# Pass the real x only to one of the func calls.
return control_flow_ops.merge([
func(control_flow_ops.switch(x, tf.equal(sel, case))[1], case)
for case in range(num_cases)
])[0]
def _mean_image_subtraction(image, means):
"""Subtracts the given means from each image channel.
For example:
means = [123.68, 116.779, 103.939]
image = _mean_image_subtraction(image, means)
Note that the rank of `image` must be known.
Args:
image: a tensor of size [height, width, C].
means: a C-vector of values to subtract from each channel.
Returns:
the centered image.
Raises:
ValueError: If the rank of `image` is unknown, if `image` has a rank other
than three or if the number of channels in `image` doesn't match the
number of values in `means`.
"""
if image.get_shape().ndims != 3:
raise ValueError("Input must be of size [height, width, C>0]")
num_channels = image.get_shape().as_list()[-1]
if len(means) != num_channels:
raise ValueError("len(means) must match the number of channels")
channels = tf.split(axis=2, num_or_size_splits=num_channels, value=image)
for i in range(num_channels):
channels[i] -= means[i]
return tf.concat(axis=2, values=channels)
def vqa_v2_preprocess_image(
image,
height,
width,
mode,
resize_side=512,
distort=True,
image_model_fn="resnet_v1_152",
):
"""vqa v2 preprocess image."""
image = tf.image.convert_image_dtype(image, dtype=tf.float32)
assert resize_side > 0
if resize_side:
image = _aspect_preserving_resize(image, resize_side)
if mode == tf_estimator.ModeKeys.TRAIN:
image = tf.random_crop(image, [height, width, 3])
else:
# Central crop, assuming resize_height > height, resize_width > width.
image = tf.image.resize_image_with_crop_or_pad(image, height, width)
image = tf.clip_by_value(image, 0.0, 1.0)
if mode == tf_estimator.ModeKeys.TRAIN and distort:
image = _flip(image)
num_distort_cases = 4
# pylint: disable=unnecessary-lambda
image = _apply_with_random_selector(
image, lambda x, ordering: _distort_color(x, ordering),
num_cases=num_distort_cases)
if image_model_fn.startswith("resnet_v1"):
# resnet_v1 uses vgg preprocessing
image = image * 255.
image = _mean_image_subtraction(image, [_R_MEAN, _G_MEAN, _B_MEAN])
elif image_model_fn.startswith("resnet_v2"):
# resnet v2 uses inception preprocessing
image = tf.subtract(image, 0.5)
image = tf.multiply(image, 2.0)
return image
================================================
FILE: tensor2tensor/data_generators/wiki.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Data generator for Wikipedia title to article dataset."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import subprocess
import numpy as np
from tensor2tensor.data_generators import generator_utils
from tensor2tensor.data_generators import problem
from tensor2tensor.data_generators import text_encoder
from tensor2tensor.data_generators import text_problems
from tensor2tensor.utils import registry
import tensorflow.compat.v1 as tf
@registry.register_problem
class LanguagemodelWikiXmlV8kL1k(text_problems.ChoppedTextProblem):
"""A language model on English Wikipedia.
XML dump is chopped arbitrarily into sequences of length 1024 tokens,
without regard to article boundaries.
"""
def maybe_prepare_text(self, tmp_dir):
"""Download corpus if necessary, decompress, split into multiple text files.
Args:
tmp_dir: directory containing dataset.
Returns:
list of filepaths for local text files.
"""
compressed_filename = os.path.basename(self.corpus_url)
compressed_filepath = os.path.join(tmp_dir, compressed_filename)
decompressed_filepath = compressed_filepath[:-4]
split_file_prefix = decompressed_filepath + "-part-"
split_filepattern = split_file_prefix + "?????"
split_files = sorted(tf.gfile.Glob(split_filepattern))
if not split_files:
if not tf.gfile.Exists(decompressed_filepath):
if not tf.gfile.Exists(compressed_filepath):
generator_utils.maybe_download(
tmp_dir, compressed_filepath, self.corpus_url)
assert not subprocess.call(["bunzip2", compressed_filepath])
assert tf.gfile.Exists(decompressed_filepath)
assert not subprocess.call([
"split", "--line-bytes=4M", "--suffix-length=5",
"--numeric-suffixes", decompressed_filepath, split_file_prefix])
split_files = sorted(tf.gfile.Glob(split_filepattern))
assert split_files
return split_files
def train_text_filepaths(self, tmp_dir):
all_files = self.maybe_prepare_text(tmp_dir)
return [f for i, f in enumerate(all_files) if i % self.dev_fraction != 0]
def dev_text_filepaths(self, tmp_dir):
all_files = self.maybe_prepare_text(tmp_dir)
return [f for i, f in enumerate(all_files) if i % self.dev_fraction == 0]
@property
def dev_fraction(self):
return 5000
@property
def corpus_url(self):
return ("https://archive.org/download/enwiki-20171201/"
"enwiki-20171201-pages-articles.xml.bz2")
@property
def approx_vocab_size(self):
return 2**13 # 8192
@property
def sequence_length(self):
"""Length of each example (in tokens)."""
return 1024
@property
def max_chars_for_vocab(self):
"""Number of characters of training data to use for generating vocab."""
# magic number for backwards compatibility
return 41800829
@registry.register_problem
class LanguagemodelWikiXmlV8kL4k(LanguagemodelWikiXmlV8kL1k):
"""A language model on English Wikipedia.
XML dump is chopped arbitrarily into sequences of length 4096 tokens,
without regard to article boundaries.
"""
@property
def sequence_length(self):
"""Length of each example (in tokens)."""
return 4096
class LanguagemodelWikiScramble(LanguagemodelWikiXmlV8kL1k):
"""Language modeling on English wikipedia.
"targets" is a sequence of sequence_length tokens - a fragment of an article.
"inputs" is a copy of "targets", but with a random scramble_fraction of the
tokens randomly permuted.
This dataset is intended to test parallel (non-autoregressive) prediction
of the target sequence given the input sequence.
"""
def example_generator(self, encoder, tmp_dir, task_id):
for x in super(LanguagemodelWikiScramble, self).example_generator(
encoder, tmp_dir, task_id):
x["inputs"] = self.scramble(x["targets"])
yield x
@property
def scramble_fraction(self):
raise NotImplementedError()
@property
def has_inputs(self):
return True
@property
def input_space_id(self):
return problem.SpaceID.EN_TOK
@property
def targeted_vocab_size(self):
return 2**13 # 8192
@property
def remainder_policy(self):
"""What to do with leftover tokens."""
return "drop"
def scramble(self, seq):
seq = np.array(seq)
num_permute = int(self.sequence_length * self.scramble_fraction)
full_permutation = np.random.permutation(self.sequence_length)
inverse_full_permutation = np.argsort(full_permutation)
partial_permutation = np.random.permutation(num_permute)
seq = seq[full_permutation]
seq = np.concatenate(
(seq[:num_permute][partial_permutation], seq[num_permute:]))
seq = seq[inverse_full_permutation]
seq = list(seq)
return seq
@registry.register_problem
class LanguagemodelWikiScrambleL128(LanguagemodelWikiScramble):
"""Sequence length 128, 50% scrambled."""
@property
def sequence_length(self):
return 128
@property
def scramble_fraction(self):
return 0.5
@registry.register_problem
class LanguagemodelWikiScrambleL1k(LanguagemodelWikiScramble):
"""Sequence length 1024, 50% scrambled."""
@property
def sequence_length(self):
return 1024
@property
def scramble_fraction(self):
return 0.5
@registry.register_problem
class LanguagemodelWikiNorefV8kL1k(LanguagemodelWikiXmlV8kL1k):
"""A language model on English Wikipedia.
References and internal links are removed from the raw XML.
Special pages (non-articles) are dropped.
This more closely resembles plain text, though there are still some xml
elements, like tables.
Each article is prefixed by a line containing the title and length in
characters - e.g.
title: "Price of Tea in China" length: 12345
During inference time, you can forward generate starting with such a header
in order to obtain a randomly generated article with a given title and
(approximate) length.
Result is chopped arbitrarily into sequences of length 1024 tokens,
without regard to article boundaries.
"""
def filepath_to_unicode_strings(self, filepath):
"""Overrides the base class to clean up the xml dump before tokenizing."""
dump = text_encoder.to_unicode_ignore_errors(tf.gfile.Open(filepath).read())
pages = _dump_to_pages(dump)
ret = u""
for p in pages:
title = _page_to_title(p)
text = _page_to_text(p)
text = _remove_triple_quotes(
_remove_double_brackets(_remove_references(text)))
if u":" in title:
# not a regular article
continue
if len(text) <= 140:
# Probably a redirect or something like that. Skip it.
continue
ret += u"title: \"%s\" length: %d\n%s\n" % (title, len(text), text)
yield ret
@property
def max_chars_for_vocab(self):
"""Number of characters of training data to use for generating vocab."""
# magic number for backwards compatibility
return 21240483
def _dump_to_pages(dump):
"""Extract pages from an xml dump.
Args:
dump: a unicode string
Returns:
a list of unicode strings
"""
pos = 0
ret = []
start_tag = u"\n"
end_tag = u"\n"
while True:
start_pos = dump.find(start_tag, pos)
if start_pos == -1:
break
start_pos += len(start_tag)
end_pos = dump.find(end_tag, start_pos)
if end_pos == -1:
break
ret.append(dump[start_pos:end_pos])
pos = end_pos + len(end_tag)
return ret
def _page_to_title(page):
"""Extract the title from a page.
Args:
page: a unicode string
Returns:
a unicode string
"""
# print("page=%s" % page)
start_tag = u""
end_tag = u""
start_pos = page.find(start_tag)
end_pos = page.find(end_tag)
assert start_pos != -1
assert end_pos != -1
start_pos += len(start_tag)
return page[start_pos:end_pos]
def _page_to_text(page):
"""Extract the text from a page.
Args:
page: a unicode string
Returns:
a unicode string
"""
# text start tag looks like ""
start_pos = page.find(u"", start_pos)
assert end_tag_pos != -1
end_tag_pos += len(u">")
end_pos = page.find(u"")
if end_pos == -1:
return u""
return page[end_tag_pos:end_pos]
def _find_and_replace(text, start_string, end_string, replace_fn):
"""Remove everything found between instances of start_string and end_string.
Replace each such instance with replace_fn(removed_text)
e.g. _find_and_replace(u"the [[fat]] cat [[sat]]", u"[[", u"]]", lambda x: x)
= u"the fat cat sat"
Args:
text: a unicode string
start_string: a unicode string
end_string: a unicode string
replace_fn: a unary function from unicode string to unicode string
Returns:
a string
"""
ret = u""
current_pos = 0
while True:
start_pos = text.find(start_string, current_pos)
if start_pos == -1:
ret += text[current_pos:]
break
ret += text[current_pos:start_pos]
end_pos = text.find(end_string, start_pos + len(start_string))
if end_pos == -1:
break
ret += replace_fn(text[start_pos + len(start_string):end_pos])
current_pos = end_pos + len(end_string)
return ret
def _remove_references(text):
"""Strip out references from wikipedia xml."""
return _find_and_replace(text, u"<ref", u"</ref>", lambda s: "")
def _remove_triple_quotes(text):
"""Strip out triple quotes from wikipedia xml."""
return _find_and_replace(text, u"'''", u"'''", lambda s: s)
def _remove_double_brackets(text):
"""Remove double brackets (internal links) but leave the viewable text.
Args:
text: a unicode string
Returns:
a unicode string
"""
def replacement_fn(s):
if u":" in s:
# this is probably a category or something like that.
return ""
# keep the part after the bar.
bar_pos = s.find(u"|")
if bar_pos == -1:
return s
return s[bar_pos + 1:]
return _find_and_replace(text, u"[[", u"]]", replacement_fn)
@registry.register_problem
class LanguagemodelWikiNorefV8kL16k(LanguagemodelWikiNorefV8kL1k):
"""A language model on English Wikipedia.
References removed. Chopped into segments of 16k tokens.
"""
@property
def sequence_length(self):
"""Length of each example (in tokens)."""
return 2**14
@registry.register_problem
class LanguagemodelWikiNorefV32kL1k(LanguagemodelWikiNorefV8kL1k):
"""32k vocab."""
@property
def approx_vocab_size(self):
return 2**15 # 32768
@property
def max_chars_for_vocab(self):
return 100 * (10 ** 6)
@registry.register_problem
class LanguagemodelWikiNorefV32kL16k(LanguagemodelWikiNorefV32kL1k):
"""A language model on English Wikipedia.
References removed. Chopped into segments of 16k tokens.
"""
@property
def sequence_length(self):
"""Length of each example (in tokens)."""
return 2**14
@registry.register_problem
class LanguagemodelWikiNorefV128kL1k(LanguagemodelWikiNorefV8kL1k):
"""128k vocab."""
@property
def approx_vocab_size(self):
return 2**17 # 131072
@property
def max_chars_for_vocab(self):
return 100 * (10 ** 6)
================================================
FILE: tensor2tensor/data_generators/wiki_lm.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Data generators for untokenized wikipedia LM dataset."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import six
from tensor2tensor.data_generators import generator_utils
from tensor2tensor.data_generators import problem
from tensor2tensor.data_generators import text_problems
from tensor2tensor.utils import registry
import tensorflow.compat.v1 as tf
def concat_generator(filename, up_threshold, low_threshold=10):
"""Generate concatenated lines from file upto up_threshold characters."""
txt = ""
for line in tf.gfile.Open(filename):
line = line.strip()
if len(txt) + len(line) + 1 >= up_threshold:
ret = txt
txt = ""
# We don't yield very short long parts to prevent noisy examples.
if len(ret) > low_threshold and len(ret) < up_threshold:
yield {"targets": ret}
if not txt:
txt = line
else:
txt = " ".join([txt, line])
def mix_generators(generator_list):
"""Given python generators, generate from one, then from another, etc."""
i = 0
l = len(generator_list)
stopiters_seen = 0
while stopiters_seen <= l:
try:
yield six.next(generator_list[i % l])
i += 1
stopiters_seen = 0
except StopIteration:
i += 1
stopiters_seen += 1
# File names and Google drive ids for the training/eval/test Wikipedia data.
_EN_TRAIN_NAME_ID = ("enwiki_train.txt.gz", "1-l02fI15ieMIZk8EnXhzhsvuEYRoznZ8")
_EN_EVAL_NAME_ID = ("enwiki_eval.txt.gz", "1odhDxWKtAPKXwxRw1KCrmlrVewxdXYq7")
_EN_TEST_NAME_ID = ("enwiki_test.txt.gz", "1i1Bg6XqvdRl1LuOiIWbg7ww8Y02Ip5VK")
_DE_TRAIN_NAME_ID = ("dewiki_train.txt.gz", "1FzEwoPonw9xlwX34vLPFInUF8F4X5yJy")
_DE_EVAL_NAME_ID = ("dewiki_eval.txt.gz", "1EKwRRPHyWny0RJ-aqSGMcNfjAlzFl51B")
_DE_TEST_NAME_ID = ("dewiki_test.txt.gz", "1Kr13Y7y_OD3JtUM9riXpFQP9UiHDkcFY")
_FR_TRAIN_NAME_ID = ("frwiki_train.txt.gz", "1etUIEZxMQKORwLGkssE5wlfCxxkeo8WV")
_FR_EVAL_NAME_ID = ("frwiki_eval.txt.gz", "13qrR5ZnHRgIMdcURVpixKL9gTO23GcPc")
_FR_TEST_NAME_ID = ("frwiki_test.txt.gz", "1mQpHRkAV9KXt68de69RwR8dkDi8EEusV")
_RO_TRAIN_NAME_ID = ("rowiki_train.txt.gz", "1wUJTEAlQeDcAwFnBxa8PzE-DCiXSU_W7")
_RO_EVAL_NAME_ID = ("rowiki_eval.txt.gz", "1uIPy2ZgkyArPy_gnsILENjgv4QQmSKtx")
_RO_TEST_NAME_ID = ("rowiki_test.txt.gz", "1kphjN4jXTbw8HyRYKaRE2zY4D7Fr-p7-")
@registry.register_problem
class LanguagemodelEnWiki32k(text_problems.Text2SelfProblem):
"""A language model on the untokenized wikipedia corpus, English."""
train_names_ids = [_EN_TRAIN_NAME_ID]
eval_names_ids = [_EN_EVAL_NAME_ID]
test_names_ids = [_EN_TEST_NAME_ID]
@property
def approx_vocab_size(self):
return 32000
@property
def max_samples_for_vocab(self):
return 128000
@property
def combine_characters_threshold(self):
"""Threshold for upto how many characters to combine in examples."""
return 512*8 # So we should have 512 tokens on average, maybe more.
def is_generate_per_split(self):
return True
@property
def dataset_splits(self):
"""Splits of data to produce and number of output shards for each."""
return [{
"split": problem.DatasetSplit.TRAIN,
"shards": 100,
}, {
"split": problem.DatasetSplit.EVAL,
"shards": 1,
}, {
"split": problem.DatasetSplit.TEST,
"shards": 1,
}]
def generate_samples(self, data_dir, tmp_dir, dataset_split):
"""Generate samples."""
if dataset_split == problem.DatasetSplit.TRAIN:
file_names_ids = self.train_names_ids
elif dataset_split == problem.DatasetSplit.TEST:
file_names_ids = self.test_names_ids
else:
file_names_ids = self.eval_names_ids
wiki_generators = []
for (fname, fid) in file_names_ids:
url = "https://drive.google.com/uc?export=download&id=" + fid
download_path = generator_utils.maybe_download_from_drive(
tmp_dir, fname, url)
wiki_file = os.path.join(tmp_dir, fname[:-3])
if not tf.gfile.Exists(wiki_file):
generator_utils.gunzip_file(download_path, wiki_file)
wiki_generators.append(
concat_generator(wiki_file, self.combine_characters_threshold))
for example in mix_generators(wiki_generators):
yield example
@registry.register_problem
class LanguagemodelEnWiki64k(LanguagemodelEnWiki32k):
"""As above, with 64k vocabulary."""
@property
def approx_vocab_size(self):
return 64000
@registry.register_problem
class LanguagemodelEnWiki64kShorter(LanguagemodelEnWiki64k):
"""With 64k vocabulary and shorter truncation lengths."""
@property
def combine_characters_threshold(self):
"""Threshold for upto how many characters to combine in examples."""
return 384*8
@property
def use_vocab_from_other_problem(self):
return LanguagemodelEnWiki64k()
@registry.register_problem
class LanguagemodelDeWiki32k(LanguagemodelEnWiki32k):
"""A language model on the untokenized wikipedia corpus, German."""
train_names_ids = [_DE_TRAIN_NAME_ID]
eval_names_ids = [_DE_EVAL_NAME_ID]
test_names_ids = [_DE_TEST_NAME_ID]
@registry.register_problem
class LanguagemodelDeWiki64k(LanguagemodelDeWiki32k):
"""As above, with 64k vocabulary."""
@property
def approx_vocab_size(self):
return 64000
@registry.register_problem
class LanguagemodelFrWiki32k(LanguagemodelEnWiki32k):
"""A language model on the untokenized wikipedia corpus, French."""
train_names_ids = [_FR_TRAIN_NAME_ID]
eval_names_ids = [_FR_EVAL_NAME_ID]
test_names_ids = [_FR_TEST_NAME_ID]
@registry.register_problem
class LanguagemodelFrWiki64k(LanguagemodelFrWiki32k):
"""As above, with 64k vocabulary."""
@property
def approx_vocab_size(self):
return 64000
@registry.register_problem
class LanguagemodelRoWiki32k(LanguagemodelEnWiki32k):
"""A language model on the untokenized wikipedia corpus, Romanian."""
train_names_ids = [_RO_TRAIN_NAME_ID]
eval_names_ids = [_RO_EVAL_NAME_ID]
test_names_ids = [_RO_TEST_NAME_ID]
@registry.register_problem
class LanguagemodelRoWiki64k(LanguagemodelRoWiki32k):
"""As above, with 64k vocabulary."""
@property
def approx_vocab_size(self):
return 64000
@registry.register_problem
class LanguagemodelDeEnFrRoWiki64k(LanguagemodelEnWiki32k):
"""A language model on untokenized Wikipedia, 4 languages together."""
train_names_ids = [_DE_TRAIN_NAME_ID, _FR_TRAIN_NAME_ID,
_EN_TRAIN_NAME_ID, _RO_TRAIN_NAME_ID]
eval_names_ids = [_DE_EVAL_NAME_ID, _FR_EVAL_NAME_ID,
_EN_EVAL_NAME_ID, _RO_EVAL_NAME_ID]
test_names_ids = [_DE_TEST_NAME_ID, _FR_TEST_NAME_ID,
_EN_TEST_NAME_ID, _RO_TEST_NAME_ID]
@property
def approx_vocab_size(self):
return 64000
@property
def max_samples_for_vocab(self):
return 256000 # Samples are intertwined, take more to cover 4 languages.
@registry.register_problem
class LanguagemodelDeEnFrRoWiki64kFitbPacked1k(
LanguagemodelDeEnFrRoWiki64k):
"""4 languages fill-in-the-blanks text-to-text problem."""
@property
def use_vocab_from_other_problem(self):
return LanguagemodelDeEnFrRoWiki64k()
@property
def has_inputs(self):
return True
def generate_samples(self, data_dir, tmp_dir, dataset_split):
for example in super(
LanguagemodelDeEnFrRoWiki64kFitbPacked1k, self).generate_samples(
data_dir, tmp_dir, dataset_split):
a, b = generator_utils.random_deinterleave(example["targets"])
yield {"inputs": a, "targets": b}
@property
def num_training_examples(self):
return 3597800
@property
def packed_length(self):
return 1024
@property
def inputs_prefix(self):
return "wiki fill "
@property
def targets_prefix(self):
return "wiki fill "
================================================
FILE: tensor2tensor/data_generators/wiki_multi_problems.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Data generators for Wiki LM and MNLI combined datasets."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from tensor2tensor.data_generators import cnn_dailymail
from tensor2tensor.data_generators import multi_problem
from tensor2tensor.data_generators import multi_problem_v2
from tensor2tensor.data_generators import multinli
from tensor2tensor.data_generators import squad
from tensor2tensor.data_generators import text_problems
from tensor2tensor.data_generators import translate_ende
from tensor2tensor.data_generators import translate_enfr
from tensor2tensor.data_generators import translate_enro
from tensor2tensor.data_generators import wiki_lm
from tensor2tensor.utils import registry
@registry.register_problem
class LanguagemodelEnWikiLMMultiNLISubwords(multi_problem.MultiProblem):
"""Wiki LM and MNLI mixed problem class."""
def __init__(self, was_reversed=False, was_copy=False):
super(LanguagemodelEnWikiLMMultiNLISubwords, self).__init__(
was_reversed, was_copy)
self.task_list.append(wiki_lm.LanguagemodelEnWiki32k())
self.task_list.append(multinli.MultiNLIWikiLMSharedVocab())
@property
def vocab_type(self):
return text_problems.VocabType.SUBWORD
@registry.register_problem
class LanguagemodelEnWikiLMMultiNLISubwordsV2(
multi_problem_v2.MultiText2TextProblem):
"""Wiki LM and MNLI mixed problem class."""
def __init__(self, was_reversed=False, was_copy=False):
problems = [
wiki_lm.LanguagemodelEnWiki32k(),
multinli.MultiNLIWikiLMSharedVocab(),
]
schedule = multi_problem_v2.constant_schedule([0.5, 0.5])
super(LanguagemodelEnWikiLMMultiNLISubwordsV2, self).__init__(
problems, schedule, was_reversed=was_reversed, was_copy=was_copy)
@property
def has_inputs(self):
return False
@property
def use_vocab_from_other_problem(self):
return wiki_lm.LanguagemodelEnWiki32k()
@property
def vocab_type(self):
return text_problems.VocabType.SUBWORD
@registry.register_problem
class LanguagemodelMultiWikiTranslatePacked1k(
multi_problem_v2.MultiText2TextProblem):
"""Wiki-LM, Translation, MNLI, SQUAD mixed problem class."""
def __init__(self, was_reversed=False, was_copy=False):
problems = []
rates = []
for rate, also_reverse, cls in self.problems_and_rates:
for r in [False, True] if also_reverse else [False]:
problems.append(cls(was_reversed=r))
rates.append(rate)
pmf = multi_problem_v2.epoch_rates_to_pmf(problems, epoch_rates=rates)
schedule = multi_problem_v2.constant_schedule(pmf)
super(LanguagemodelMultiWikiTranslatePacked1k, self).__init__(
problems, schedule, was_reversed=was_reversed, was_copy=was_copy)
@property
def problems_and_rates(self):
"""Returns a list of (weight, also_reverse, problem_class) triples."""
return [
(1.0, True, wiki_lm.LanguagemodelDeEnFrRoWiki64kFitbPacked1k),
(1.0, True, translate_ende.TranslateEndeWmtMulti64kPacked1k),
(1.0, True, translate_enfr.TranslateEnfrWmtMulti64kPacked1k),
(1.0, True, translate_enro.TranslateEnroWmtMultiTiny64kPacked1k),
(1.0, True, cnn_dailymail.SummarizeCnnDailymailMulti64kPacked1k),
(1.0, False, multinli.MultiNLIText2textMulti64kPacked1k),
(1.0, False, squad.SquadText2textMulti64kPacked1k),
]
@property
def has_inputs(self):
return True
@property
def use_vocab_from_other_problem(self):
return wiki_lm.LanguagemodelDeEnFrRoWiki64k()
@property
def vocab_type(self):
return text_problems.VocabType.SUBWORD
@property
def packed_length(self):
return 1024
@registry.register_problem
class LanguagemodelMultiWikiTranslatePacked1kV2(
LanguagemodelMultiWikiTranslatePacked1k):
"""Higher rates for rarer problems."""
@property
def problems_and_rates(self):
"""Returns a list of (weight, also_reverse, problem_class) triples."""
return [
(1.0, True, wiki_lm.LanguagemodelDeEnFrRoWiki64kFitbPacked1k),
(3.0, True, translate_ende.TranslateEndeWmtMulti64kPacked1k),
(1.0, True, translate_enfr.TranslateEnfrWmtMulti64kPacked1k),
(100.0, True, translate_enro.TranslateEnroWmtMultiTiny64kPacked1k),
(1.0, True, cnn_dailymail.SummarizeCnnDailymailMulti64kPacked1k),
(10.0, False, multinli.MultiNLIText2textMulti64kPacked1k),
(10.0, False, squad.SquadText2textMulti64kPacked1k),
]
@registry.register_problem
class LanguagemodelEnWikiLMMultiNLISubwords64k(multi_problem.MultiProblem):
"""Wiki LM and MNLI mixed problem class."""
def __init__(self, was_reversed=False, was_copy=False):
super(LanguagemodelEnWikiLMMultiNLISubwords64k, self).__init__(
was_reversed, was_copy)
self.task_list.append(wiki_lm.LanguagemodelEnWiki64k())
self.task_list.append(multinli.MultiNLIWikiLMSharedVocab64k())
@property
def vocab_type(self):
return text_problems.VocabType.SUBWORD
@registry.register_problem
class LanguagemodelEnWikiLMShortMultiNLISubwords64k(multi_problem.MultiProblem):
"""Wiki LM and MNLI mixed problem class."""
def __init__(self, was_reversed=False, was_copy=False):
super(LanguagemodelEnWikiLMShortMultiNLISubwords64k, self).__init__(
was_reversed, was_copy)
self.task_list.append(wiki_lm.LanguagemodelEnWiki64kShorter())
self.task_list.append(multinli.MultiNLIWikiLMSharedVocab64k())
@property
def vocab_type(self):
return text_problems.VocabType.SUBWORD
@registry.register_problem
class LanguagemodelEnWikiLMSummarizeCnndmSubwords(multi_problem.MultiProblem):
"""Wiki LM and CNN/DM summarization mixed problem class."""
def __init__(self, was_reversed=False, was_copy=False):
super(LanguagemodelEnWikiLMSummarizeCnndmSubwords, self).__init__(
was_reversed, was_copy)
self.task_list.append(wiki_lm.LanguagemodelEnWiki32k())
self.task_list.append(
cnn_dailymail.SummarizeCnnDailymailWikiLMSharedVocab())
@property
def vocab_type(self):
return text_problems.VocabType.SUBWORD
@registry.register_problem
class LanguagemodelEnWikiLMSummarizeCnndmSubwords64k(
multi_problem.MultiProblem):
"""Wiki LM and CNN/DM summarization mixed problem class."""
def __init__(self, was_reversed=False, was_copy=False):
super(LanguagemodelEnWikiLMSummarizeCnndmSubwords64k, self).__init__(
was_reversed, was_copy)
self.task_list.append(wiki_lm.LanguagemodelEnWiki64k())
self.task_list.append(
cnn_dailymail.SummarizeCnnDailymailWikiLMSharedVocab64k())
@property
def vocab_type(self):
return text_problems.VocabType.SUBWORD
@registry.register_problem
class LanguagemodelMultiWikiTranslateFr(multi_problem.MultiProblem):
"""Wiki multi-lingual LM and En-Fr translation."""
def __init__(self, was_reversed=False, was_copy=False):
super(LanguagemodelMultiWikiTranslateFr, self).__init__(
was_reversed, was_copy)
self.task_list.append(wiki_lm.LanguagemodelDeEnFrRoWiki64k())
self.task_list.append(translate_enfr.TranslateEnfrWmtMulti64k())
@property
def vocab_type(self):
return text_problems.VocabType.SUBWORD
@registry.register_problem
class LanguagemodelMultiWikiTranslate(multi_problem.MultiProblem):
"""Wiki multi-lingual LM and multiple translations."""
def __init__(self, was_reversed=False, was_copy=False):
super(LanguagemodelMultiWikiTranslate, self).__init__(
was_reversed, was_copy)
self.task_list.append(wiki_lm.LanguagemodelDeEnFrRoWiki64k())
self.task_list.append(translate_ende.TranslateEndeWmtMulti64k())
self.task_list.append(translate_enfr.TranslateEnfrWmtMulti64k())
self.task_list.append(translate_enro.TranslateEnroWmtMultiTiny64k())
self.task_list.append(translate_ende.TranslateEndeWmtMulti64k(
was_reversed=True))
self.task_list.append(translate_enfr.TranslateEnfrWmtMulti64k(
was_reversed=True))
self.task_list.append(translate_enro.TranslateEnroWmtMultiTiny64k(
was_reversed=True))
self.task_list.append(
cnn_dailymail.SummarizeCnnDailymailWikiLMMultiVocab64k())
self.task_list.append(multinli.MultiNLIWikiLMMultiVocab64k())
self.task_list.append(squad.SquadConcatMulti64k())
@property
def vocab_type(self):
return text_problems.VocabType.SUBWORD
@registry.register_problem
class LanguagemodelEnWikiLMSummarizeFrac1CnndmSubwords64k(
multi_problem.MultiProblem):
"""Wiki LM and CNN/DM summarization mixed problem class."""
def __init__(self, was_reversed=False, was_copy=False):
super(LanguagemodelEnWikiLMSummarizeFrac1CnndmSubwords64k, self).__init__(
was_reversed, was_copy)
self.task_list.append(wiki_lm.LanguagemodelEnWiki64k())
self.task_list.append(
cnn_dailymail.SummarizeFrac1CnnDailymailWikiLMSharedVocab64k())
@property
def vocab_type(self):
return text_problems.VocabType.SUBWORD
@registry.register_problem
class LanguagemodelEnWikiLMSummarizeFrac2CnndmSubwords64k(
multi_problem.MultiProblem):
"""Wiki LM and CNN/DM summarization mixed problem class."""
def __init__(self, was_reversed=False, was_copy=False):
super(LanguagemodelEnWikiLMSummarizeFrac2CnndmSubwords64k, self).__init__(
was_reversed, was_copy)
self.task_list.append(wiki_lm.LanguagemodelEnWiki64k())
self.task_list.append(
cnn_dailymail.SummarizeFrac2CnnDailymailWikiLMSharedVocab64k())
@property
def vocab_type(self):
return text_problems.VocabType.SUBWORD
@registry.register_problem
class LanguagemodelEnWikiLMSummarizeFrac5CnndmSubwords64k(
multi_problem.MultiProblem):
"""Wiki LM and CNN/DM summarization mixed problem class."""
def __init__(self, was_reversed=False, was_copy=False):
super(LanguagemodelEnWikiLMSummarizeFrac5CnndmSubwords64k, self).__init__(
was_reversed, was_copy)
self.task_list.append(wiki_lm.LanguagemodelEnWiki64k())
self.task_list.append(
cnn_dailymail.SummarizeFrac5CnnDailymailWikiLMSharedVocab64k())
@property
def vocab_type(self):
return text_problems.VocabType.SUBWORD
@registry.register_problem
class LanguagemodelEnWikiLMSummarizeFrac10CnndmSubwords64k(
multi_problem.MultiProblem):
"""Wiki LM and CNN/DM summarization mixed problem class."""
def __init__(self, was_reversed=False, was_copy=False):
super(LanguagemodelEnWikiLMSummarizeFrac10CnndmSubwords64k, self).__init__(
was_reversed, was_copy)
self.task_list.append(wiki_lm.LanguagemodelEnWiki64k())
self.task_list.append(
cnn_dailymail.SummarizeFrac10CnnDailymailWikiLMSharedVocab64k())
@property
def vocab_type(self):
return text_problems.VocabType.SUBWORD
@registry.register_problem
class LanguagemodelEnWikiLMSummarizeFrac20CnndmSubwords64k(
multi_problem.MultiProblem):
"""Wiki LM and CNN/DM summarization mixed problem class."""
def __init__(self, was_reversed=False, was_copy=False):
super(LanguagemodelEnWikiLMSummarizeFrac20CnndmSubwords64k, self).__init__(
was_reversed, was_copy)
self.task_list.append(wiki_lm.LanguagemodelEnWiki64k())
self.task_list.append(
cnn_dailymail.SummarizeFrac20CnnDailymailWikiLMSharedVocab64k())
@property
def vocab_type(self):
return text_problems.VocabType.SUBWORD
@registry.register_problem
class LanguagemodelEnWikiLMSummarizeFrac50CnndmSubwords64k(
multi_problem.MultiProblem):
"""Wiki LM and CNN/DM summarization mixed problem class."""
def __init__(self, was_reversed=False, was_copy=False):
super(LanguagemodelEnWikiLMSummarizeFrac50CnndmSubwords64k, self).__init__(
was_reversed, was_copy)
self.task_list.append(wiki_lm.LanguagemodelEnWiki64k())
self.task_list.append(
cnn_dailymail.SummarizeFrac50CnnDailymailWikiLMSharedVocab64k())
@property
def vocab_type(self):
return text_problems.VocabType.SUBWORD
@registry.register_problem
class LanguagemodelEnWikiLMSquadConcatSubwords(multi_problem.MultiProblem):
"""Wiki LM and MNLI mixed problem class."""
def __init__(self, was_reversed=False, was_copy=False):
super(LanguagemodelEnWikiLMSquadConcatSubwords, self).__init__(
was_reversed, was_copy)
self.task_list.append(wiki_lm.LanguagemodelEnWiki32k())
self.task_list.append(multinli.SquadConcatSharedVocab())
@property
def vocab_type(self):
return text_problems.VocabType.SUBWORD
================================================
FILE: tensor2tensor/data_generators/wiki_revision.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
r"""Data extraction/preprocessing for processing wiki history dumps for GEC.
We use a set of heuristics to distill prose from the wikipedia xml. We produce
source-target pairs of text reflecting wikipedia edits.
WikiRevision problem - fragment of older revision -> fragment of newer revision.
This implements data extraction from wikipedia as desribed in the paper,
Weakly Supervised Grammatical Error Correction using Iterative Decoding
(https://arxiv.org/pdf/1811.01710.pdf).
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import math
import random
from absl import flags
from tensor2tensor.data_generators import generator_utils
from tensor2tensor.data_generators import text_encoder
from tensor2tensor.data_generators import text_problems
from tensor2tensor.data_generators import wiki_revision_utils
from tensor2tensor.utils import metrics
from tensor2tensor.utils import registry
import tensorflow.compat.v1 as tf
FLAGS = flags.FLAGS
flags.DEFINE_integer("wiki_revision_num_train_shards", 50,
"Set the number of training shards to be output.")
flags.DEFINE_integer("wiki_revision_num_dev_shards", 1,
"Set the number of dev shards to be output.")
flags.DEFINE_string(
"wiki_revision_data_prefix", "",
"Specify the prefix for input data. Expects 7z compressed Wikipedia XML "
"files, available at https://dumps.wikimedia.org/enwiki/latest/.")
flags.DEFINE_string(
"wiki_revision_vocab_file", "",
"Specify a wordpieces vocabulary with which to encode the text. Will "
"generate one from data if not specified.")
flags.DEFINE_integer(
"wiki_revision_max_examples_per_shard", 0,
"Use this to set a cap on examples per shard. "
"0 is no cap.")
# Data filtration heuristics:
flags.DEFINE_integer("wiki_revision_max_page_size_exp", 26,
"Exponent for 2**X byte cap on page size.")
flags.DEFINE_float(
"wiki_revision_max_equal_to_diff_ratio", 0,
"Max ratio between count of equal, diff chars for generated "
"examples. Ratio of 1 means examples with more diff chars "
"than equal chars will be tossed out.")
flags.DEFINE_float(
"wiki_revision_revision_skip_factor", 1.5,
"If >1, process only logarithmically many revisions. "
"This avoids blowup in runtime due to many-revision pages. "
"See wiki_revision_utils.include_revision for details.")
flags.DEFINE_float("wiki_revision_percent_identical_examples", 0.04,
"Percent of generated examples for which source == target.")
flags.DEFINE_bool(
"wiki_revision_introduce_errors", True, "Add errors to the data."
"See wiki_revision_utils.introduce_errors for details.")
@registry.register_problem
class WikiRevision(text_problems.Text2TextProblem):
"""Old segment -> revised segment.
Data filtration heuristics:
wiki_revision_max_page_size_exp:
pages above this # of bytes are thrown out
wiki_revision_revision_skip_factor:
rate of logarithmic downsampling of revision history list
wiki_revision_percent_identical_examples:
how many identitcal examples to admit, as percent of total examples
wiki_revision_introduce_errors:
whether or not to introduce spelling-type errors on the source side
wiki_revision_max_equal_to_diff_ratio:
whether or not to introduce spelling-type errors on the source side
Vocab size=32k
Maximum input/target length = 1024 wordpiece tokens
"""
num_identity_examples = 0
num_total_examples = 0
num_identity_examples = 0
num_pages = 0
num_revisions_total = 0
num_revisions_admitted = 0
num_examples_thrown_out_identity = 0
num_examples_thrown_out_too_long = 0
num_examples_thrown_out_edit_distance = 0
num_examples_with_introduced_error = 0
num_introduced_errors = 0
num_source_tokens = 0
num_target_tokens = 0
corpus_files = None
@property
def approx_vocab_size(self):
return 2**15 # 32K
@property
def strip(self):
"""Whether to strip wikipedia-stuff to get plain text."""
return True
@property
def wiki_revision_skip_factor(self):
"""If this value is >1.0, process only logarithmically many revisions."""
return FLAGS.wiki_revision_revision_skip_factor
@property
def max_segment_length(self):
"""Maximum number of input/target wordpiece tokens."""
return 256
@property
def max_examples_per_shard(self):
"""Maximum number of examples to generate per shard. 0=unlimited."""
return FLAGS.wiki_revision_max_examples_per_shard
def aggregate_job_stats(self):
# Aggregate job stats for output.
stat = []
# Run stats.
stat.append("Flags for job:\n"
"Dev shards: {}\n"
"Train shards: {}\n"
"Revision skip factor: {}\n"
"Max page size: 2**{}\n"
"Introduce errors: {}\n"
"Max edit ratio: {}\n"
"Percent Identical Examples: {}\n"
"".format(FLAGS.wiki_revision_num_dev_shards,
FLAGS.wiki_revision_num_train_shards,
FLAGS.wiki_revision_revision_skip_factor,
FLAGS.wiki_revision_max_page_size_exp,
FLAGS.wiki_revision_introduce_errors,
FLAGS.wiki_revision_max_equal_to_diff_ratio,
FLAGS.wiki_revision_percent_identical_examples))
# File stats.
stat.append("corpus files: {}\n"
"\tnames: {}\n"
"\tpages per input file: {:.1f}\n"
"".format(
len(self.corpus_files), self.corpus_files,
(0 if not self.corpus_files else
self.num_pages / len(self.corpus_files))))
# Page stats.
stat.append(
"pages processed: {}\n"
"\trevisions per page: {:.2f}, total: {}\n"
"\trevisions admitted per page: {:.2f}, percent of total: {:.2f}\n"
"".format(
self.num_pages, (0 if not self.num_pages else
self.num_revisions_total / self.num_pages),
self.num_revisions_total,
(0 if not self.num_pages else
self.num_revisions_admitted / self.num_pages),
(0 if not self.num_revisions_total else
100 * self.num_revisions_admitted / self.num_revisions_total)))
# Revision stats.
stat.append(
"revisions admitted: {}\n"
"\texamples generated per revision: {:.2f}\n"
"".format(self.num_revisions_admitted,
(0 if not self.num_revisions_admitted else
self.num_total_examples / self.num_revisions_admitted)))
# Example stats.
stat.append(
"examples generated: {}\n"
"\twith error introduced: {}, percent of total: {:.2f}\n"
"\ttotal errors introduced: {}, errors per errorred example: {:.2f}\n"
"\texamples thrown out: {}\n"
"\t\ttoo long: {}\n"
"\t\tidentity: {}\n"
"\t\tedit distance: {}\n"
"\tremaining identity examples: {}\n"
"\tratio identity (actual, desired): {:.3f}, {}\n"
"".format(
self.num_total_examples, self.num_examples_with_introduced_error,
(0 if not self.num_total_examples else 100 *
self.num_examples_with_introduced_error / self.num_total_examples),
self.num_introduced_errors,
(0 if not self.num_examples_with_introduced_error else
self.num_introduced_errors /
self.num_examples_with_introduced_error),
self.num_examples_thrown_out_too_long +
self.num_examples_thrown_out_identity +
self.num_examples_thrown_out_edit_distance,
self.num_examples_thrown_out_too_long,
self.num_examples_thrown_out_identity,
self.num_examples_thrown_out_edit_distance,
self.num_identity_examples,
(0 if not self.num_total_examples else
self.num_identity_examples / self.num_total_examples),
FLAGS.wiki_revision_percent_identical_examples))
# Token stats.
stat.append("tokens generated: {}\n"
"\tsource: {}\n"
"\ttarget: {}\n"
"\tper example: {:.2f}\n"
"\t\tsource: {:.2f}\n"
"\t\ttarget: {:.2f}\n"
"".format(self.num_source_tokens + self.num_target_tokens,
self.num_source_tokens, self.num_target_tokens,
(0 if not self.num_total_examples else
(self.num_source_tokens + self.num_target_tokens) /
self.num_total_examples),
(0 if not self.num_total_examples else
self.num_source_tokens / self.num_total_examples),
(0 if not self.num_total_examples else
self.num_target_tokens / self.num_total_examples)))
return "\n".join(stat)
def generate_data(self, data_dir, tmp_dir, task_id=-1):
if task_id == -1 or task_id is None:
for i in range(FLAGS.wiki_revision_num_train_shards +
FLAGS.wiki_revision_num_dev_shards):
self.generate_data(data_dir, tmp_dir, i)
return
tf.logging.info(
"Flags for job (task_id {}): "
"Dev shards: {}, Train shards: {}, "
"Revision skip factor: {}, Max page size: 2**{}, Introduce errors: {},"
"Percent Identical Examples: {}"
"".format(task_id, FLAGS.wiki_revision_num_dev_shards,
FLAGS.wiki_revision_num_train_shards,
FLAGS.wiki_revision_revision_skip_factor,
FLAGS.wiki_revision_max_page_size_exp,
FLAGS.wiki_revision_introduce_errors,
FLAGS.wiki_revision_percent_identical_examples))
if FLAGS.wiki_revision_vocab_file:
encoder = wiki_revision_utils.get_encoder_from_vocab(
FLAGS.wiki_revision_vocab_file)
else:
encoder = wiki_revision_utils.get_or_generate_vocabulary(
data_dir, tmp_dir, FLAGS.wiki_revision_data_prefix,
FLAGS.wiki_revision_max_page_size_exp, self.approx_vocab_size,
self.strip)
random.seed(123)
if task_id < FLAGS.wiki_revision_num_train_shards:
out_file = self.training_filepaths(
data_dir, FLAGS.wiki_revision_num_train_shards,
shuffled=False)[task_id]
else:
out_file = self.dev_filepaths(
data_dir, FLAGS.wiki_revision_num_dev_shards,
shuffled=False)[task_id - FLAGS.wiki_revision_num_train_shards]
tf.logging.info("Generating files for path: %s", out_file)
self.corpus_files = wiki_revision_utils.corpus_files_for_shard(
task_id, FLAGS.wiki_revision_num_train_shards,
FLAGS.wiki_revision_num_dev_shards, FLAGS.wiki_revision_data_prefix)
example_generator = self.generator(encoder, self.corpus_files, tmp_dir)
packed_example_generator = self._maybe_pack_examples(example_generator)
generator_utils.generate_files(packed_example_generator, [out_file])
generator_utils.shuffle_dataset([out_file])
tf.logging.info(
"Job stats: identity examples: {}, total examples {}, ratio: {}".format(
self.num_identity_examples, self.num_total_examples,
(1 + self.num_identity_examples) / (1 + self.num_total_examples)))
job_stats_string = self.aggregate_job_stats()
out_dir, filename = out_file.replace("-unshuffled", "").rsplit("/", 1)
stats_prefix = "/stats_"
stats_file_path = "".join([out_dir, stats_prefix, filename])
if tf.gfile.Exists(
stats_file_path) and tf.gfile.Open(stats_file_path).size() != 0:
tf.logging.info("Skipping writing stats because output file exists.")
else:
with tf.gfile.Open(stats_file_path, "w") as out:
tf.logging.info("Writing job stats to {}".format(stats_file_path))
out.write(job_stats_string)
tf.logging.info(job_stats_string)
def generator(self, encoder, corpus_files, tmp_dir):
for page in wiki_revision_utils.corpus_page_generator(
corpus_files, tmp_dir, FLAGS.wiki_revision_max_page_size_exp):
self.num_pages += 1
examples = self.page_to_examples(page, encoder)
for x in examples:
yield x
if self.num_total_examples % 100000 == 0:
tf.logging.info(
u"page count={} num_total_examples={} id={} title={}".format(
self.num_pages, self.num_total_examples, page["id"],
page["title"]))
if (self.max_examples_per_shard and
self.num_total_examples >= self.max_examples_per_shard):
tf.logging.info(
"Examples per shard {} >= max_examples_per_shard {}. Shutting down."
.format(self.num_total_examples, self.max_examples_per_shard))
break
tf.logging.info(
"Total pages: {}, total examples: {}, examples per page: {}".format(
self.num_pages, self.num_total_examples, 0 if not self.num_pages
else self.num_total_examples / self.num_pages))
def page_to_examples(self, page, encoder):
revisions = page["revisions"]
self.num_revisions_total += len(revisions)
if len(revisions) < 2:
return []
revisions = [
wiki_revision_utils.get_text(r)
for n, r in enumerate(revisions)
if wiki_revision_utils.include_revision(
n, self.wiki_revision_skip_factor) or n + 1 == len(revisions)
]
self.num_revisions_admitted += len(revisions)
ret = []
for i in range(len(revisions) - 1):
old_revision = revisions[i]
new_revision = revisions[i + 1]
if FLAGS.wiki_revision_introduce_errors:
old_revision_text, num_added_err = wiki_revision_utils.introduce_errors(
revisions[i])
if num_added_err:
self.num_introduced_errors += num_added_err
self.num_examples_with_introduced_error += 1
else:
old_revision_text = revisions[i]
new_revision_text = revisions[i + 1]
if encoder:
# Encode text into list of ids, if a text encoder is present.
old_revision = encoder.encode(old_revision_text)
new_revision = encoder.encode(new_revision_text)
else:
# Retain text (as list of characters), if a text encoder is not present.
old_revision = old_revision_text
new_revision = new_revision_text
ret.extend(
self.make_examples(
encoder,
old_revision,
new_revision,
max_length=self.max_segment_length,
percent_identical_examples=FLAGS
.wiki_revision_percent_identical_examples))
return ret
def make_examples(self,
encoder,
old_snapshot,
new_snapshot,
max_length=1024,
percent_identical_examples=0.01,
max_length_distance=0):
"""Produce training examples based on a pair of snapshots.
Aligns the snapshots, then chops at a random subset of the alignment points
to create (old snippet -> new snippet) examples.
Most negative examples (those with no changes) are discarded, but we
keep some of them, maintaining a proportion in the final data
determined by percent_identical_examples.
Args:
encoder: the subword text encoder
old_snapshot: a list of ids
new_snapshot: a list of ids
max_length: an integer. Maximum length of "inputs" and "targets".
percent_identical_examples: a float
max_length_distance: an integer. Max token edit dist for admitted examples
Returns:
a list of feature dictionaries. The dictionaries have
"inputs" and "targets" populated. text_encoder.EOS is appended to both.
"""
ret = []
eos_sequence = [text_encoder.EOS_ID]
# Pick a per-token cut probability with a log-uniform distribution between
# 1/4 and 1/(max_length / 2)
bound1 = -math.log(4.0)
bound2 = -math.log(max_length / 2.0)
cut_prob = math.exp(random.random() * (bound2 - bound1) + bound1)
opcodes = wiki_revision_utils.fast_match_sequences(old_snapshot,
new_snapshot)
cut_points = [(0, 0)]
for tag, i1, i2, j1, j2 in opcodes:
if tag == "equal":
for i in range(i1, i2 + 1):
if random.random() < cut_prob:
cut_points.append((i, i + j1 - i1))
cut_points.append((len(old_snapshot), len(new_snapshot)))
src_tgt_pairs = []
for cut_number in range(len(cut_points) - 1):
i1, j1 = cut_points[cut_number]
i2, j2 = cut_points[cut_number + 1]
old_segment = old_snapshot[i1:i2]
new_segment = new_snapshot[j1:j2]
src_tgt_pairs.append((old_segment, new_segment))
src_tgt_pairs, thrown_edit_count = wiki_revision_utils.edit_distance_filter(
wiki_revision_utils.throw_empty_pairs(src_tgt_pairs),
FLAGS.wiki_revision_max_equal_to_diff_ratio)
self.num_examples_thrown_out_edit_distance += thrown_edit_count
for source, target in src_tgt_pairs:
# Add EOS segment.
old_segment = source + eos_sequence
new_segment = target + eos_sequence
if len(old_segment) <= max_length and len(new_segment) <= max_length:
if max_length_distance and (abs(len(old_segment) - len(new_segment)) >
max_length_distance):
self.num_examples_thrown_out_edit_distance += 1
continue
if old_segment == new_segment:
# If current proportion of identity is below target
# percent_identical_examples, then roll for a 50% chance to add an
# identitical example. Random roll preserves nondeterminism.
# percent_identical_examples, then add identitical example.
# Random roll preserves nondeterminism in selecting identity examples.
if (((self.num_identity_examples) / (1 + self.num_total_examples)) >
percent_identical_examples) or random.random() > 0.5:
self.num_examples_thrown_out_identity += 1
continue
else:
self.num_identity_examples += 1
self.num_total_examples += 1
self.num_source_tokens += len(old_segment) - 1
self.num_target_tokens += len(new_segment) - 1
ret.append({"inputs": old_segment, "targets": new_segment})
else:
self.num_examples_thrown_out_too_long += 1
return ret
def eval_metrics(self):
return [
metrics.Metrics.ACC,
metrics.Metrics.ACC_TOP5,
metrics.Metrics.ACC_PER_SEQ,
metrics.Metrics.NEG_LOG_PERPLEXITY,
]
@property
def invert_prob(self):
"""Ratio of e^2 positive forward to backward examples."""
return 1.0 / (1.0 + math.exp(2.0))
@registry.register_problem
class WikiRevisionPacked1k(WikiRevision):
"""Packed version for TPU."""
@property
def packed_length(self):
return 1024
@registry.register_problem
class WikiRevisionPacked256(WikiRevision):
"""Packed version for TPU."""
@property
def packed_length(self):
return 256
@property
def max_segment_length(self):
return 256
================================================
FILE: tensor2tensor/data_generators/wiki_revision_utils.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Utilties for data generation for Wikipedia Revision problem.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import math
import os
import random
import re
import subprocess
from tensor2tensor.data_generators import generator_utils
from tensor2tensor.data_generators import text_encoder
import tensorflow.compat.v1 as tf
def to_unicode(s):
return s.decode("utf-8")
def include_revision(revision_num, skip_factor=1.1):
"""Decide whether to include a revision.
If the number of revisions is large, we exclude some revisions to avoid
a quadratic blowup in runtime, since the article is likely also large.
We make the ratio between consecutive included revision numbers
appproximately equal to "factor".
Args:
revision_num: an integer
skip_factor: a floating point number >= 1.0
Returns:
a boolean
"""
if skip_factor <= 1.0:
return True
return (int(math.log1p(revision_num) / math.log(skip_factor)) != int(
math.log(revision_num + 2.0) / math.log(skip_factor)))
def file_page_generator(my_file, max_page_size=2**28):
"""Read wikipedia pages from a history dump.
Since some pages can be terabytes in size (with all the revisions),
we limit page size to max_page_size bytes.
Args:
my_file: an open file object.
max_page_size: an integer
Yields:
strings
"""
page_start = " \n"
page_end = " \n"
chunk_size = max_page_size
page_start = " \n"
page_end = " \n"
leftovers = ""
while True:
chunk = my_file.read(chunk_size)
try:
chunk = to_unicode(chunk)
except UnicodeDecodeError:
chunk = ""
if not chunk:
break
chunk = leftovers + chunk
current_pos = 0
while True:
start_pos = chunk.find(page_start, current_pos)
if start_pos == -1:
break
end_pos = chunk.find(page_end, start_pos)
if end_pos == -1:
if len(chunk) - start_pos > max_page_size:
leftovers = ""
else:
leftovers = chunk[start_pos:]
break
raw_page = chunk[start_pos + len(page_start):end_pos]
if len(raw_page) < max_page_size:
ret = parse_page(raw_page)
if ret:
yield ret
current_pos = end_pos + len(page_end)
def get_title(page):
"""Extract the title from a page.
Args:
page: a string
Returns:
a string
"""
start_pos = page.find("")
end_pos = page.find("")
assert start_pos != -1
assert end_pos != -1
start_pos += len("")
return page[start_pos:end_pos]
def get_id(page):
"""Extract the id from a page.
Args:
page: a string
Returns:
an integer
"""
start_pos = page.find("<id>")
end_pos = page.find("</id>")
assert start_pos != -1
assert end_pos != -1
start_pos += len("<id>")
return int(page[start_pos:end_pos])
def get_revisions(page):
"""Extract the revisions of a page.
Args:
page: a string
Returns:
a list of strings
"""
start_string = " <revision>\n"
end_string = " </revision>\n"
ret = []
current_pos = 0
while True:
start_pos = page.find(start_string, current_pos)
if start_pos == -1:
break
end_pos = page.find(end_string, start_pos)
assert end_pos != -1
ret.append(page[start_pos + len(start_string):end_pos])
current_pos = end_pos + len(end_string)
return ret
def parse_page(raw_page):
"""Create a dictionary with title, id, and list of revisions.
The dictionary contains:
"title": a string
"id": an integer
"revisions": a list of strings
Args:
raw_page: a string
Returns:
a dictionary, or None in the case of an error.
"""
ret = {"title": get_title(raw_page), "id": get_id(raw_page)}
if ":" in ret["title"]:
return None
ret["revisions"] = get_revisions(raw_page)
return ret
def maybe_copy_file_to_directory(source_filepath, target_directory):
"""Copy a file to a directory if it is not already there.
Returns the target filepath.
Args:
source_filepath: a string
target_directory: a string
Returns:
a string
"""
if not tf.gfile.Exists(target_directory):
tf.logging.info("Creating directory %s" % target_directory)
os.mkdir(target_directory)
target_filepath = os.path.join(target_directory,
os.path.basename(source_filepath))
if not tf.gfile.Exists(target_filepath):
tf.logging.info("Copying %s to %s" % (source_filepath, target_filepath))
tf.gfile.Copy(source_filepath, target_filepath)
statinfo = os.stat(target_filepath)
tf.logging.info("Successfully copied %s, %s bytes." % (target_filepath,
statinfo.st_size))
else:
tf.logging.info("Not copying, file already found: %s" % target_filepath)
return target_filepath
def corpus_page_generator(corpus_files, tmp_dir, max_page_size_exp):
"""Generate pages from a list of .7z encoded history dumps.
Args:
corpus_files: a list of strings
tmp_dir: a string
max_page_size_exp: an integer
Yields:
strings
"""
for remote_filepath in corpus_files:
filepath = maybe_copy_file_to_directory(remote_filepath, tmp_dir)
tf.logging.info("Reading from " + filepath)
command = ["7z", "x", "-so", filepath]
tf.logging.info("Running command: %s", command)
p = subprocess.Popen(command, stdout=subprocess.PIPE, bufsize=-1)
for page in file_page_generator(p.stdout, 2**max_page_size_exp):
yield page
def get_text(revision, strip=True):
"""Extract the text from a revision.
Args:
revision: a string
strip: a boolean
Returns:
a string
"""
# text start tag looks like "<text ..otherstuff>"
start_pos = revision.find("<text")
assert start_pos != -1
end_tag_pos = revision.find(">", start_pos)
assert end_tag_pos != -1
end_tag_pos += len(">")
end_pos = revision.find("</text>")
if end_pos == -1:
ret = ""
else:
ret = revision[end_tag_pos:end_pos]
if strip:
ret = strip_text(ret)
return ret
def strip_text(text):
"""Strip wikipedia-stuff out of text, making it mostly prose.
The reason for this is to learn a model that is good at editing prose.
Args:
text: a string
Returns:
a string
"""
return _remove_boring_lines(
_remove_triple_quotes(
_remove_double_brackets(
_remove_references(_remove_curly_braces(text)))))
def _find_and_replace(text, start_string, end_string, replace_fn):
"""Remove everything found between instances of start_string and end_string.
Replace each such instance with replace_fn(removed_text)
e.g. _find_and_replace("the [[fat]] cat [[sat]]", "[[", "]]", lambda x: x)
= "the fat cat sat"
Args:
text: a string
start_string: a string
end_string: a string
replace_fn: a unary function from string to string
Returns:
a string
"""
ret = ""
current_pos = 0
while True:
start_pos = text.find(start_string, current_pos)
if start_pos == -1:
ret += text[current_pos:]
break
ret += text[current_pos:start_pos]
end_pos = text.find(end_string, start_pos + len(start_string))
if end_pos == -1:
break
ret += replace_fn(text[start_pos + len(start_string):end_pos])
current_pos = end_pos + len(end_string)
return ret
def _remove_references(text):
return _find_and_replace(text, "<ref", "</ref>", lambda s: "")
def _remove_triple_quotes(text):
return _find_and_replace(text, "'''", "'''", lambda s: s)
def _remove_curly_braces(text):
"""Remove everything in curly braces.
Curly braces may be nested, so we keep track of depth.
Args:
text: a string
Returns:
a string
"""
current_pos = 0
depth = 0
ret = ""
for match in re.finditer("[{}]", text):
if depth == 0:
ret += text[current_pos:match.start()]
depth += 1 if text[match.start()] == "{" else -1
current_pos = match.end()
if depth != 0:
# Many articles have mismatched braces, but it still seems better to remove
# them than not.
pass
else:
ret += text[current_pos:]
return ret
def _remove_double_brackets(text):
"""Remove double brackets, but leave the viewable text.
Args:
text: a string
Returns:
a string
"""
def replacement_fn(s):
if ":" in s:
# this is probably a category or something like that.
return ""
# keep the part after the bar.
bar_pos = s.find("|")
if bar_pos == -1:
return s
return s[bar_pos + 1:]
return _find_and_replace(text, "[[", "]]", replacement_fn)
def _remove_boring_lines(text):
"""Remove lines that do not start with a letter or a quote.
From inspecting the data, this seems to leave in most prose and remove
most weird stuff.
Args:
text: a string
Returns:
a string
"""
lines = text.split("\n")
filtered = [line for line in lines if re.match("[a-zA-z\"\']", line)]
return "\n".join(filtered)
def all_corpus_files(data_prefix):
return sorted(tf.gfile.Glob(data_prefix + "*"))
def corpus_files_for_shard(shard_num, train_shards, dev_shards, data_prefix):
corpus_files = [
filename for i, filename in enumerate(all_corpus_files(data_prefix))
if i % (train_shards + dev_shards) == shard_num
]
tf.logging.info("Corpus files for shard %s: %s", shard_num, corpus_files)
assert shard_num < (train_shards + dev_shards)
return corpus_files
def vocab_filename(approx_vocab_size, strip):
return "vocab.wiki_revision%s.%d" % (".strip" if strip else "",
approx_vocab_size)
def get_or_generate_vocabulary(data_dir,
tmp_dir,
data_prefix,
max_page_size_exp,
approx_vocab_size=32768,
strip=True):
"""Get or generate the vocabulary.
Args:
data_dir: a string
tmp_dir: a string
data_prefix: a string
max_page_size_exp: an integer
approx_vocab_size: an integer
strip: a boolean
Returns:
a TextEncoder
"""
num_pages_for_vocab_generation = approx_vocab_size // 3
vocab_file = vocab_filename(approx_vocab_size, strip)
def my_generator(data_prefix):
"""Line generator for vocab."""
count = 0
for page in corpus_page_generator(
all_corpus_files(data_prefix)[::-1], tmp_dir, max_page_size_exp):
revisions = page["revisions"]
if revisions:
text = get_text(revisions[-1], strip=strip)
yield text
count += 1
if count % 100 == 0:
tf.logging.info("reading pages for vocab %d" % count)
if count > num_pages_for_vocab_generation:
break
return generator_utils.get_or_generate_vocab_inner(data_dir, vocab_file,
approx_vocab_size,
my_generator(data_prefix))
def get_encoder_from_vocab(vocab_filepath):
"""Get encoder from vocab file.
If vocab is not found in output dir, it will be copied there by
copy_vocab_to_output_dir to clarify the vocab used to generate the data.
Args:
vocab_filepath: path to vocab, either local or cns
Returns:
A SubwordTextEncoder vocabulary object. None if the output_parallel_text
is set.
"""
if not tf.gfile.Exists(vocab_filepath):
raise ValueError("Vocab file does not exist: {}.".format(vocab_filepath))
tf.logging.info("Found vocab file: %s", vocab_filepath)
encoder = text_encoder.SubwordTextEncoder(vocab_filepath)
return encoder
def throw_empty_pairs(src_tgt_pairs):
"""Filter [src,tgt] tuple from input list of pairs if either element is empty.
Args:
src_tgt_pairs: list of (src,tgt) pairs
Returns:
subset of input pair list for which all elements are non-empty
"""
return [x for x in src_tgt_pairs if x[0] and x[1]]
def edit_distance_filter(source_target_input, max_equal_to_diff_ratio=0):
"""Filter out examples that exceed max_edit_ratio between source and target.
Args:
source_target_input: a list of [source, target] pairs
max_equal_to_diff_ratio: cutoff for ratio of equal chars / diff chars
between source and target
Returns:
source_target_output: filtered subset of [source, target] input pairs
thrown_out_count: number of examples filtered out
"""
thrown_out_count = 0
source_target_output = []
if not max_equal_to_diff_ratio:
return source_target_input, thrown_out_count
for src_tgt in source_target_input:
opcodes = fast_match_sequences(*src_tgt)
diff_char_count = 0
equal_char_count = 0
for tag, i1, i2, j1, j2 in opcodes:
if tag == "diff":
# max() prevents double-counting substitutions.
diff_char_count += max(i2 - i1, j2 - j1)
else:
equal_char_count += i2 - i1
if diff_char_count <= max_equal_to_diff_ratio * equal_char_count:
source_target_output.append(src_tgt)
else:
thrown_out_count += 1
return source_target_output, thrown_out_count
def introduce_errors(s,
corruption_rate=3e-3,
infill_marker="|?|",
max_infill_len=8):
"""Artificially add spelling errors and infill markers.
This function should be applied to the inputs of a correction model.
The artificial errors are particularly useful to train a network to
correct spelling when the training data does not contain many
natural errors.
Also replaces some substrings with an "infill" marker. e.g.
"the fat cat sat on the mat" -> "the fat ca??? the mat"
This causes the trained model to learn infilling (predicting what text
to insert at the current cursor position).
Args:
s: a string (the uncorrupted text)
corruption_rate: a floating point value. Probability of introducing an
error/infill at each character.
infill_marker: a string
max_infill_len: an optional integer - maximum number of characters to remove
and replace by an infill marker. None means no infilling.
Returns:
a string
"""
num_errors = 0
ret = []
operations = [
"delete", # delete a character
"insert", # insert a random character from the input string
"replace", # replace a character with a random character from
# the input string
"transpose", # transpose two adjacent characters
]
if max_infill_len:
operations.append("infill")
pos = 0
while pos < len(s):
if random.random() >= corruption_rate:
ret.append(s[pos])
pos += 1
continue
num_errors += 1
operation = operations[random.randint(0, len(operations) - 1)]
if operation == "delete":
pos += 1
elif operation == "insert":
ret.append(s[random.randint(0, len(s) - 1)])
elif operation == "replace":
ret.append(s[random.randint(0, len(s) - 1)])
pos += 1
elif operation == "transpose":
ret.append(s[pos + 1] if pos + 1 < len(s) else "")
ret.append(s[pos])
pos += 2
else:
assert operation == "infill"
ret.append(infill_marker)
pos += random.randint(0, max_infill_len)
return "".join(ret), num_errors
def fast_match_sequences(a,
b,
a_start=0,
a_end=None,
b_start=0,
b_end=None,
min_match_length=3,
max_recursion_depth=128):
"""Compute diffs between two sequences.
This function is similar in functionality and spirit to
difflib.SequenceMatcher.get_opcodes, but it seems to run faster.
if a_start, a_end, b_start, b_end are specified, then we compute diffs of
the segments a[a_start:a_end] and b[b_start:b_end]. Returned indices
are relative to the full sequence.
We try to match the longest matching segments first, but due to heuristics
in finding the matches, this is not guaranteed.
Matching segments shorter than min_match_length are counted as part of the
surrounding differing segments, unless they are at the beginning or end of
both sequences. This helps eliminate junk matches.
Args:
a: a sequence
b: a sequence
a_start: an optional integer
a_end: an optional integer
b_start: an optional integer
b_end: an optional integer
min_match_length: an integer
max_recursion_depth: an integer - avoids crashes in weird corner cases
involving pairs of long repetitive sequences.
Returns:
a list of 5-tuples (tag, i1, i2, j1, j2).
Each tuple represents the alignment of segment a[i1:i2] with b[j1:j2].
tag is either "equal" or "diff". Note that the tags differ from those
returned by difflib.SequenceMatcher.get_opcodes.
"""
if a_end is None:
a_end = len(a)
if b_end is None:
b_end = len(b)
if a_start == a_end and b_start == b_end:
return []
if a_start == a_end or b_start == b_end:
return [("diff", a_start, a_end, b_start, b_end)]
# Compute an index from value to first occurrence in the b segment.
# Technically, we should index and explore all occurrences of a value,
# but that might be much slower.
b_index = {}
for j in range(b_end - 1, b_start - 1, -1):
b_index[b[j]] = j
# we will look for the longest match we can find.
max_match_length = 0
a_pos = a_start
while a_pos < a_end:
val = a[a_pos]
b_pos = b_index.get(val)
if b_pos is None:
a_pos += 1
continue
else:
a_match_start = a_pos
a_match_end = a_pos + 1
b_match_start = b_pos
b_match_end = b_pos + 1
while (a_match_start > a_start and b_match_start > b_start and
a[a_match_start - 1] == b[b_match_start - 1]):
a_match_start -= 1
b_match_start -= 1
while (a_match_end < a_end and b_match_end < b_end and
a[a_match_end] == b[b_match_end]):
a_match_end += 1
b_match_end += 1
# Compute the length of the matching segment. We prefer the longest.
match_length = a_match_end - a_match_start
# Extra credit for matching at the beginning or end of the sequence.
if a_match_start == 0 and b_match_start == 0:
match_length += min_match_length
if a_match_end == len(a) and b_match_end == len(b):
match_length += min_match_length
if match_length > max_match_length:
max_match_length = match_length
best_match = (a_match_start, a_match_end, b_match_start, b_match_end)
# advance a_pos to the end of this match to avoid wasting time
# rediscovering this match.
a_pos = a_match_end
if max_match_length < min_match_length or max_recursion_depth == 0:
return [("diff", a_start, a_end, b_start, b_end)]
a_match_start, a_match_end, b_match_start, b_match_end = best_match
return (fast_match_sequences(
a, b, a_start, a_match_start, b_start, b_match_start, min_match_length,
max_recursion_depth - 1) + [
("equal", a_match_start, a_match_end, b_match_start, b_match_end)
] + fast_match_sequences(a, b, a_match_end, a_end, b_match_end, b_end,
min_match_length, max_recursion_depth - 1))
================================================
FILE: tensor2tensor/data_generators/wikifact/README.md
================================================
# Assessing the Factual Accuracy of Generated Text
This directory will contain the code and scripts to generate data and train
models from the paper *Assessing the Factual Accuracy of Generated Text*.
================================================
FILE: tensor2tensor/data_generators/wikisum/README.md
================================================
# Generating Wikipedia by Summarizing Long Sequences
This directory contains the code and scripts to generate the dataset from the
paper [Generating Wikipedia by Summarizing Long
Sequences](https://arxiv.org/abs/1801.10198). The task is to generate a
Wikipedia article based on the contents of the cited references in that article
and the top 10 Google search results for the article's title.
There are 2 sources for the reference URLs used:
1. [CommonCrawl](http://commoncrawl.org/), an open-source crawl of the web. The
advantage of using CommonCrawl is that the dataset is perfectly reproducible.
However, there is limited coverage of the reference URLs.
1. Live web fetches. Coverage is considerably increased, but the content is
subject to change.
This document provides instructions for producing both datasets.
## Support files
Some files that are used in dataset generation have already been generated and
uploaded to Google Cloud Storage as `gs://tensor2tensor-data/wikisum`.
**URLs:** The dataset contains ~90M URLs total (~2.3M Wikipedia articles, each
with ~40 reference URLs). The URLs in the dataset are available in sharded JSON
files here: `gs://tensor2tensor-data/wikisum/wiki_urls/`.
**Wikipedia Articles:** We have processed the Wikipedia articles slightly to
extract the title, section breaks, and section headings. The processed Wikipedia
content is available in sharded `TFRecord` files containing serialized
`tensorflow.Example` protocol buffers here:
`gs://tensor2tensor-data/wikisum/wiki_content/`. The sharding is determined by a
hash of the Wikpedia article's title. The `Example`s contain features `[url,
title, section_titles, section_texts]`.
**CommonCrawl References Index:** To enable efficiently extracting the reference
URLs from CommonCrawl, we provide a JSON file per CommonCrawl file which maps a
reference URL contained in that CommonCrawl file to a list of shard ids:
`gs://tensor2tensor-data/wikisum/commoncrawl_metadata/`. These shards are the
ones that contain one or more Wikipedia articles that cite this reference. The
scripts in this directory will use this information to efficiently join the
reference with their Wikipedia articles.
*Note*: You can use [`gsutil`](https://cloud.google.com/storage/docs/gsutil) to
view the support files.
## Data generation
Data generation will first extract reference content (from either CommonCrawl or
the web), then generate a vocabulary, join the references with their Wikipedia
articles, run TF-IDF to rank reference paragraphs for a given article, and then
encode the references and the Wikipedia article with the vocabulary and write
the encoded training or evaluation example out to disk.
The output of data generation is a set of `TFRecord` files containing serialized
`tensorflow.Example` protocol buffers, with feature keys `"inputs"` and
`"targets"`. The inputs are the reference tokens, and the targets are the
Wikipedia article tokens.
In both cases, you must use multiple machines to extract references and produce
the final data to disk because of the size of the data. See `parallel_launch.py`
which is a script that will launch N machines in parallel on GCP. You can use it
as a guide if you'd like to launch on other infrastructure.
There are 3 jobs to run:
1. Extract references: `get_references_commoncrawl.py` for `WikisumCommoncrawl`
and `get_references_web.py` for `WikisumWeb`.
1. Build vocabulary (single-machine): `generate_vocab.py`
1. Produce Examples: `produce_examples.py`
With 1,000 machines with a good internet connection, data generation takes well
under 24 hours.
## Setup if using `parallel_launch.py` to launch on Google Cloud Platform
First, [install the `gcloud` CLI](https://cloud.google.com/sdk/downloads).
```
# Initialize the CLI
gcloud init
# Login
gcloud auth login
# Update the CLI
gcloud components update
# Set the default project and zone
gcloud config set core/project myproject
gcloud config set compute/zone us-central1-c
```
You'll also need to request the requisite
[quotas](https://console.cloud.google.com/iam-admin/quotas) in the zone you'll
be launching the machines in (whatever default zone you set above):
* In-use IP addresses: 1,000
* Internal IP addresses: 1,000
* Persistent Disk Standard (GB): 10,000
* CPUs: 4,000
**Running the commands below will launch instances on Google Cloud Platform and
you will incur charges.** If any of the commands go bad, immediately delete any
stranded instances. `delete_instances.sh` helps you delete instances in bulk
from the command-line, or you can delete many instances at once from the
[GCP Console](https://console.cloud.google.com/).
### Cost estimates
These are rough (and **not** guaranteed) estimates of cost if you were to launch
on GCP.
Pricing is taken from
[here](https://cloud.google.com/compute/pricing#custommachinetypepricing).
* `WikisumCommoncrawl`
* `get_references_commoncrawl`: $50 (1k machines, 1 CPU, 2G memory, 1 hour)
* `produce_examples`: $25 (1k machines, 1 CPU, 3G memory, 30 minutes)
* `WikisumWeb`
* `get_references_web`: $600 (1k machines, 4 CPU, 4G memory, 4 hours)
* `produce_examples`: $25 (1k machines, 1 CPU, 3G memory, 30 minutes)
## Commands to generate `WikisumCommoncrawl`
```
pip install tensor2tensor -U --user
# Set to your own GCS bucket
BUCKET=gs://my-gcs-bucket/wikisum_commoncrawl
# Extract references from CommonCrawl
python -m tensor2tensor.data_generators.wikisum.parallel_launch \
--num_instances=1000 \
--cpu=1 --mem=2 \
--name=wikisum-cc-refs \
--log_dir=$BUCKET/logs \
--setup_command="pip install tensor2tensor tensorflow -U -q --user" \
--command_prefix="python -m tensor2tensor.data_generators.wikisum.get_references_commoncrawl --num_tasks=1000 --out_dir=$BUCKET/wiki_references --task_id"
# Generate vocabulary file
python -m tensor2tensor.data_generators.wikisum.generate_vocab \
--out_dir=$BUCKET/data \
--refs_dir=$BUCKET/wiki_references \
--for_commoncrawl
# Produce examples
python -m tensor2tensor.data_generators.wikisum.parallel_launch \
--num_instances=1000 \
--cpu=1 --mem=3 \
--name=wikisum-cc-produce \
--log_dir=$BUCKET/logs \
--setup_command="pip install tensor2tensor tensorflow -U -q --user" \
--command_prefix="python -m tensor2tensor.data_generators.wikisum.produce_examples --out_dir=$BUCKET/data --refs_dir=$BUCKET/wiki_references --num_tasks=1000 --for_commoncrawl --task_id"
# Validate data
python -m tensor2tensor.data_generators.wikisum.validate_data \
--out_dir=$BUCKET/data \
--for_commoncrawl
```
## Commands to generate `WikisumWeb`
```
pip install tensor2tensor -U --user
# Set to your own GCS bucket
BUCKET=gs://my-gcs-bucket/wikisum_web
# Fetch references from web
python -m tensor2tensor.data_generators.wikisum.parallel_launch \
--num_instances=1000 \
--cpu=4 --mem=4 \
--name=wikisum-web-refs \
--log_dir=$BUCKET/logs \
--setup_command="pip3 install tensorflow tensor2tensor aiohttp cchardet aiodns bs4 -U -q --user" \
--command_prefix="python3 -m tensor2tensor.data_generators.wikisum.get_references_web --out_dir=$BUCKET/wiki_references --shard_id"
# Generate vocabulary file
python -m tensor2tensor.data_generators.wikisum.generate_vocab \
--out_dir=$BUCKET/data \
--refs_dir=$BUCKET/wiki_references
# Produce examples
python -m tensor2tensor.data_generators.wikisum.parallel_launch \
--num_instances=1000 \
--cpu=1 --mem=3 \
--name=wikisum-web-produce \
--log_dir=$BUCKET/logs \
--setup_command="pip install tensor2tensor tensorflow -U -q --user" \
--command_prefix="python -m tensor2tensor.data_generators.wikisum.produce_examples --out_dir=$BUCKET/data --refs_dir=$BUCKET/wiki_references --num_tasks=1000 --task_id"
# Validate data
python -m tensor2tensor.data_generators.wikisum.validate_data \
--out_dir=$BUCKET/data
```
## Training
**TODO(rsepassi)**: Put actual results achieved on `wikisum_web` and/or
`wikisum_commoncrawl` and with what `hparams_set`.
```
PROBLEM=wikisum_web # or wikisum_commoncrawl
t2t-trainer \
--problem=$PROBLEM \
--model=transformer \
--hparams_set=transformer_base \
--train_steps=250000 \
--eval_steps=100 \
--data_dir=$DATA_DIR \
--output_dir=$TRAIN_DIR
```
## Dataset Metadata
The following table is necessary for this dataset to be indexed by search
engines such as <a href="https://g.co/datasetsearch">Google Dataset Search</a>.
<div itemscope itemtype="http://schema.org/Dataset">
<table>
<tr>
<th>property</th>
<th>value</th>
</tr>
<tr>
<td>name</td>
<td><code itemprop="name">wikisum</code></td>
</tr>
<tr>
<td>alternateName</td>
<td><code itemprop="alternateName">WikisumCommonCrawl</code></td>
</tr>
<tr>
<td>alternateName</td>
<td><code itemprop="alternateName">WikisumWeb</code></td>
</tr>
<tr>
<td>alternateName</td>
<td><code itemprop="alternateName">wkisum_commoncrawl</code></td>
</tr>
<tr>
<td>alternateName</td>
<td><code itemprop="alternateName">wikisum_web</code></td>
</tr>
<tr>
<td>url</td>
<td><code itemprop="url">https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/wikisum</code></td>
</tr>
<tr>
<td>sameAs</td>
<td><code itemprop="sameAs">https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/data_generators/wikisum</code></td>
</tr>
<tr>
<td>description</td>
<td><code itemprop="description">The dataset from the
paper [Generating Wikipedia by Summarizing Long
Sequences](https://arxiv.org/abs/1801.10198). The task is to generate a
Wikipedia article based on the contents of the cited references in that article
and the top 10 Google search results for the article's title.\n
\n
There are 2 sources for the reference URLs used:
\n
1. [CommonCrawl](http://commoncrawl.org/), an open-source crawl of the web. The
advantage of using CommonCrawl is that the dataset is perfectly reproducible.
However, there is limited coverage of the reference URLs.
1. Live web fetches. Coverage is considerably increased, but the content is
subject to change.\n
\n
The dataset includes:\n
\n
**URLs:** The dataset contains ~90M URLs total (~2.3M Wikipedia articles, each
with ~40 reference URLs). The URLs in the dataset are available in sharded JSON
files.\n
\n
**Wikipedia Articles:** We have processed the Wikipedia articles slightly to
extract the title, section breaks, and section headings. The processed Wikipedia
content is available in sharded `TFRecord` files containing serialized
`tensorflow.Example` protocol buffers.\n
\n
**CommonCrawl References Index:** To enable efficiently extracting the reference
URLs from CommonCrawl, we provide a JSON file per CommonCrawl file which maps a
reference URL contained in that CommonCrawl file to a list of shard ids.
These shards are the ones that contain one or more Wikipedia articles that cite
this reference.</code></td>
</tr>
<tr>
<td>citation</td>
<td><code itemprop="citation">https://identifiers.org/arxiv:1801.10198</code></td>
</tr>
<tr>
<td>provider</td>
<td>
<div itemscope itemtype="http://schema.org/Organization" itemprop="provider">
<table>
<tr>
<th>property</th>
<th>value</th>
</tr>
<tr>
<td>name</td>
<td><code itemprop="name">Google</code></td>
</tr>
<tr>
<td>sameAs</td>
<td><code itemprop="sameAs">https://en.wikipedia.org/wiki/Google</code></td>
</tr>
</table>
</div>
</td>
</tr>
</table>
</div>
================================================
FILE: tensor2tensor/data_generators/wikisum/__init__.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
================================================
FILE: tensor2tensor/data_generators/wikisum/delete_instances.sh
================================================
#!/bin/bash
# Delete Google Compute Engine instances with naming structure $NAME-$INDEX
# (e.g. machines created with parallel_launch.py).
# Example usage:
# delete_instances.sh fetch-ref-urls 1000
NAME=$1
MAX=$2
MIN=${3:-0}
LOG_F=/tmp/delete-$NAME-logs.txt
echo "Deleting $MAX instances starting with $NAME-$MIN"
for i in $(seq $MIN $MAX)
do
gcloud compute instances delete --quiet $NAME-$i > $LOG_F 2>&1 &
if [[ $(( i % 100 )) == 0 ]]
then
# Give it some room to breathe every 100
sleep 30
fi
done
echo "Delete commands launched. Logs redirected to $LOG_F"
================================================
FILE: tensor2tensor/data_generators/wikisum/generate_vocab.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Generate vocab from references and wikis."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from tensor2tensor.data_generators.wikisum import wikisum
import tensorflow.compat.v1 as tf
flags = tf.flags
FLAGS = flags.FLAGS
flags.DEFINE_string("out_dir", None, "Directory to write vocab to.")
flags.DEFINE_string("wikis_dir",
"gs://tensor2tensor-data/wikisum/wiki_content/",
"Directory with wiki_content.tfrecords shards.")
flags.DEFINE_string("refs_dir", None,
"Directory with process_X folders with reference shards.")
flags.DEFINE_bool("for_commoncrawl", False,
"Whether to use WikisumCommoncrawl or WikisumWeb.")
def main(_):
if FLAGS.for_commoncrawl:
problem = wikisum.WikisumCommoncrawl()
else:
problem = wikisum.WikisumWeb()
problem.generate_vocab(FLAGS.out_dir, FLAGS.wikis_dir, FLAGS.refs_dir)
if __name__ == "__main__":
tf.logging.set_verbosity(tf.logging.INFO)
tf.app.run()
================================================
FILE: tensor2tensor/data_generators/wikisum/get_references_commoncrawl.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Extract references from CommonCrawl files."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import tempfile
from tensor2tensor.data_generators.wikisum import utils
from tensor2tensor.data_generators.wikisum import wikisum
import tensorflow.compat.v1 as tf
flags = tf.flags
FLAGS = flags.FLAGS
flags.DEFINE_integer("num_tasks", 1000, "Number of parallel tasks.")
flags.DEFINE_integer("task_id", 0, "Task id in a parallel run.")
flags.DEFINE_string("metadata_dir",
"gs://tensor2tensor-data/wikisum/commoncrawl_metadata/",
"Path to metadata files specifying what references are in "
"which CommonCrawl files.")
flags.DEFINE_string("out_dir", None, "Directory to write references to.")
flags.DEFINE_string("commoncrawl_wet_dir", None,
"Path to CommonCrawl wet.gz files locally. If not "
"provided, will download.")
def main(_):
assert FLAGS.out_dir
assert FLAGS.metadata_dir
out_dir = os.path.join(FLAGS.out_dir, "process_%d" % FLAGS.task_id)
tf.gfile.MakeDirs(out_dir)
with utils.timing("get_refs_commoncrawl"):
# Get all WET files
if FLAGS.commoncrawl_wet_dir:
wet_files = tf.gfile.Glob(
os.path.join(FLAGS.commoncrawl_wet_dir, "*.wet.gz"))
else:
tmp_dir = tempfile.gettempdir()
wet_files = list(
utils.wet_download_urls(utils.WET_PATHS_BY_DATE["0917"], tmp_dir))
# Shard and select this task's work
wet_files.sort()
wet_files = utils.shard(wet_files, FLAGS.num_tasks)[FLAGS.task_id]
tf.logging.info("Sharded out WET files. Processing %d files",
len(wet_files))
wikisum.extract_references_from_wets(wet_files, FLAGS.metadata_dir, out_dir)
if __name__ == "__main__":
tf.logging.set_verbosity(tf.logging.INFO)
tf.app.run()
================================================
FILE: tensor2tensor/data_generators/wikisum/get_references_web.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
# pylint: disable=line-too-long
r"""Fetch reference URLs from all groups for a single shard id.
Because of an SSL memory leak in Python 3.5, fetching too many URLs in the same
Python process will OOM. This script wraps get_references_web_single_group.py
and calls it through subprocess for each group in the shard, where each group is
~5k URLs.
Launch with parallel_launch.py
Each job should finish in ~5 hours with the settings below.
GCS_BUCKET=gs://my-bucket
python parallel_launch.py \
--num_instances=1000 \
--cpu=4 \
--mem=4 \
--name=get-refs-web \
--code_dir=./ \
--log_dir=$GCS_BUCKET/logs \
--setup_command="pip3 install aiohttp cchardet aiodns bs4 -q --user" \
--command_prefix="python3 wikisum/get_references_web.py --out_dir=$GCS_BUCKET/wiki_references --shard_id"
"""
# pylint: enable=line-too-long
import math
import os
import subprocess as sp
from tensor2tensor.data_generators.wikisum import get_references_web_single_group as fetch
from tensor2tensor.data_generators.wikisum import utils
import tensorflow.compat.v1 as tf
flags = tf.flags
FLAGS = flags.FLAGS
flags.DEFINE_string(
"command",
"python3 -m "
"tensor2tensor.data_generators.wikisum.get_references_web_single_group",
"Command to run get_references_web_single_group, without flags.")
def main(_):
shard_urls = fetch.get_urls_for_shard(FLAGS.urls_dir, FLAGS.shard_id)
num_groups = int(math.ceil(len(shard_urls) / fetch.URLS_PER_CLIENT))
tf.logging.info("Launching get_references_web_single_group sequentially for "
"%d groups in shard %d. Total URLs: %d",
num_groups, FLAGS.shard_id, len(shard_urls))
command_prefix = FLAGS.command.split() + [
"--urls_dir=%s" % FLAGS.urls_dir,
"--shard_id=%d" % FLAGS.shard_id,
"--debug_num_urls=%d" % FLAGS.debug_num_urls,
]
with utils.timing("all_groups_fetch"):
for i in range(num_groups):
command = list(command_prefix)
out_dir = os.path.join(FLAGS.out_dir, "process_%d" % i)
command.append("--out_dir=%s" % out_dir)
command.append("--group_id=%d" % i)
try:
# Even on 1 CPU, each group should finish within an hour.
sp.check_call(command, timeout=60*60)
except sp.TimeoutExpired:
tf.logging.error("Group %d timed out", i)
if __name__ == "__main__":
tf.logging.set_verbosity(tf.logging.INFO)
tf.app.run()
================================================
FILE: tensor2tensor/data_generators/wikisum/get_references_web_single_group.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Fetch reference URLs for a single group_id within a single shard_id.
See get_references_web.py to fetch URLs for all groups in within a single
shard_id.
Requires Python 3.5
pip3 install aiohttp cchardet aiodns bs4 tensorflow
"""
import datetime
import json
import math
import multiprocessing
import os
import random
import asyncio
import aiohttp
import tensorflow as tf
from tensor2tensor.data_generators.wikisum import html
from tensor2tensor.data_generators.wikisum import utils
flags = tf.flags
FLAGS = flags.FLAGS
flags.DEFINE_string("urls_dir", "gs://tensor2tensor-data/wikisum/wiki_urls/",
"Directory with wiki_urls.json files.")
flags.DEFINE_string("out_dir", None, "Directory to write reference files.")
flags.DEFINE_integer("max_parallel_requests", 50,
"Number of web requests to make in parallel.")
# Identify which URLs to fetch
flags.DEFINE_integer("shard_id", 0, "ID of URL shard to process.")
flags.DEFINE_integer("group_id", 0, "ID of group within the shard to process.")
flags.DEFINE_bool("log_samples", False,
"Whether to write out samples of the text extraction.")
flags.DEFINE_integer("log_every", 1000,
"How often to log and write out samples.")
flags.DEFINE_integer("debug_num_urls", 0,
"If >0, limits number of URLs fetched per input shard. "
"For debugging purposes only.")
WIKI_URLS_FILE = "wiki_urls.json-%05d-of-01000"
REF_SHARD_FILE = "references.tfrecords.gz-%05d-of-01000"
# Note that this program leaks memory, likely due to a bug in Python's SSL
# implementation that leaks sockets. This constant is used here and in
# get_references_web.py to limit the number of requests made by a single
# Python process. The more requests made, the more memory required due to the
# leak.
# TODO(rsepassi): Document memory impact of changing this.
URLS_PER_CLIENT = 5000
def concat_tfrecord_files(fnames, out_fname, rm_after=True):
with tf.gfile.Open(out_fname, "wb") as out_f:
for fname in fnames:
with tf.gfile.Open(fname, "rb") as in_f:
while True:
read = in_f.read(1000)
if not read:
break
out_f.write(read)
if rm_after:
tf.gfile.Remove(fname)
def shard(items, num_shards):
"""Split items into num_shards groups."""
sharded = []
num_per_shard = len(items) // num_shards
start = 0
for _ in range(num_shards):
sharded.append(items[start:start + num_per_shard])
start += num_per_shard
remainder = len(items) % num_shards
start = len(items) - remainder
for i in range(remainder):
sharded[i].append(items[start + i])
assert sum([len(fs) for fs in sharded]) == len(items)
return sharded
def mp_get_text(url, html):
return url, html.get_text_from_html(html)
def encode(s):
return bytes(s, "utf-8")
def make_example_from_ref(url, ref):
try:
url = encode(url)
ref = encode(ref)
except UnicodeEncodeError:
return None
features = {
"url":
tf.train.Feature(bytes_list=tf.train.BytesList(value=[url])),
"content":
tf.train.Feature(
bytes_list=tf.train.BytesList(value=[ref])),
}
return tf.train.Example(features=tf.train.Features(feature=features))
def tfrecord_fname(out_dir, shard_id, idx=None):
fname = os.path.join(out_dir, REF_SHARD_FILE % shard_id)
if idx is not None:
fname += ".%d" % idx
return fname
def make_tfrecord_writer(fname):
opts = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.GZIP)
return tf.python_io.TFRecordWriter(fname, opts)
def write_ref_content(url, ref, f):
if not ref:
return False
ex = make_example_from_ref(url, ref)
if ex is None:
return False
f.write(ex.SerializeToString())
return True
async def fetch_url(url, session, side_data):
text = None
try:
async with session.get(url, timeout=10, verify_ssl=False) as response:
if response.status == 200:
text = await response.text()
else:
tf.logging.error("Status %d, url: %s", response.status, url)
except:
# Request can fail for many reasons.
pass
return text, side_data
async def throttled_fetch_url(url, sem, session, side_data):
async with sem:
return await fetch_url(url, session, side_data)
async def fetch_urls(urls,
out_fname,
logging_fnames=None):
tasks = []
connector = aiohttp.TCPConnector(limit_per_host=1)
async with aiohttp.ClientSession(
connector=connector, cookie_jar=aiohttp.DummyCookieJar()) as session:
# Async fetch the urls
sem = asyncio.Semaphore(FLAGS.max_parallel_requests)
for url in urls:
side_data = {"url": url}
task = asyncio.ensure_future(
throttled_fetch_url(url, sem, session, side_data))
tasks.append(task)
tf.logging.info("Async requested %d urls", len(urls))
# Setup output files
file_handles = []
out_f = make_tfrecord_writer(out_fname)
file_handles.append(out_f)
logging_fnames = logging_fnames or {}
samples_f = None
if "samples" in logging_fnames:
samples_f = tf.gfile.Open(logging_fnames["samples"], "w")
file_handles.append(samples_f)
refs_written = [0] # Made a list so can be mutated
def text_extraction_callback(callback_arg):
url, text = callback_arg
written = write_ref_content(url, text, out_f)
if not written:
return
if not refs_written[0] % FLAGS.log_every:
timestamp = datetime.datetime.now().strftime("%H:%M")
tf.logging.info("%s: Wrote ref %d in group", timestamp, refs_written[0])
if samples_f is not None:
samples_f.write(url)
samples_f.write("\n")
samples_f.write(text)
samples_f.write("\n\n---\n\n")
refs_written[0] += 1
try:
# Process each URL as it comes in.
# Using a multiprocessing Pool because the text extraction is expensive
# and so we distribute across cores.
pool = multiprocessing.Pool()
results = []
for task in asyncio.as_completed(tasks):
html, side_data = await task
url = side_data["url"]
if not html:
continue
res = pool.apply_async(mp_get_text, (url, html), {},
text_extraction_callback)
results.append(res)
for res in results:
try:
res.get(timeout=10)
except multiprocessing.TimeoutError:
pass
finally:
for f in file_handles:
f.close()
return refs_written[0]
def get_urls_per_shard(urls_files):
total_urls = 0
per_shard = {}
for urls_file in urls_files:
ref_urls = set()
shard_id = int(os.path.basename(urls_file)[15:20])
with tf.gfile.Open(urls_file) as f:
wiki_urls = json.loads(f.read())
for _, wiki_info in wiki_urls.items():
ref_urls |= set(wiki_info["refs"])
per_shard[shard_id] = list(ref_urls)
total_urls += len(ref_urls)
return per_shard, total_urls
def get_urls_for_shard(urls_dir, shard_id):
urls_file = os.path.join(urls_dir, WIKI_URLS_FILE % shard_id)
urls_per_shard, _ = get_urls_per_shard([urls_file])
assert len(urls_per_shard) == 1
return urls_per_shard[shard_id]
def get_urls_for_shard_group(urls_dir, shard_id, group_id):
shard_urls = get_urls_for_shard(urls_dir, shard_id)
# Deterministic sort and shuffle to prepare for sharding
shard_urls.sort()
random.seed(123)
random.shuffle(shard_urls)
groups = shard(shard_urls, int(math.ceil(len(shard_urls) / URLS_PER_CLIENT)))
group_urls = groups[group_id]
if FLAGS.debug_num_urls:
group_urls = group_urls[:FLAGS.debug_num_urls]
return group_urls
def main(_):
urls = get_urls_for_shard_group(
FLAGS.urls_dir, FLAGS.shard_id, FLAGS.group_id)
tf.logging.info("Fetching %d URLs for shard %d, group %d",
len(urls), FLAGS.shard_id, FLAGS.group_id)
tf.gfile.MakeDirs(FLAGS.out_dir)
out_fname = tfrecord_fname(FLAGS.out_dir, FLAGS.shard_id)
with utils.timing("group_fetch"):
logging_fnames = {}
if FLAGS.log_samples:
logging_fnames["samples"] = os.path.join(
FLAGS.out_dir, "samples.%d.txt" % FLAGS.shard_id)
loop = asyncio.get_event_loop()
num_written = loop.run_until_complete(asyncio.ensure_future(
fetch_urls(urls,
out_fname,
logging_fnames)))
tf.logging.info("Total URLs: %d", len(urls))
tf.logging.info("Num written: %d", num_written)
tf.logging.info("Coverage: %.1f", (num_written / len(urls)) * 100)
if __name__ == "__main__":
tf.logging.set_verbosity(tf.logging.INFO)
tf.app.run()
================================================
FILE: tensor2tensor/data_generators/wikisum/html.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Utils to parse HTML content into plaintext."""
import bs4
def get_text_from_html(html):
"""Returns a plaintext representation of HTML content."""
try:
soup = bs4.BeautifulSoup(html, "html.parser")
except: # pylint: disable=bare-except
# Some docs don't parse
return ""
# Remove script and style tags
for s in soup(["script", "style"]):
s.decompose()
return "\n".join([s for s in _soup_strings(soup)])
def _soup_strings(soup):
"""Return text strings in soup."""
paragraph_tags = set([
"caption", "details", "h1", "h2", "h3", "h4", "h5", "h6", "li", "p", "td",
"div", "span"
])
skip_children = None
for descendant in soup.descendants:
# If we've treated a tag as a contiguous paragraph, don't re-emit the
# children (see below).
if skip_children is not None:
try:
in_skip = descendant in skip_children # pylint: disable=unsupported-membership-test
except RecursionError: # pylint: disable=undefined-variable
# Possible for this check to hit a nasty infinite recursion because of
# BeautifulSoup __eq__ checks.
in_skip = True
if in_skip:
continue
else:
skip_children = None
# Treat some tags as contiguous paragraphs, regardless of other tags nested
# inside (like <a> or <b>).
if isinstance(descendant, bs4.Tag):
if descendant.name in paragraph_tags:
if descendant.find_all(paragraph_tags):
# If there are nested paragraph tags, don't treat it as a single
# contiguous tag.
continue
skip_children = list(descendant.descendants)
text = " ".join(descendant.get_text(" ", strip=True).split())
if text:
yield text
continue
if (isinstance(descendant, bs4.Comment) or
not isinstance(descendant, bs4.NavigableString)):
continue
text = " ".join(descendant.strip().split())
if text:
yield text
================================================
FILE: tensor2tensor/data_generators/wikisum/parallel_launch.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
# pylint: disable=line-too-long
r"""Launch a script in parallel on GCP.
For each instance (`--num_instances`), the script will copy the code in
`--code_dir` to the instance, run `--setup_command` and then run
`--command_prefix` joined with the task's id or a line in
`--per_instance_suffix_file`.
Note that the machines will attempt to down themselves on completion or failure.
If they do not, you can delete them manually or use delete_instances.sh to
delete many at once.
Example usage:
```
BUCKET=gs://my-bucket
python parallel_launch.py \
--num_instances=1000 \
--cpu=4 --mem=4 \
--name=wikisum-refs-web \
--code_dir=./ \
--log_dir=$BUCKET/refs_logs \
--setup_command="pip3 install aiohttp cchardet aiodns bs4 -q --user" \
--command_prefix="python3 wikisum/get_references_web.py --out_dir=$BUCKET/wiki_references --shard_id"
```
"""
# pylint: enable=line-too-long
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import contextlib
import multiprocessing as mp
import os
import socket
import subprocess as sp
import time
from tensor2tensor.utils import cloud_mlengine as cloud
import tensorflow.compat.v1 as tf
flags = tf.flags
FLAGS = flags.FLAGS
flags.DEFINE_integer("num_instances", None, "Number of instances to launch.")
flags.DEFINE_string("name", None, "Instance name prefix.")
flags.DEFINE_string("log_dir", None, "GCS bucket to copy logs out to.")
flags.DEFINE_string("code_dir", None, "Directory to copy.")
flags.DEFINE_string("setup_command", None, "Setup command to run.")
flags.DEFINE_string("command_prefix", None, "Command to run, prefix.")
flags.DEFINE_string("per_instance_suffix_file", None,
"Command to run, suffix per instance. If None, suffix will "
"be instance id.")
flags.DEFINE_integer("cpu", 1, "Number of CPUs per instance.")
flags.DEFINE_integer("mem", 4, "Memory in GB per instance.")
flags.DEFINE_integer("num_threads", 48,
"Number of threads to use to spin up jobs.")
flags.DEFINE_bool("debug_keep_up", False,
"If True, will keep the machine up. num_instances must be 1.")
flags.DEFINE_string("instance_ids", None,
"Comma-separated list of integer instance ids to launch. "
"Useful if some failed on a previous run and you only want "
"to rerun specific tasks.")
DELETE = "gcloud compute instances delete {name}"
DELETE_SELF = ("gcloud compute instances delete $(hostname) --quiet "
"--zone={zone}")
CREATE_INSTANCE = ("gcloud compute instances create {instance_name} "
"--custom-cpu {cpu} --custom-memory {mem} "
"--custom-extensions "
"--image-project=ml-images --image-family=tf-1-7 "
"--scopes=cloud-platform")
COPY_CODE = "gcloud compute scp --recurse {local_dir} {instance_name}:~/"
SSH = "gcloud compute ssh {instance_name} --command"
SCREEN = "screen -dmS test bash -c \"{command}\""
DEFAULT_ZONE = "gcloud config get-value compute/zone"
LOGS = "> ~/logs-{task_id}.txt 2>&1; gsutil cp ~/logs-{task_id}.txt {bucket}"
def remote_run(cmd, instance_name, detach=False, retries=1):
"""Run command on GCS instance, optionally detached."""
if detach:
cmd = SCREEN.format(command=cmd)
args = SSH.format(instance_name=instance_name).split()
args.append(cmd)
for i in range(retries + 1):
try:
if i > 0:
tf.logging.info("Retry %d for %s", i, args)
return sp.check_call(args)
except sp.CalledProcessError as e:
if i == retries:
raise e
def default_zone():
return cloud.shell_output(DEFAULT_ZONE).strip()
@contextlib.contextmanager
def safe_socket(timeout=2):
s = socket.socket()
s.settimeout(timeout)
try:
yield s
finally:
s.close()
def wait_for_ssh(ip):
"""Wait for SSH to be available at given IP address."""
for _ in range(12):
with safe_socket() as s:
try:
s.connect((ip, 22))
return True
except socket.timeout:
pass
time.sleep(10)
return False
def create_instance(instance_name, cpu=1, mem=4):
tf.logging.info("Creating instance %s", instance_name)
out = cloud.shell_output(CREATE_INSTANCE, instance_name=instance_name,
cpu=cpu, mem=mem)
return out.split("\n")[1:-1][0].split()[8]
def list_vm_names_and_ips():
list_out = cloud.shell_output(cloud.LIST_VM)
lines = [l.split() for l in list_out.split("\n")[1:-1]]
names_and_ips = [(l[0].strip(), l[-2].strip()) for l in lines]
return names_and_ips
def shell_run_with_retry(cmd, retries=1, **kwargs):
for i in range(retries + 1):
try:
if i > 0:
tf.logging.info("Retry %d for %s", i, cmd)
cloud.shell_run(cmd, **kwargs)
return
except sp.CalledProcessError as e:
if i == retries:
raise e
def delete_instance(instance_name):
cloud.shell_run(DELETE, name=instance_name)
def launch_instance(instance_name,
command,
existing_ip=None,
cpu=1,
mem=4,
code_dir=None,
setup_command=None):
"""Launch a GCE instance."""
# Create instance
ip = existing_ip or create_instance(instance_name, cpu=cpu, mem=mem)
tf.logging.info("Waiting for SSH %s", instance_name)
ready = wait_for_ssh(ip)
if not ready:
raise ValueError("Instance %s never ready for SSH" % instance_name)
# Copy code
if code_dir:
shell_run_with_retry(COPY_CODE, retries=2,
local_dir=code_dir, instance_name=instance_name)
# Run setup
if setup_command:
tf.logging.info("Running setup on %s", instance_name)
remote_run(setup_command, instance_name)
# Run command
tf.logging.info("Running command on %s", instance_name)
remote_run(command, instance_name, detach=True)
def main(_):
assert FLAGS.num_instances
assert FLAGS.name
zone = default_zone()
assert zone
code_dir = None
if FLAGS.code_dir:
code_dir = os.path.abspath(os.path.expanduser(FLAGS.code_dir))
# Suffixes per instance
if FLAGS.per_instance_suffix_file:
with tf.gfile.Open(FLAGS.per_instance_suffix_file) as f:
suffixes = [l.strip() for l in f.readlines()]
else:
suffixes = list(range(FLAGS.num_instances))
assert len(suffixes) == FLAGS.num_instances
vm_info = list_vm_names_and_ips()
vm_names = list(zip(*vm_info))[0] if vm_info else []
pool = mp.Pool(FLAGS.num_threads)
async_results = []
assert FLAGS.log_dir
log_dir = os.path.join(FLAGS.log_dir, FLAGS.name)
tf.gfile.MakeDirs(log_dir)
assert log_dir.startswith("gs://")
if not log_dir.endswith("/"):
log_dir += "/"
# Write a test file to make sure gcloud GCS APIs are enabled
test_filename = os.path.join(log_dir, "check_write")
with tf.gfile.Open(test_filename, "w") as f:
f.write("testing GCS write")
tf.gfile.Remove(test_filename)
instance_ids = list(range(FLAGS.num_instances))
if FLAGS.instance_ids:
instance_ids = [int(i) for i in FLAGS.instance_ids.split(",")]
tf.logging.info("Launching %d instances", len(instance_ids))
for i in instance_ids:
instance_name = "%s-%d" % (FLAGS.name, i)
existing_ip = (vm_info[vm_names.index(instance_name)][1]
if instance_name in vm_names else None)
logging = LOGS.format(task_id=i, bucket=log_dir) if log_dir else ""
delete = DELETE_SELF.format(zone=zone)
if FLAGS.debug_keep_up:
assert len(instance_ids) == 1
delete = ""
command = "{prefix} {suffix} {logging}; {delete}".format(
prefix=FLAGS.command_prefix,
suffix=suffixes[i],
delete=delete,
logging=logging)
args = (instance_name, command, existing_ip,
FLAGS.cpu, FLAGS.mem, code_dir,
FLAGS.setup_command)
res = pool.apply_async(launch_instance, args)
async_results.append((res, instance_name, i))
failed = []
for res, instance_name, i in async_results:
try:
res.get()
except Exception as e: # pylint: disable=broad-except
failed.append((instance_name, i))
tf.logging.error("Failed to launch task %s due to exception %s",
instance_name, str(e))
results = []
if failed:
ids_for_flag = ",".join([str(i) for i in list(zip(*failed))[1]])
tf.logging.error("Failed to launch %d jobs. Tasks: %s. "
"Attempting delete in case they are still up. Rerun with "
"--instance_ids='%s' to attempt relaunch.",
len(failed), str(failed), ids_for_flag)
for instance_name, _ in failed:
res = pool.apply_async(delete_instance, (instance_name,))
results.append(res)
for res in results:
try:
res.get()
except: # pylint: disable=bare-except
pass
tf.logging.info("Launching complete.")
if __name__ == "__main__":
tf.logging.set_verbosity(tf.logging.INFO)
tf.app.run()
================================================
FILE: tensor2tensor/data_generators/wikisum/produce_examples.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Produce examples given a vocab, wikis, references, and dataset URLs."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
from six.moves import range
from tensor2tensor.data_generators.wikisum import utils
from tensor2tensor.data_generators.wikisum import wikisum
import tensorflow.compat.v1 as tf
flags = tf.flags
FLAGS = flags.FLAGS
flags.DEFINE_integer("num_tasks", 1000, "Number of parallel tasks.")
flags.DEFINE_integer("task_id", 0, "Task id in a parallel run.")
flags.DEFINE_string("out_dir", None, "Directory to write to.")
flags.DEFINE_string("wikis_dir",
"gs://tensor2tensor-data/wikisum/wiki_content/",
"Directory with wiki_content.tfrecords.")
flags.DEFINE_string("refs_dir", None, "Directory with process_X dirs")
flags.DEFINE_string("urls_dir", "gs://tensor2tensor-data/wikisum/wiki_urls/",
"Directory with wiki_urls.json")
flags.DEFINE_string("vocab_dir", None, "Directory with vocab file")
flags.DEFINE_bool("for_commoncrawl", False,
"Whether to use WikisumCommoncrawl or WikisumWeb.")
def main(_):
if FLAGS.for_commoncrawl:
problem = wikisum.WikisumCommoncrawl()
else:
problem = wikisum.WikisumWeb()
out_filepaths = problem.out_filepaths(FLAGS.out_dir)
out_filepaths = utils.shard(out_filepaths, FLAGS.num_tasks)[FLAGS.task_id]
if not FLAGS.vocab_dir:
FLAGS.vocab_dir = FLAGS.out_dir
shard_ids = utils.shard(list(range(utils.NUM_SHARDS)),
FLAGS.num_tasks)[FLAGS.task_id]
with utils.timing("produce_examples"):
wikisum.produce_examples(
shard_ids=shard_ids,
wikis_dir=FLAGS.wikis_dir,
refs_dir=FLAGS.refs_dir,
urls_dir=FLAGS.urls_dir,
vocab_path=os.path.join(FLAGS.vocab_dir, problem.vocab_filename),
out_filepaths=out_filepaths)
if __name__ == "__main__":
tf.logging.set_verbosity(tf.logging.INFO)
tf.app.run()
================================================
FILE: tensor2tensor/data_generators/wikisum/test_data/para_bad1.txt
================================================
kolkata ward no 97 37
you are here : india » west bengal » kolkata » kolkata
this paragraph too short
a | b | c | d | e | f | g | h | i | j | k | l | m | n | o | p | q | r | s | t | u | v | w | x | y | z
123 123 123 123 985 9880 1230 0980 . 12398 .
- 5 . 7 % - 5 . 2 % - 15 . 1 % 4 . 7 % - 13 . 3 %
http : / / www . bbc . co . uk / sport / football / 24351521
no . - 26 beadon street .
{ { / playpopup } } { { ^ playpopup } } { { # playinvideopage } } { { / playinvideopage } } { { ^ playinvideopage } } { { / playinvideopage } } { { / playpopup } } <p> { { # playpopup } } { { / playpopup } } { { ^ playpopup } } { { # playinvideopage } } { { / playinvideopage } } { { ^ playinvideopage } } { { / playinvideopage } } { { / playpopup } } { { genre } }
denham , samuel coulter , sally 133 oct 28 1819
browse by
================================================
FILE: tensor2tensor/data_generators/wikisum/test_data/para_good1.txt
================================================
this is a very good paragraph . it even has two sentences .
the castle that was soon to figure so largely in lee’s life lay fourteen miles
to the southwest of where he sat perched atop his tank . topped with storybook
crenelations and accompanied by a rich history , schloss itter , as it’s called
in german , was first mentioned in land records as early as 1240 . since then ,
itter has passed through a number of hands . after germany’s march 1938
annexation of austria , the castle’s robust construction and relatively remote
location attracted the attention of the notoriously secretive nazis . within
months of absorbing austria into the greater reich , the german government
requisitioned castle itter for unspecified “official use”—which included housing
for several months in 1942 an organization called the “german association for
combating the dangers of tobacco . ” on february 7 , 1943 , it fell into new
hands yet again , for on that day , the structure and all its outbuildings were
requisitioned by the wehrmacht on behalf of the ss .
the url for the site is http : / / www . bbc . co . uk / sport / football / 24351521 .
================================================
FILE: tensor2tensor/data_generators/wikisum/utils.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Wikisum data generation utilities."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import collections
import contextlib
import datetime
import gzip
import os
import re
import urllib
import tensorflow.compat.v1 as tf
# pylint: disable=g-import-not-at-top
# To maintain compatibility with Python 2 and 3
try:
import cStringIO as StringIO
except ImportError:
import io as StringIO
# pylint: enable=g-import-not-at-top
# Each entry is a URL to the wet.paths.gz file for that CommonCrawl dump.
WET_PATHS_BY_DATE = {
'0917': ('https://commoncrawl.s3.amazonaws.com/crawl-data/CC-MAIN-2017-39/'
'wet.paths.gz'),
}
S3_HTTP_PREFIX = 'https://commoncrawl.s3.amazonaws.com/'
NUM_SHARDS = 1000
METADTA_SUFFIX = '.metadata.json'
def readahead(path):
return path
class WETHeader(collections.namedtuple('WETHeader', ['url', 'length'])):
URI_HEADER = 'WARC-Target-URI: '
LENGTH_HEADER = 'Content-Length: '
@classmethod
def read(cls, f):
"""Read header from file. Headers end with length and then 1 blank line."""
url = None
line = f.readline()
if not line:
# EOF
return None
while not line.startswith(cls.LENGTH_HEADER):
if line.startswith(cls.URI_HEADER):
url = line[len(cls.URI_HEADER):].strip()
line = f.readline()
# Consume empty separator
f.readline()
# Read content
length = int(line.split(':')[1])
return cls(url, length)
class WETRecord(collections.namedtuple('WETRecord', ['url', 'content'])):
@classmethod
def read(cls, f):
"""Read WETRecord from file. Records end with 2 blank lines."""
header = WETHeader.read(f)
if header is None:
# EOF
return None
content = f.read(header.length)
# Consume empty separators
f.readline()
f.readline()
return cls(header.url, content)
def wet_records_from_file_obj(f, take_ownership=False):
"""Iterate through records in WET file object."""
while True:
record = WETRecord.read(f)
if record is None:
break
if not record.url:
continue
yield record
if take_ownership:
f.close()
def wet_records(wet_filepath):
"""Generate WETRecords from filepath."""
if wet_filepath.endswith('.gz'):
fopen = gzip.open
else:
fopen = tf.gfile.GFile
with fopen(wet_filepath) as f:
for record in wet_records_from_file_obj(f):
yield record
def download(url, download_dir):
outname = os.path.join(download_dir, os.path.basename(url))
if tf.gfile.Exists(outname):
print('Found %s, skipping download' % outname)
return outname
inprogress = outname + '.incomplete'
print('Downloading %s' % url)
inprogress, _ = urllib.urlretrieve(url, inprogress)
tf.gfile.Rename(inprogress, outname)
return outname
def wet_download_urls(wet_paths_url, tmp_dir, rm_after=True):
paths_gz = download(wet_paths_url, tmp_dir)
with gzip.open(paths_gz) as f:
path = f.readline()
while path:
download_path = S3_HTTP_PREFIX + path[:-1]
yield download_path
path = f.readline()
if rm_after:
tf.gfile.Remove(paths_gz)
def wet_records_from_url(download_url, tmp_dir, rm_after=True):
wet_gz = download(download_url, tmp_dir)
try:
for wet_record in wet_records(wet_gz):
yield wet_record
finally:
if rm_after:
tf.gfile.Remove(wet_gz)
class DummyPool(object):
def __init__(self, processes=None):
pass
def apply_async(self, fn, args=None):
args = args or tuple()
return DummyResult(fn(*args))
def map(self, fn, arg_list):
return [fn(a) for a in arg_list]
class DummyResult(object):
def __init__(self, result):
self.result = result
def get(self):
return self.result
def shard(items, num_shards):
"""Split items into num_shards groups."""
sharded = []
num_per_shard = len(items) // num_shards
start = 0
for _ in range(num_shards):
sharded.append(items[start:start + num_per_shard])
start += num_per_shard
remainder = len(items) % num_shards
start = len(items) - remainder
for i in range(remainder):
sharded[i].append(items[start + i])
assert sum([len(fs) for fs in sharded]) == len(items)
return sharded
def gzip_memfile(fname):
with tf.gfile.Open(readahead(fname)) as f:
memfile = StringIO.StringIO(f.read())
return gzip.GzipFile(fileobj=memfile)
_SOME_ALPHA_RE = re.compile(r'[A-Za-z]+')
_ONLY_ALPHA_RE = re.compile(r'^[A-Za-z]*$')
def filter_paragraph(p):
"""Simple filter to remove obviously bad paragraphs (bad text extraction).
Note this needs to run very quickly as it is applied to every paragraph
in the corpus, so nothing fancy! This whole method should be linear
expected time in len(p).
Args:
p: string, paragraph
Returns:
True if we should remove the paragraph.
"""
# Expect a minimum number of words.
tokens = p.split()
if len(tokens) < 6:
return True
# Require some letters.
if not re.search(_SOME_ALPHA_RE, p):
return True
# Keep this one at the end, probably the most complicated logic.
# We try to detect sentences, which should have a minimum of 3 tokens
# with only alphabetic characters.
last = 0
found_sentence = False
num_alpha = 0
for i, x in enumerate(tokens):
if x == '.':
if i - last > 3 and num_alpha >= 3:
found_sentence = True
break
last = i
num_alpha = 0
if re.match(_ONLY_ALPHA_RE, x):
num_alpha += 1
if not found_sentence:
return True
return False
@contextlib.contextmanager
def timing(name=''):
"""Log start, end, and duration."""
start = datetime.datetime.now()
timestamp = start.strftime('%H:%M')
tf.logging.info('Starting job [%s] at %s', name, timestamp)
yield
end = datetime.datetime.now()
timestamp = end.strftime('%H:%M')
tf.logging.info('Finished job [%s] at %s', name, timestamp)
duration = end - start
duration_mins = duration.total_seconds() / 60
tf.logging.info('Total time [%s] (m): %d', name, int(duration_mins))
================================================
FILE: tensor2tensor/data_generators/wikisum/utils_test.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Tests for tensor2tensor.data_generators.wikisum.utils."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
from tensor2tensor.data_generators.wikisum import utils
import tensorflow.compat.v1 as tf
pkg_dir = os.path.abspath(__file__)
pkg_dir, _ = os.path.split(pkg_dir)
_TESTDATA = os.path.join(pkg_dir, "test_data")
def _get_testdata(filename):
with tf.io.gfile.GFile(filename) as f:
return f.read()
class UtilsTest(tf.test.TestCase):
def test_filter_paragraph(self):
for bad in tf.io.gfile.glob(os.path.join(_TESTDATA, "para_bad*.txt")):
for p in _get_testdata(bad).split("\n"):
self.assertTrue(utils.filter_paragraph(p),
msg="Didn't filter %s" % p)
for good in tf.io.gfile.glob(os.path.join(_TESTDATA, "para_good*.txt")):
for p in _get_testdata(good).split("\n"):
p = _get_testdata(good)
self.assertFalse(utils.filter_paragraph(p), msg="Filtered %s" % p)
if __name__ == "__main__":
tf.test.main()
================================================
FILE: tensor2tensor/data_generators/wikisum/validate_data.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Aggregate stats from produce_examples."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import json
import os
import numpy as np
import six
from six.moves import zip
from tensor2tensor.data_generators.wikisum import wikisum
import tensorflow.compat.v1 as tf
flags = tf.flags
FLAGS = flags.FLAGS
flags.DEFINE_string("out_dir", None, "Directory with data and stats files.")
flags.DEFINE_bool("for_commoncrawl", False,
"Whether to use WikisumCommoncrawl or WikisumWeb.")
flags.DEFINE_bool("rm_per_shard_stats", True,
"Whether to remove the per-shard stats files after writing "
"out the aggregated stats.")
def aggregate_stats(stats_files):
"""Aggregate stats in per-shard stats files."""
all_stats = {}
for fname in stats_files:
with tf.gfile.Open(fname) as f:
stats = json.loads(f.read())
for k, v in six.iteritems(stats):
if k not in all_stats:
if isinstance(v, list):
all_stats[k] = []
else:
all_stats[k] = 0
if isinstance(v, list):
all_stats[k].extend(v)
else:
all_stats[k] += v
stats = all_stats
ref_coverage = float(stats["total_found_refs"]) / stats["total_original_refs"]
len_bounds = [0, 2, 10, 100, 1000, 5000, 10000, 20000, 50000, 100000, 1000000]
len_counts, len_bounds = np.histogram(stats["ref_lengths"], len_bounds)
len_dist = len_counts.astype(np.float32) / len_counts.sum()
wiki_coverage = (float(stats["num_wikis_written"]) /
stats["total_original_wikis"])
wikis_skipped_no_ref = (float(stats["wikis_skipped_no_refs"]) /
stats["total_original_wikis"])
wikis_skipped_no_lead = (float(stats["wikis_skipped_short_lead"]) /
stats["total_original_wikis"])
wiki_ref_coverage = [
float(found) / orig for found, orig
in zip(stats["wiki_found_refs"], stats["wiki_original_refs"]) if found
]
coverage_bounds = np.arange(21).astype(np.float32) / 20
coverage_counts, coverage_bounds = np.histogram(wiki_ref_coverage,
coverage_bounds)
coverage_dist = coverage_counts.astype(np.float32) / coverage_counts.sum()
agg_stats = dict(
total_original_wikis=stats["total_original_wikis"],
total_original_refs=stats["total_original_refs"],
wiki_coverage=wiki_coverage,
wikis_skipped_no_ref=wikis_skipped_no_ref,
wikis_skipped_no_lead=wikis_skipped_no_lead,
overall_ref_coverage=ref_coverage,
per_wiki_ref_coverage_dist=list((coverage_dist * 100).astype(int)),
per_wiki_ref_coverage_bounds=list((coverage_bounds * 100).astype(int)),
ref_len_dist=list((len_dist * 100).astype(int)),
ref_len_bounds=list(len_bounds),
)
return agg_stats
def filename_to_task_id(fname):
"""Map filename to the task id that created it assuming 1k tasks."""
# This matches the order and size in WikisumBase.out_filepaths
fname = os.path.basename(fname)
shard_id_increment = {
"train": 0,
"dev": 800,
"test": 900,
}
parts = fname.split("-")
split = parts[1]
shard_id = parts[2]
task_id = int(shard_id) + shard_id_increment[split]
return task_id
def get_length(fname):
return tf.gfile.Stat(fname).length
def validate_data_files(problem, data_files, min_size):
"""Validate presence and minimum size of files."""
# Check that all files are present
data_dir = os.path.split(data_files[0])[0]
out_filepaths = problem.out_filepaths(data_dir)
missing_filepaths = set(out_filepaths) - set(data_files)
if missing_filepaths:
tf.logging.error("Missing %d data files", len(missing_filepaths))
# Check that each file is at least 100M
too_small = []
for data_file in data_files:
length = get_length(data_file)
if length < min_size:
too_small.append(data_file)
if too_small:
tf.logging.error("%d files too small", len(too_small))
bad_files = too_small + list(missing_filepaths)
return bad_files
def main(_):
if FLAGS.for_commoncrawl:
problem = wikisum.WikisumCommoncrawl()
else:
problem = wikisum.WikisumWeb()
prefix = problem.dataset_filename()
data_files = tf.gfile.Glob(os.path.join(FLAGS.out_dir, "%s*" % prefix))
missing_files = validate_data_files(
problem, data_files,
min_size=(60 if FLAGS.for_commoncrawl else 120) * 1e6)
task_ids = [filename_to_task_id(fname) for fname in missing_files]
ids_for_flag = ",".join([str(i) for i in task_ids])
tf.logging.error("You should (re)generate %d of the data files. "
"Rerun produce_examples with --instance_ids='%s'.",
len(missing_files), ids_for_flag)
# Compute and write out aggregated stats
stats_files = tf.gfile.Glob(os.path.join(FLAGS.out_dir, "stats*"))
agg_stats = aggregate_stats(stats_files)
if not FLAGS.for_commoncrawl:
coverage = agg_stats["overall_ref_coverage"] * 100
if not coverage > 80:
tf.logging.error("Overall reference coverage is expected to be > 80%. "
"It is %0.1f. You may want to rerun get_references_web.",
coverage)
with tf.gfile.Open(
os.path.join(FLAGS.out_dir, "stats.json"), "w") as f:
f.write(json.dumps(agg_stats))
if FLAGS.rm_per_shard_stats and not missing_files:
for fname in stats_files:
tf.gfile.Remove(fname)
if __name__ == "__main__":
tf.app.run()
================================================
FILE: tensor2tensor/data_generators/wikisum/wikisum.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Wikipedia Summarization Problems."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import collections
import json
import math
import os
import re
import string
import tempfile
import six
from tensor2tensor.data_generators import generator_utils
from tensor2tensor.data_generators import problem
from tensor2tensor.data_generators import text_encoder
from tensor2tensor.data_generators import tokenizer
from tensor2tensor.data_generators.wikisum import utils as cc_utils
from tensor2tensor.layers import modalities
from tensor2tensor.utils import metrics
from tensor2tensor.utils import registry
import tensorflow.compat.v1 as tf
PROCESS_FOLDER_PREFIX = "process"
REF_SHARD_FILE_PREFIX = "references.tfrecords.gz"
REF_SHARD_FILE = REF_SHARD_FILE_PREFIX + "-%05d-of-01000"
# Support files
BASE_SUPPORT_DIR = "gs://tensor2tensor-data/wikisum"
WIKI_CONTENT_DIR = os.path.join(BASE_SUPPORT_DIR, "wiki_content")
WIKI_URLS_DIR = os.path.join(BASE_SUPPORT_DIR, "wiki_urls")
WET_METADATA_DIR = os.path.join(BASE_SUPPORT_DIR, "commoncrawl_metadata")
WIKI_CONTENT_FILE = "wiki_content.tfrecords-%05d-of-01000"
WIKI_URLS_FILE = "wiki_urls.json-%05d-of-01000"
EOT = "<EOT>" # end-of-title string
_MIN_REFS = 1
_MIN_LEADSECTION_TOKENS = 1
class WikisumBase(problem.Problem):
"""Base class for Wikisum problems."""
def example_reading_spec(self):
data_fields = {
"inputs": tf.VarLenFeature(tf.int64),
"targets": tf.VarLenFeature(tf.int64),
"section_boundaries": tf.VarLenFeature(tf.int64),
}
data_items_to_decoders = None
return (data_fields, data_items_to_decoders)
@property
def target_vocab_size(self):
return 2**15
@property
def vocab_filename(self):
return "vocab.%s.%d" % (self.dataset_filename(), self.target_vocab_size)
def feature_encoders(self, data_dir):
vocab_filename = os.path.join(data_dir, self.vocab_filename)
encoder = text_encoder.SubwordTextEncoder(vocab_filename)
# Shared encoder for inputs and targets
return {"inputs": encoder, "targets": encoder}
def hparams(self, defaults, unused_model_hparams):
p = defaults
p.stop_at_eos = True
p.vocab_size = {
"inputs": self._encoders["inputs"].vocab_size,
"targets": self._encoders["targets"].vocab_size,
}
p.modality = {
"inputs": modalities.ModalityType.SYMBOL,
"targets": modalities.ModalityType.SYMBOL,
}
def eval_metrics(self):
return super(WikisumBase, self).eval_metrics() + [
metrics.Metrics.ROUGE_2_F, metrics.Metrics.ROUGE_L_F
]
def generate_lines_for_vocab(self, wikis_dir, refs_dir, max_chars=10**7):
total_chars = 0
ref_files_by_shard = _references_files_by_shard(refs_dir)
for shard_id in range(cc_utils.NUM_SHARDS):
# Wikipedia articles
for wiki in _wiki_articles(shard_id, wikis_dir):
yield _normalize_text(wiki.title) + EOT
for section in wiki.sections:
yield _format_title(_normalize_text(section.title))
yield _normalize_text(section.text)
total_chars += len(section.title)
total_chars += len(section.text)
# References
for i, content in enumerate(
six.itervalues(_references_content(ref_files_by_shard[shard_id]))):
for line in content.split("\n"):
if line:
yield _normalize_text(line)
total_chars += len(line)
# Make sure we use at least 1k references
if i >= 1000 and total_chars >= max_chars:
break
if total_chars >= max_chars:
tf.logging.info("Seen enough chars: %d; finished.", max_chars)
break
tf.logging.info("Built vocabulary using %d chars", total_chars)
def generate_vocab(self, data_dir, wikis_dir, refs_dir):
# Produce a SubwordTextEncoder from a subset of the data
return generator_utils.get_or_generate_vocab_inner(
data_dir, self.vocab_filename, self.target_vocab_size,
self.generate_lines_for_vocab(wikis_dir, refs_dir))
def generate_data(self, data_dir, tmp_dir, task_id=-1):
tf.logging.warn("See wikisum/README.md for instructions to generate data.")
def out_filepaths(self, data_dir):
train_shards = 800
dev_shards = 100
test_shards = 100
train_filepaths = self.training_filepaths(
data_dir, train_shards, shuffled=True)
dev_filepaths = self.dev_filepaths(data_dir, dev_shards, shuffled=True)
test_filepaths = self.test_filepaths(data_dir, test_shards, shuffled=True)
out_filepaths = train_filepaths + dev_filepaths + test_filepaths
out_filepaths.sort()
assert len(out_filepaths) == cc_utils.NUM_SHARDS
return out_filepaths
@registry.register_problem
class WikisumCommoncrawl(WikisumBase):
"""Wikipedia references->article summarization task based on CommonCrawl."""
pass
@registry.register_problem
class WikisumWeb(WikisumBase):
"""Wikipedia references->article summarization task based on web data."""
pass
@registry.register_problem
class WikisumCommoncrawlLeadSection(WikisumCommoncrawl):
"""Wikipedia references->lead section summarization task."""
def preprocess_example(self, example, mode, hparams):
example["targets"] = _truncate_to_lead_section(example)
return super(WikisumCommoncrawlLeadSection, self).preprocess_example(
example, mode, hparams)
def dataset_filename(self):
return WikisumCommoncrawl.name
def generate_data(self, data_dir, tmp_dir, task_id=-1):
tf.logging.warn("Problem %s reuses data from problem %s", self.name,
WikisumCommoncrawl.name)
@registry.register_problem
class WikisumWebLeadSection(WikisumWeb):
"""Wikipedia references->lead section summarization task."""
def preprocess_example(self, example, mode, hparams):
example["targets"] = _truncate_to_lead_section(example)
return super(WikisumWebLeadSection, self).preprocess_example(
example, mode, hparams)
def dataset_filename(self):
return WikisumWeb.name
def generate_data(self, data_dir, tmp_dir, task_id=-1):
tf.logging.warn("Problem %s reuses data from problem %s", self.name,
WikisumWeb.name)
def make_ref_shard_files(out_dir):
tf.gfile.MakeDirs(out_dir)
opts = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.GZIP)
files = [
tf.python_io.TFRecordWriter(
os.path.join(out_dir, REF_SHARD_FILE % i), opts)
for i in range(cc_utils.NUM_SHARDS)
]
return files
def _truncate_to_lead_section(example):
wiki = example["targets"]
lead_boundary = example["section_boundaries"][0]
# Concat a new EOS to the lead since the original one gets truncated.
lead = tf.concat((wiki[:lead_boundary], [text_encoder.EOS_ID]), 0)
return lead
def _make_example_from_record(record):
features = {
"url":
tf.train.Feature(bytes_list=tf.train.BytesList(value=[record.url])),
"content":
tf.train.Feature(
bytes_list=tf.train.BytesList(value=[record.content])),
}
return tf.train.Example(features=tf.train.Features(feature=features))
def _shard_id_for_file(sharded_filename):
suffix = "00000-of-00000"
parts = sharded_filename[-len(suffix):].split("-")
assert len(parts) == 3
return int(parts[0])
def _references_files_by_shard(refs_dir):
process_dirs = _process_folders(refs_dir)
shards = collections.defaultdict(list)
for d in process_dirs:
ref_files = tf.gfile.Glob(os.path.join(d, REF_SHARD_FILE_PREFIX) + "*")
for f in ref_files:
shards[_shard_id_for_file(f)].append(f)
return shards
def _references_content(ref_files):
"""Returns dict<str ref_url, str ref_content>."""
example_spec = {
"url": tf.FixedLenFeature([], tf.string),
"content": tf.FixedLenFeature([], tf.string),
}
data = {}
for ex in generator_utils.tfrecord_iterator(
ref_files, gzipped=True, example_spec=example_spec):
data[ex["url"]] = text_encoder.to_unicode(ex["content"])
return data
def _wiki_urls_for_shard(shard_id, urls_dir=None):
"""Urls for chunk: dict<str wiki_url, list<str> ref_urls>."""
urls_dir = urls_dir or WIKI_URLS_DIR
urls_filepath = os.path.join(urls_dir, WIKI_URLS_FILE % shard_id)
with tf.gfile.GFile(urls_filepath) as f:
return json.loads(f.read())
class WikipediaSection(
collections.namedtuple("WikipediaSection", ["title", "text"])):
pass
class WikipediaArticle(
collections.namedtuple("WikipediaArticle", ["url", "title", "sections"])):
pass
def _wiki_articles(shard_id, wikis_dir=None):
"""Generates WikipediaArticles from GCS that are part of shard shard_id."""
if not wikis_dir:
wikis_dir = WIKI_CONTENT_DIR
with tf.Graph().as_default():
dataset = tf.data.TFRecordDataset(
cc_utils.readahead(
os.path.join(wikis_dir, WIKI_CONTENT_FILE % shard_id)),
buffer_size=16 * 1000 * 1000)
def _parse_example(ex_ser):
"""Parse serialized Example containing Wikipedia article content."""
features = {
"url": tf.VarLenFeature(tf.string),
"title": tf.VarLenFeature(tf.string),
"section_titles": tf.VarLenFeature(tf.string),
"section_texts": tf.VarLenFeature(tf.string),
}
ex = tf.parse_single_example(ex_ser, features)
for k in ex.keys():
ex[k] = ex[k].values
ex["url"] = ex["url"][0]
ex["title"] = ex["title"][0]
return ex
dataset = dataset.map(_parse_example, num_parallel_calls=32)
dataset = dataset.prefetch(100)
record_it = dataset.make_one_shot_iterator().get_next()
with tf.Session() as sess:
while True:
try:
ex = sess.run(record_it)
except tf.errors.OutOfRangeError:
break
sections = [
WikipediaSection(title=text_encoder.to_unicode(title),
text=text_encoder.to_unicode(text))
for title, text in zip(ex["section_titles"], ex["section_texts"])
]
yield WikipediaArticle(
url=text_encoder.to_unicode(ex["url"]),
title=text_encoder.to_unicode(ex["title"]),
sections=sections)
def _token_counts(text, token_set=None):
counts = collections.defaultdict(int)
for token in tokenizer.encode(text_encoder.native_to_unicode(text)):
if token_set and token not in token_set:
continue
counts[token] += 1
return counts
def _normalize_text(text):
text = text.lower()
# Space around punctuation
text = re.sub("[%s]" % re.escape(string.punctuation), r" \g<0> ", text)
text = re.sub(r"\s+", " ", text)
text = text.strip()
return text
def _tokens_to_score(tokens):
return {t for t in tokens if re.search("[a-z0-9]", t)}
def rank_reference_paragraphs(wiki_title, references_content, normalize=True):
"""Rank and return reference paragraphs by tf-idf score on title tokens."""
normalized_title = _normalize_text(wiki_title)
title_tokens = _tokens_to_score(
set(tokenizer.encode(text_encoder.native_to_unicode(normalized_title))))
ref_paragraph_info = []
doc_counts = collections.defaultdict(int)
for ref in references_content:
for paragraph in ref.split("\n"):
normalized_paragraph = _normalize_text(paragraph)
if cc_utils.filter_paragraph(normalized_paragraph):
# Skip paragraph
continue
counts = _token_counts(normalized_paragraph, title_tokens)
for token in title_tokens:
if counts[token]:
doc_counts[token] += 1
content = normalized_paragraph if normalize else paragraph
info = {"content": content, "counts": counts}
ref_paragraph_info.append(info)
for info in ref_paragraph_info:
score = 0.
for token in title_tokens:
term_frequency = info["counts"][token]
inv_doc_frequency = (
float(len(ref_paragraph_info)) / max(doc_counts[token], 1))
score += term_frequency * math.log(inv_doc_frequency)
info["score"] = score
ref_paragraph_info.sort(key=lambda el: el["score"], reverse=True)
return [info["content"] for info in ref_paragraph_info]
def produce_examples(shard_ids, wikis_dir, refs_dir, urls_dir, vocab_path,
out_filepaths):
"""Produce examples from shard_ids to out_filepaths."""
# * Join the Wikipedia articles with their references
# * Run Tf-idf to sort reference paragraphs
# * Encode the Wikipedia and reference text with the vocabulary
# * Write out TFRecords of tensorflow.Example
tf.logging.info("Processing %d input shards into %d output files.",
len(shard_ids), len(out_filepaths))
vocab = text_encoder.SubwordTextEncoder(vocab_path)
eot_ids = vocab.encode(EOT)
def example_generator():
"""Generate Example dicts."""
stats = dict(total_original_wikis=0, total_original_refs=0,
total_found_refs=0, ref_lengths=[], wiki_original_refs=[],
wiki_found_refs=[], wikis_skipped_no_refs=0,
wikis_skipped_short_lead=0, num_wikis_written=0)
ref_files_by_shard = _references_files_by_shard(refs_dir)
for shard_id in shard_ids:
tf.logging.info("Processing shard %d", shard_id)
wiki_urls = _wiki_urls_for_shard(shard_id, urls_dir)
tf.logging.info("Loaded wiki URLs for shard")
refs_content = _references_content(ref_files_by_shard[shard_id])
tf.logging.info("Loaded reference content for shard")
for i, wiki in enumerate(_wiki_articles(shard_id, wikis_dir)):
if not i % 1000:
tf.logging.info("Processing wiki index %d for shard %d", i, shard_id)
stats["total_original_wikis"] += 1
# Get reference content
wiki_ref_content = []
ref_urls = wiki_urls[wiki.url]["refs"]
stats["total_original_refs"] += len(ref_urls)
stats_wiki_original_refs = len(ref_urls)
stats_wiki_found_refs = 0
for ref_url in ref_urls:
ref_content = refs_content.get(ref_url)
if not ref_content:
continue
stats["total_found_refs"] += 1
stats["ref_lengths"].append(len(ref_content))
stats_wiki_found_refs += 1
wiki_ref_content.append(ref_content)
stats["wiki_original_refs"].append(stats_wiki_original_refs)
stats["wiki_found_refs"].append(stats_wiki_found_refs)
if not wiki_ref_content or len(wiki_ref_content) < _MIN_REFS:
# No/few refs were found
stats["wikis_skipped_no_refs"] += 1
continue
# Rank reference paragraphs with TFIDF
wiki_title = _normalize_text(wiki.title)
ranked_paragraphs = rank_reference_paragraphs(wiki_title,
wiki_ref_content)
# Construct inputs from Wiki title and references
inputs = []
inputs.extend(vocab.encode(wiki_title))
inputs.extend(eot_ids)
for paragraph in ranked_paragraphs:
if len(inputs) >= 1e6:
break
paragraph += " "
inputs.extend(vocab.encode(paragraph))
# Construct targets from article sections
targets, section_boundaries = _encode_wiki_sections(
wiki.sections, vocab)
# Skip if lead section is too short
if (not section_boundaries or
section_boundaries[0] < _MIN_LEADSECTION_TOKENS):
stats["wikis_skipped_short_lead"] += 1
continue
inputs.append(text_encoder.EOS_ID)
targets.append(text_encoder.EOS_ID)
stats["num_wikis_written"] += 1
yield {
"inputs": inputs,
"targets": targets,
"section_boundaries": section_boundaries,
}
tf.logging.info("Total: %d, Skipped: %d",
stats["num_wikis_written"],
stats["total_original_wikis"] - stats["num_wikis_written"])
tf.logging.info("Total refs: %d, Skipped refs: %d",
stats["total_found_refs"],
stats["total_original_refs"] - stats["total_found_refs"])
stats_fname = os.path.join(os.path.split(out_filepaths[0])[0],
"stats.%d.json" % shard_ids[0])
with tf.gfile.Open(stats_fname, "w") as f:
f.write(json.dumps(stats))
generator_utils.generate_files(example_generator(), out_filepaths)
def _format_title(title):
return " == %s == " % title
def _encode_wiki_sections(sections, vocab):
"""Encodes sections with vocab. Returns ids and section boundaries."""
ids = []
section_boundaries = []
for i, section in enumerate(sections):
if i > 0:
# Skip including article title
ids.extend(vocab.encode(_format_title(_normalize_text(section.title))))
ids.extend(vocab.encode(_normalize_text(section.text)))
section_boundaries.append(len(ids))
return ids, section_boundaries
def _process_folders(tmp_dir):
return tf.gfile.Glob(os.path.join(tmp_dir, PROCESS_FOLDER_PREFIX) + "*")
def extract_references_from_wets(wet_files, metadata_dir, out_dir,
tmp_dir=None):
"""Extract references from WET files into sharded output files."""
# Setup output files
shard_files = make_ref_shard_files(out_dir)
num_refs = 0
for i, wet_file in enumerate(wet_files):
num_refs_in_wet = 0
tf.logging.info("Processing file %d", i)
# Read metadata file
metadata_fname = os.path.join(
metadata_dir, os.path.basename(wet_file)) + cc_utils.METADTA_SUFFIX
with tf.gfile.Open(cc_utils.readahead(metadata_fname)) as f:
wet_metadata = json.loads(f.read())
if not wet_metadata:
# No references in this WET file
continue
if wet_file.startswith("http"):
# download
if not tmp_dir:
tmp_dir = tempfile.gettempdir()
record_gen = cc_utils.wet_records_from_url(wet_file, tmp_dir)
else:
# local
record_gen = cc_utils.wet_records_from_file_obj(
cc_utils.gzip_memfile(wet_file), take_ownership=True)
for wet_record in record_gen:
shard_ids = wet_metadata.get(wet_record.url)
if not shard_ids:
# URL not in dataset
continue
# Serialize and write out
ex = _make_example_from_record(wet_record)
ex_str = ex.SerializeToString()
for shard_id in shard_ids:
shard_files[shard_id].write(ex_str)
num_refs += 1
num_refs_in_wet += 1
tf.logging.info("Wrote out %d references for this WET", num_refs_in_wet)
tf.logging.info("Wrote out %d references total", num_refs)
# Cleanup
for shard_file in shard_files:
shard_file.close()
================================================
FILE: tensor2tensor/data_generators/wikitext103.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Data generators for wikitext-103.
Wikitext-103: Long term dependency language modeling dataset
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import collections
import os
import zipfile
from tensor2tensor.data_generators import generator_utils
from tensor2tensor.data_generators import problem
from tensor2tensor.data_generators import text_encoder
from tensor2tensor.data_generators import text_problems
from tensor2tensor.utils import registry
import tensorflow.compat.v1 as tf
def _build_vocab(filename, vocab_dir, vocab_name):
"""Reads a file to build a vocabulary.
Args:
filename: file to read list of words from.
vocab_dir: directory where to save the vocabulary.
vocab_name: vocab file name.
Returns:
text encoder.
"""
vocab_path = os.path.join(vocab_dir, vocab_name)
if not tf.gfile.Exists(vocab_path):
with tf.gfile.GFile(filename, "r") as f:
data = f.read().split()
counter = collections.Counter(data)
count_pairs = sorted(counter.items(), key=lambda x: (-x[1], x[0]))
words, _ = list(zip(*count_pairs))
encoder = text_encoder.TokenTextEncoder(None, vocab_list=words)
encoder.store_to_file(vocab_path)
else:
encoder = text_encoder.TokenTextEncoder(vocab_path)
return encoder
def _maybe_download_corpus(tmp_dir, vocab_type):
"""Download and unpack the corpus.
Args:
tmp_dir: directory containing dataset.
vocab_type: which vocabulary are we using.
Returns:
The list of names of files.
"""
if vocab_type == text_problems.VocabType.CHARACTER:
dataset_url = ("https://s3.amazonaws.com/research.metamind.io/wikitext"
"/wikitext-103-raw-v1.zip")
dir_name = "wikitext-103-raw"
else:
dataset_url = ("https://s3.amazonaws.com/research.metamind.io/wikitext"
"/wikitext-103-v1.zip")
dir_name = "wikitext-103"
fname = os.path.basename(dataset_url)
compressed_filepath = generator_utils.maybe_download(tmp_dir, fname,
dataset_url)
zip_ref = zipfile.ZipFile(compressed_filepath, "r")
zip_ref.extractall(tmp_dir)
zip_ref.close()
files = os.path.join(tmp_dir, dir_name, "*")
train_file, valid_file, test_file = None, None, None
for f in tf.gfile.Glob(files):
fname = os.path.basename(f)
if "train" in fname:
train_file = f
elif "valid" in fname:
valid_file = f
elif "test" in fname:
test_file = f
assert train_file, "Training file not found"
assert valid_file, "Validation file not found"
assert test_file, "Testing file not found"
return train_file, valid_file, test_file
@registry.register_problem
class LanguagemodelWikitext103(text_problems.Text2SelfProblem):
"""Wikitext103 dataset token-level."""
@property
def dataset_splits(self):
return [{
"split": problem.DatasetSplit.TRAIN,
"shards": 10,
}, {
"split": problem.DatasetSplit.EVAL,
"shards": 1,
}, {
"split": problem.DatasetSplit.TEST,
"shards": 1,
}]
@property
def is_generate_per_split(self):
return True
@property
def vocab_type(self):
return text_problems.VocabType.TOKEN
def generate_samples(self, data_dir, tmp_dir, dataset_split):
train_file, valid_file, test_file = _maybe_download_corpus(
tmp_dir, self.vocab_type)
if dataset_split == problem.DatasetSplit.TRAIN:
filepath = train_file
if self.vocab_type == text_problems.VocabType.TOKEN:
_build_vocab(train_file, data_dir, self.vocab_filename)
elif dataset_split == problem.DatasetSplit.EVAL:
filepath = valid_file
elif dataset_split == problem.DatasetSplit.TEST:
filepath = test_file
def _generate_samples():
with tf.gfile.GFile(filepath, "r") as f:
for line in f:
line = " ".join(line.strip().split())
if line:
yield {"targets": line}
return _generate_samples()
@registry.register_problem
class LanguagemodelWikitext103Characters(LanguagemodelWikitext103):
"""Wikitext-103, character-level."""
@property
def vocab_type(self):
return text_problems.VocabType.CHARACTER
@registry.register_problem
class LanguagemodelWikitext103L4k(LanguagemodelWikitext103):
"""Wikitext-103, token-level, with examples up to 4,096 tokens long."""
def generate_samples(self, data_dir, tmp_dir, dataset_split):
samples_by_line = super(LanguagemodelWikitext103L4k,
self).generate_samples(data_dir, tmp_dir,
dataset_split)
def _generate_samples():
tokens = []
for sample in samples_by_line:
sample_tokens = sample["targets"].split()
if len(tokens) + len(sample_tokens) < self.sequence_length:
tokens.extend(sample_tokens)
else:
yield {"targets": " ".join(tokens)}
tokens = sample_tokens
return _generate_samples()
def max_length(self, model_hparams):
return model_hparams.split_to_length or self.sequence_length
@property
def sequence_length(self):
"""Length of each example (in tokens)."""
return 4096
@registry.register_problem
class LanguagemodelWikitext103L16k(LanguagemodelWikitext103L4k):
"""Wikitext-103, token-level, with examples up to 16,384 tokens long."""
@property
def sequence_length(self):
"""Length of each example (in tokens)."""
return 16384
================================================
FILE: tensor2tensor/data_generators/wnli.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Data generators for the Winograd NLI dataset."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import zipfile
from tensor2tensor.data_generators import generator_utils
from tensor2tensor.data_generators import problem
from tensor2tensor.data_generators import text_encoder
from tensor2tensor.data_generators import text_problems
from tensor2tensor.utils import registry
import tensorflow.compat.v1 as tf
EOS = text_encoder.EOS
@registry.register_problem
class WinogradNLI(text_problems.TextConcat2ClassProblem):
"""Winograd NLI classification problems."""
# Link to data from GLUE: https://gluebenchmark.com/tasks
_WNLI_URL = ("https://firebasestorage.googleapis.com/v0/b/"
"mtl-sentence-representations.appspot.com/o/"
"data%2FWNLI.zip?alt=media&token=068ad0a0-ded7-"
"4bd7-99a5-5e00222e0faf")
@property
def is_generate_per_split(self):
return True
@property
def dataset_splits(self):
return [{
"split": problem.DatasetSplit.TRAIN,
"shards": 1,
}, {
"split": problem.DatasetSplit.EVAL,
"shards": 1,
}]
@property
def approx_vocab_size(self):
return 2**13 # 8k vocab suffices for this small dataset.
@property
def vocab_filename(self):
return "vocab.wnli.%d" % self.approx_vocab_size
@property
def num_classes(self):
return 2
def class_labels(self, data_dir):
del data_dir
# Note this binary classification is different from usual MNLI.
return ["not_entailment", "entailment"]
def _maybe_download_corpora(self, tmp_dir):
wnli_filename = "WNLI.zip"
wnli_finalpath = os.path.join(tmp_dir, "WNLI")
if not tf.gfile.Exists(wnli_finalpath):
zip_filepath = generator_utils.maybe_download(
tmp_dir, wnli_filename, self._WNLI_URL)
zip_ref = zipfile.ZipFile(zip_filepath, "r")
zip_ref.extractall(tmp_dir)
zip_ref.close()
return wnli_finalpath
def example_generator(self, filename):
for idx, line in enumerate(tf.gfile.Open(filename, "rb")):
if idx == 0: continue # skip header
line = text_encoder.to_unicode_utf8(line.strip())
_, s1, s2, l = line.split("\t")
inputs = [s1, s2]
yield {
"inputs": inputs,
"label": int(l)
}
def generate_samples(self, data_dir, tmp_dir, dataset_split):
wnli_dir = self._maybe_download_corpora(tmp_dir)
if dataset_split == problem.DatasetSplit.TRAIN:
filesplit = "train.tsv"
else:
filesplit = "dev.tsv"
filename = os.path.join(wnli_dir, filesplit)
for example in self.example_generator(filename):
yield example
@registry.register_problem
class WinogradNLICharacters(WinogradNLI):
"""Winograd NLI classification problems, character level"""
@property
def vocab_type(self):
return text_problems.VocabType.CHARACTER
def global_task_id(self):
return problem.TaskID.EN_NLI
================================================
FILE: tensor2tensor/data_generators/wsj_parsing.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Data generators for parsing data-sets."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
from absl import flags
from tensor2tensor.data_generators import problem
from tensor2tensor.data_generators import text_problems
from tensor2tensor.utils import registry
import tensorflow.compat.v1 as tf
flags.DEFINE_string("parsing_path", "", "Path to parsing files in tmp_dir.")
FLAGS = flags.FLAGS
@registry.register_problem
class WsjParsing(text_problems.Text2textTmpdir):
"""Generate vocabulary and training data for parsing.
"""
# These files are used for vocab generation
TRAIN_FILES = ("wsj.train.text.txt", "wsj.train.tags.txt")
# These files are used for generating encoded samples
TRAIN_FILES_TREE = "wsjTrain.trees"
EVAL_FILES_TREE = "wsjEval.trees"
def generate_samples(self, data_dir, tmp_dir, dataset_split):
del data_dir
is_training = dataset_split == problem.DatasetSplit.TRAIN
tree_file = self.TRAIN_FILES_TREE if is_training else self.EVAL_FILES_TREE
tree_file_path = os.path.join(tmp_dir, tree_file)
with tf.gfile.GFile(tree_file_path, mode="r") as cur_tree_file:
for line in cur_tree_file:
(words, tags) = words_and_tags_from_wsj_tree(line)
yield {"inputs": words, "targets": tags}
def generate_encoded_samples(self, data_dir, tmp_dir, dataset_split):
generator = self.generate_samples(data_dir, tmp_dir, dataset_split)
encoder = self.get_or_create_vocab(data_dir, tmp_dir)
return text_problems.text2text_generate_encoded(generator, encoder,
has_inputs=self.has_inputs)
def generate_text_for_vocab(self, data_dir, tmp_dir):
files = [os.path.join(tmp_dir, f) for f in self.TRAIN_FILES]
inputs_file, targets_file = files
for i, sample in enumerate(text_problems.text2text_txt_iterator(inputs_file,
targets_file
)):
yield sample["inputs"]
yield sample["targets"]
if self.max_samples_for_vocab and (i + 1) >= self.max_samples_for_vocab:
break
@property
def max_samples_for_vocab(self):
return 1000
def words_and_tags_from_wsj_tree(tree_string):
"""Generates linearized trees and tokens from the wsj tree format.
It uses the linearized algorithm described in https://arxiv.org/abs/1412.7449.
Args:
tree_string: tree in wsj format
Returns:
tuple: (words, linearized tree)
"""
stack, tags, words = [], [], []
for tok in tree_string.strip().split():
if tok[0] == "(":
symbol = tok[1:]
tags.append(symbol)
stack.append(symbol)
else:
assert tok[-1] == ")"
stack.pop() # Pop the POS-tag.
while tok[-2] == ")":
tags.append("/" + stack.pop())
tok = tok[:-1]
words.append(tok[:-1])
return str.join(" ", words), str.join(" ", tags[1:-1]) # Strip "TOP" tag.
def token_generator(tree_path, source_token_vocab, target_token_vocab,
eos=None):
"""Generator for parsing as a sequence-to-sequence task that uses tokens.
This generator assumes the files at source_path and target_path have
the same number of lines and yields dictionaries of "inputs" and "targets"
where inputs and targets are token ids from source and target lines
converted to integers using the token_map.
Args:
tree_path: path to the file with WSJ format trees, one per line.
source_token_vocab: GenericVocabulary object for source vocabulary.
target_token_vocab: GenericVocabulary object for target vocabulary.
eos: integer to append at the end of each sequence (default: None).
Yields:
A dictionary {"inputs": source-line, "targets": target-line} where
the lines are integer lists converted from tokens in the file lines.
"""
eos_list = [] if eos is None else [eos]
with tf.gfile.GFile(tree_path, mode="r") as tree_file:
tree_line = tree_file.readline()
while tree_line:
source, target = words_and_tags_from_wsj_tree(tree_line)
source_ints = source_token_vocab.encode(source.strip()) + eos_list
target_ints = target_token_vocab.encode(target.strip()) + eos_list
yield {"inputs": source_ints, "targets": target_ints}
tree_line = tree_file.readline()
def parsing_token_generator(data_dir, tmp_dir, train, source_vocab_size,
target_vocab_size):
"""Generator for parsing as a sequence-to-sequence task that uses tokens.
This generator assumes the files parsing_{train,dev}.trees, which contain
trees in WSJ format.
Args:
data_dir: path to the data directory.
tmp_dir: path to temporary storage directory.
train: whether we're training or not.
source_vocab_size: source vocab size.
target_vocab_size: target vocab size.
Returns:
A generator to a dictionary of inputs and outputs.
"""
# TODO(lukaszkaiser): Correct these calls to generate vocabularies. No data
# sources are being passed.
del (data_dir, tmp_dir, train, source_vocab_size, target_vocab_size)
assert False, "Vocabulary generation not implemented"
# source_symbolizer_vocab = generator_utils.get_or_generate_vocab(
# data_dir, tmp_dir, "wsj_source.vocab.%d" % source_vocab_size,
# source_vocab_size)
# target_symbolizer_vocab = generator_utils.get_or_generate_vocab(
# data_dir, tmp_dir, "wsj_target.vocab.%d" % target_vocab_size,
# target_vocab_size)
# filename = "%s_%s.trees" % (FLAGS.parsing_path, "train" if train else "dev")
# tree_filepath = os.path.join(tmp_dir, filename)
# return token_generator(tree_filepath, source_symbolizer_vocab,
# target_symbolizer_vocab, 1)
================================================
FILE: tensor2tensor/data_generators/yelp_full.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Yelp dataset."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import tarfile
from tensor2tensor.data_generators import generator_utils
from tensor2tensor.data_generators import problem
from tensor2tensor.data_generators import text_problems
from tensor2tensor.utils import registry
import tensorflow.compat.v1 as tf
@registry.register_problem
class SentimentYelpFull(text_problems.Text2ClassProblem):
"""Yelp dataset."""
URL = "https://s3.amazonaws.com/fast-ai-nlp/yelp_review_full_csv.tgz"
@property
def is_generate_per_split(self):
return True
@property
def dataset_splits(self):
return [{
"split": problem.DatasetSplit.TRAIN,
"shards": 10,
}, {
"split": problem.DatasetSplit.EVAL,
"shards": 1,
}]
@property
def approx_vocab_size(self):
return 2**13 # 8k vocab suffices for this small dataset.
@property
def num_classes(self):
return 5
def class_labels(self, data_dir):
del data_dir
return ["1", "2", "3", "4", "5"]
def doc_generator(self, yelp_dir, dataset, include_label=False):
file_path = os.path.join(yelp_dir, dataset + ".csv")
with tf.gfile.Open(file_path) as yelp_f:
lines = yelp_f.readlines()
for line in lines:
label = line[1]
doc = line[5:-2].strip()
if include_label:
yield doc, label
else:
yield doc
def generate_samples(self, data_dir, tmp_dir, dataset_split):
"""Generate examples."""
# Download and extract
compressed_filename = os.path.basename(self.URL)
download_path = generator_utils.maybe_download(tmp_dir, compressed_filename,
self.URL)
yelp_dir = os.path.join(tmp_dir, "yelp_review_full_csv")
if not tf.gfile.Exists(yelp_dir):
with tarfile.open(download_path, "r:gz") as tar:
tar.extractall(tmp_dir)
# Generate examples
train = dataset_split == problem.DatasetSplit.TRAIN
dataset = "train" if train else "test"
for doc, label in self.doc_generator(yelp_dir, dataset, include_label=True):
yield {
"inputs": doc,
"label": int(label),
}
@registry.register_problem
class SentimentYelpFullCharacters(SentimentYelpFull):
"""Yelp dataset, character level."""
@property
def vocab_type(self):
return text_problems.VocabType.CHARACTER
def global_task_id(self):
return problem.TaskID.EN_CHR_SENT
================================================
FILE: tensor2tensor/data_generators/yelp_polarity.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Yelp dataset."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import tarfile
from tensor2tensor.data_generators import generator_utils
from tensor2tensor.data_generators import problem
from tensor2tensor.data_generators import text_problems
from tensor2tensor.utils import registry
import tensorflow.compat.v1 as tf
@registry.register_problem
class SentimentYelpPolarity(text_problems.Text2ClassProblem):
"""Yelp dataset."""
URL = "https://s3.amazonaws.com/fast-ai-nlp/yelp_review_polarity_csv.tgz"
@property
def is_generate_per_split(self):
return True
@property
def dataset_splits(self):
return [{
"split": problem.DatasetSplit.TRAIN,
"shards": 10,
}, {
"split": problem.DatasetSplit.EVAL,
"shards": 1,
}]
@property
def approx_vocab_size(self):
return 2**13 # 8k vocab suffices for this small dataset.
@property
def num_classes(self):
return 2
def class_labels(self, data_dir):
del data_dir
return ["1", "2"]
def doc_generator(self, yelp_dir, dataset, include_label=False):
file_path = os.path.join(yelp_dir, dataset + ".csv")
with tf.gfile.Open(file_path) as yelp_f:
lines = yelp_f.readlines()
for line in lines:
label = line[1]
doc = line[5:-2].strip()
if include_label:
yield doc, label
else:
yield doc
def generate_samples(self, data_dir, tmp_dir, dataset_split):
"""Generate examples."""
# Download and extract
compressed_filename = os.path.basename(self.URL)
download_path = generator_utils.maybe_download(tmp_dir, compressed_filename,
self.URL)
yelp_dir = os.path.join(tmp_dir, "yelp_review_polarity_csv")
if not tf.gfile.Exists(yelp_dir):
with tarfile.open(download_path, "r:gz") as tar:
tar.extractall(tmp_dir)
# Generate examples
train = dataset_split == problem.DatasetSplit.TRAIN
dataset = "train" if train else "test"
for doc, label in self.doc_generator(yelp_dir, dataset, include_label=True):
yield {
"inputs": doc,
"label": int(label),
}
@registry.register_problem
class SentimentYelpPolarityCharacters(SentimentYelpPolarity):
"""Yelp dataset, character level."""
@property
def vocab_type(self):
return text_problems.VocabType.CHARACTER
def global_task_id(self):
return problem.TaskID.EN_CHR_SENT
================================================
FILE: tensor2tensor/envs/__init__.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Environments defined in T2T. Imports here force registration."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from tensor2tensor.envs import gym_env_problem
from tensor2tensor.envs import tic_tac_toe_env
from tensor2tensor.envs import tic_tac_toe_env_problem
================================================
FILE: tensor2tensor/envs/env_problem.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Base class for envs that store their history.
EnvProblem subclasses Problem and also implements the Gym interface (step,
reset, render, close, seed)
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from absl import logging
from gym.core import Env
import numpy as np
import six
from tensor2tensor.data_generators import generator_utils
from tensor2tensor.data_generators import problem
from tensor2tensor.envs import gym_spaces_utils
from tensor2tensor.envs import trajectory
from tensor2tensor.layers import modalities
from tensor2tensor.utils import contrib
import tensorflow.compat.v1 as tf
# Names for data fields in stored tf.Examples.
TIMESTEP_FIELD = "timestep"
ACTION_FIELD = "action"
RAW_REWARD_FIELD = "raw_reward"
PROCESSED_REWARD_FIELD = "reward"
DONE_FIELD = "done"
OBSERVATION_FIELD = "observation"
class EnvProblem(Env, problem.Problem):
"""Base class of an env which generates data like a problem class.
EnvProblem is both a gym Env and a Problem, since it subclasses both.
Conceptually it contains `batch_size` environments on which step (and reset)
are called. The data that is generated by the repeated application of step and
reset is stored within this class and is persisted on disk when we call
`generate_data` on it.
Subclasses *should* override the following functions:
- initialize_environments
- observation_space
- action_space
- reward_range
- _reset
- _step
- _render
In addition, they should ovveride the following functions, which are used in
the `hparams` function to return modalities and vocab_sizes.
- input_modality
- input_vocab_size
- target_modality
- target_vocab_size
- action_modality
- reward_modality
NON NATIVELY BATCHED ENVS:
The implementation for cases where the env is not batched by default is
`gym_env_problem.GymEnvProblem`.
NATIVELY BATCHED ENVS:
If however, our env is a neural network, which can be batched by default, we
should:
# 1 - Give it a gym style interface, by overriding observation_space and
action_space.
# 2 - Override `_reset` and `_step` to do the reset and step in a natively
batched manner.
# 3 - More generally any function that iterates over the self._env list will
need to be overridden, ex: `_verify_same_spaces` and `initialize_environments`
KNOWN LIMITATIONS:
- observation_space and action_space should be subclasses of gym.spaces
- not all subclasses of gym.spaces are supported
"""
def __init__(self,
batch_size=None,
discrete_rewards=True,
parallelism=1,
**env_kwargs):
"""Initializes this class by creating the envs and managing trajectories.
Args:
batch_size: (int or None) How many envs to make in the non natively
batched mode.
discrete_rewards: (bool) whether to round the rewards to the nearest
integer.
parallelism: (int) If this is greater than one then we run the envs in
parallel using multi-threading.
**env_kwargs: (dict) Additional kwargs to pass to the environments.
"""
# Call the super's ctor.
problem.Problem.__init__(self, was_reversed=False, was_copy=False)
# An env generates data when it is given actions by an agent which is either
# a policy or a human -- this is supposed to be the `id` of the agent.
#
# In practice, this is used only to store (and possibly retrieve) history
# to an appropriate directory.
self._agent_id = "default"
# If set, we discretize the rewards and treat them as integers.
self._discrete_rewards = discrete_rewards
# A data structure to hold the `batch_size` currently active trajectories
# and also the ones that are completed, i.e. done.
self._trajectories = None
self._batch_size = None
self._parallelism = None
# The parallelism is passes in via env_kwargs because it will be used by
# `GymEnvProblem` to paralellize env actions across a batch.
env_kwargs["parallelism"] = parallelism
if batch_size is not None:
self.initialize(batch_size=batch_size, **env_kwargs)
@property
def batch_size(self):
# TODO(afrozm): I've added this here since it is being used in a lot of
# places in ppo_learner.py -- re-evaluate if needed.
return self._batch_size
@property
def trajectories(self):
return self._trajectories
@trajectories.setter
def trajectories(self, trajectories_):
assert self.trajectories.batch_size == trajectories_.batch_size
self._trajectories = trajectories_
def initialize(self, batch_size=1, **kwargs):
self.initialize_environments(batch_size=batch_size, **kwargs)
self._batch_size = batch_size
# This data structure stores the history of each env.
#
# NOTE: Even if the env is a NN and can step in all batches concurrently, it
# is still valuable to store the trajectories separately.
self._trajectories = trajectory.BatchTrajectory(batch_size=batch_size)
# Assert that *all* the above are now set, we should do this since
# subclasses can override `initialize_environments`.
self.assert_common_preconditions()
assert self.observation_space is not None
assert self.action_space is not None
assert self.reward_range is not None
def initialize_environments(self, batch_size=1, parallelism=1, **kwargs):
"""Initializes the environments.
Args:
batch_size: (int) Number of envs to initialize.
parallelism: (int) If this is greater than one then we allow the
implementation to use multi-threading to step the envs.
**kwargs: (dict) Any additional args needed to initialize the envs.
"""
raise NotImplementedError
def assert_common_preconditions(self):
pass
@property
def observation_space(self):
raise NotImplementedError
@property
def observation_spec(self):
"""The spec for reading an observation stored in a tf.Example."""
return gym_spaces_utils.gym_space_spec(self.observation_space)
def process_observations(self, observations):
"""Processes observations prior to saving in the trajectories.
Args:
observations: (np.ndarray) observations to be processed.
Returns:
processed observation
"""
return observations
@property
def action_space(self):
raise NotImplementedError
@property
def action_spec(self):
"""The spec for reading an observation stored in a tf.Example."""
return gym_spaces_utils.gym_space_spec(self.action_space)
@property
def action_modality(self):
raise NotImplementedError
@property
def num_actions(self):
"""Returns the number of actions in a discrete action space."""
return gym_spaces_utils.cardinality(self.action_space)
@property
def reward_range(self):
# We clip rewards to this range before processing them further, as described
# in `process_rewards`.
raise NotImplementedError
@property
def is_reward_range_finite(self):
min_reward, max_reward = self.reward_range
return (min_reward != -np.inf) and (max_reward != np.inf)
@property
def discrete_rewards(self):
return self._discrete_rewards
def process_rewards(self, rewards):
"""Clips the rewards, optionally rounds them and casts to integer.
Args:
rewards: numpy array of raw (float) rewards.
Returns:
processed_rewards: numpy array of np.int64
"""
min_reward, max_reward = self.reward_range
# Clips at min and max reward.
rewards = np.clip(rewards, min_reward, max_reward)
if self._discrete_rewards:
# Round to (nearest) int and convert to integral type.
rewards = np.around(rewards, decimals=0).astype(np.int64)
return rewards
@property
def is_processed_rewards_discrete(self):
"""Returns true if `self.process_rewards` returns discrete rewards."""
# Subclasses can override, but it should match their self.process_rewards.
# This check is a little hackily.
return self.process_rewards(0.0).dtype == np.int64
@property
def num_rewards(self):
"""Returns the number of distinct rewards.
Returns:
Returns None if the reward range is infinite or the processed rewards
aren't discrete, otherwise returns the number of distinct rewards.
"""
# Pre-conditions: reward range is finite.
# : processed rewards are discrete.
if not self.is_reward_range_finite:
logging.warn("Infinite reward range, `num_rewards returning None`")
return None
if not self.is_processed_rewards_discrete:
logging.warn(
"Processed rewards are not discrete, `num_rewards` returning None")
return None
min_reward, max_reward = self.reward_range
return max_reward - min_reward + 1
@property
def input_modality(self):
raise NotImplementedError
@property
def reward_modality(self):
raise NotImplementedError
@property
def input_vocab_size(self):
raise NotImplementedError
@property
def target_modality(self):
raise NotImplementedError
@property
def target_vocab_size(self):
raise NotImplementedError
@property
def unwrapped(self):
return self
def seed(self, seed=None):
return [seed]
def close(self):
pass
def _reset(self, indices):
"""Resets environments at indices shouldn't pre-process or record.
Args:
indices: list of indices of underlying envs to call reset on.
Returns:
np.ndarray of stacked observations from the reset-ed envs.
"""
raise NotImplementedError
def truncate(self, indices=None, num_to_keep=1):
"""Truncates trajectories at the specified indices."""
if indices is None:
indices = np.arange(self.batch_size)
self.trajectories.truncate_trajectories(indices, num_to_keep=num_to_keep)
def reset(self, indices=None):
"""Resets environments at given indices.
Subclasses should override _reset to do the actual reset if something other
than the default implementation is desired.
NOTE: With `indices` as None the recorded trajectories are also erased since
the expecation is that we want to re-use the whole env class from
scratch.
Args:
indices: Indices of environments to reset. If None all envs are reset as
well as trajectories are erased.
Returns:
Batch of initial observations of reset environments.
"""
if indices is None:
self.trajectories.reset_batch_trajectories()
indices = np.arange(self.batch_size)
# If this is empty (not None) then don't do anything, no env was done.
if indices.size == 0:
logging.warning(
"`reset` called with empty indices array, this is a no-op.")
return None
# Pre-conditions: common_preconditions, see `assert_common_preconditions`.
self.assert_common_preconditions()
observations = self._reset(indices)
processed_observations = self.process_observations(observations)
# Record history.
self.trajectories.reset(indices, processed_observations)
return processed_observations
def _render(self, indices, mode="human"):
"""Renders the environments with the given mode on the specified indices.
Args:
indices: array of indices.
mode: rendering mode.
Returns:
a list of return values from the environments rendered.
"""
raise NotImplementedError
def render(self, indices=None, mode="human"):
"""Renders the environments with the given mode on the specified indices.
Args:
indices: array of indices, calls render on everything if indices is None.
mode: rendering mode.
Returns:
a list of return values from the environments rendered.
"""
if indices is None:
indices = np.arange(self.batch_size)
return self._render(indices, mode)
def _step(self, actions):
"""Takes a step in all environments, shouldn't pre-process or record.
Args:
actions: (np.ndarray) with first dimension equal to the batch size.
Returns:
a tuple of stacked raw observations, raw rewards, dones and infos.
"""
raise NotImplementedError
def step(self, actions, infos=None):
"""Takes a step in all environments.
Subclasses should override _step to do the actual reset if something other
than the default implementation is desired.
Args:
actions: Batch of actions.
infos: (optional) a dictionary of keys and values, where all the values
have the first dimension as batch_size.
Returns:
(preprocessed_observations, processed_rewards, dones, env_infos).
"""
# Pre-conditions: common_preconditions, see `assert_common_preconditions`.
# : len(actions) == len(self._envs)
self.assert_common_preconditions()
assert self.batch_size == len(actions)
observations, raw_rewards, dones, env_infos = self._step(actions)
# Process rewards.
raw_rewards = raw_rewards.astype(np.float32)
processed_rewards = self.process_rewards(raw_rewards)
# Process observations.
processed_observations = self.process_observations(observations)
# Record history.
self.trajectories.step(processed_observations, raw_rewards,
processed_rewards, dones, actions,
infos=infos)
return processed_observations, processed_rewards, dones, env_infos
def example_reading_spec(self):
"""Data fields to store on disk and their decoders."""
# Subclasses can override and/or extend.
processed_reward_type = tf.float32
if self.is_processed_rewards_discrete:
processed_reward_type = tf.int64
data_fields = {
TIMESTEP_FIELD: tf.FixedLenFeature((1,), tf.int64),
RAW_REWARD_FIELD: tf.FixedLenFeature((1,), tf.float32),
PROCESSED_REWARD_FIELD: tf.FixedLenFeature((1,), processed_reward_type),
DONE_FIELD: tf.FixedLenFeature((1,), tf.int64), # we wrote this as int.
# Special treatment because we need to determine type and shape, also
# enables classes to override.
OBSERVATION_FIELD: self.observation_spec,
ACTION_FIELD: self.action_spec,
}
data_items_to_decoders = {
field: contrib.slim().tfexample_decoder.Tensor(field)
for field in data_fields
}
return data_fields, data_items_to_decoders
def hparams(self, defaults, model_hparams):
# Usually when using the environment in a supervised setting, given the
# observation we are predicting the reward.
p = defaults
# Have to add these the 'proper' way, otherwise __str__ doesn't show them.
if "modality" not in p:
p.add_hparam("modality", {})
if "vocab_size" not in p:
p.add_hparam("vocab_size", {})
# TODO(afrozm): Document what all of these keys are and are supposed to do.
p.modality.update({
"inputs": self.input_modality,
"targets": self.target_modality,
"input_reward": self.reward_modality,
"target_reward": self.reward_modality,
"input_action": self.action_modality,
"target_action": self.action_modality,
"target_policy": modalities.ModalityType.IDENTITY,
"target_value": modalities.ModalityType.IDENTITY,
})
p.vocab_size.update({
"inputs": self.input_vocab_size,
"targets": self.target_vocab_size,
"input_reward": self.num_rewards,
"target_reward": self.num_rewards,
"input_action": self.num_actions,
"target_action": self.num_actions,
"target_policy": None,
"target_value": None,
})
p.input_space_id = problem.SpaceID.GENERIC
p.target_space_id = problem.SpaceID.GENERIC
@property
def agent_id(self):
return self._agent_id
@agent_id.setter
def agent_id(self, agent_id):
# Lets us call agent_id with integers that we increment.
agent_id = str(agent_id)
# We use `-` in self.dataset_filename, disallow it here for convenience.
if "-" in agent_id:
raise ValueError("agent_id shouldn't have - in it.")
self._agent_id = agent_id
def dataset_filename(self):
return "{}-{}".format(self.name, self.agent_id)
@property
def num_shards(self):
return {
problem.DatasetSplit.TRAIN: 10,
problem.DatasetSplit.EVAL: 1,
}
def _generate_time_steps(self, trajectory_list):
"""A generator to yield single time-steps from a list of trajectories."""
for single_trajectory in trajectory_list:
assert isinstance(single_trajectory, trajectory.Trajectory)
# Skip writing trajectories that have only a single time-step -- this
# could just be a repeated reset.
if single_trajectory.num_time_steps <= 1:
continue
for index, time_step in enumerate(single_trajectory.time_steps):
# The first time-step doesn't have reward/processed_reward, if so, just
# setting it to 0.0 / 0 should be OK.
raw_reward = time_step.raw_reward
if not raw_reward:
raw_reward = 0.0
processed_reward = time_step.processed_reward
if not processed_reward:
processed_reward = 0
action = time_step.action
if action is None:
# The last time-step doesn't have action, and this action shouldn't be
# used, gym's spaces have a `sample` function, so let's just sample an
# action and use that.
action = self.action_space.sample()
action = gym_spaces_utils.gym_space_encode(self.action_space, action)
if six.PY3:
# py3 complains that, to_example cannot handle np.int64 !
action_dtype = self.action_space.dtype
if action_dtype in [np.int64, np.int32]:
action = list(map(int, action))
elif action_dtype in [np.float64, np.float32]:
action = list(map(float, action))
# same with processed_reward.
processed_reward = int(processed_reward)
assert time_step.observation is not None
yield {
TIMESTEP_FIELD: [index],
ACTION_FIELD:
action,
# to_example errors on np.float32
RAW_REWARD_FIELD: [float(raw_reward)],
PROCESSED_REWARD_FIELD: [processed_reward],
# to_example doesn't know bools
DONE_FIELD: [int(time_step.done)],
OBSERVATION_FIELD:
gym_spaces_utils.gym_space_encode(self.observation_space,
time_step.observation),
}
def generate_data(self, data_dir, tmp_dir, task_id=-1):
# List of files to generate data in.
# NOTE: We don't want to shuffle, so we mark the files as shuffled.
files_list = []
for split, num_shards in self.num_shards.items():
files_list.extend(self.data_filepaths(split, data_dir, num_shards, True))
# At this point some trajectories haven't finished. However we still want to
# write those down.
# A simple way of doing this is to call `self.reset()` here, this will make
# all the envs take one (extra) step, but would be a clean way to do it.
#
# self.reset()
self.trajectories.complete_all_trajectories()
# Write the completed data into these files
num_completed_trajectories = self.trajectories.num_completed_trajectories
num_shards = len(files_list)
if num_completed_trajectories < num_shards:
logging.warning(
"Number of completed trajectories [%d] is less than "
"the number of shards [%d], some shards maybe empty.",
num_completed_trajectories, num_shards)
for i, f in enumerate(files_list[:num_completed_trajectories]):
# Start at index i of completed trajectories and take every `num_shards`
# trajectory. This ensures that the data is approximately a balanced
# partition of completed trajectories, also because of the above slicing
# of files_list, i will be a valid index into completed_trajectories.
trajectories_to_write = self.trajectories.completed_trajectories[
i::num_shards]
# Convert each trajectory from `trajectories_to_write` to a sequence of
# time-steps and then send that generator to `generate_files`.
# `cycle_every_n` isn't needed since file list given to it is a singleton.
generator_utils.generate_files(
self._generate_time_steps(trajectories_to_write), [f])
def print_state(self):
for t in self.trajectories.trajectories:
print("---------")
if not t.is_active:
print("trajectory isn't active.")
continue
last_obs = t.last_time_step.observation
print(str(last_obs))
================================================
FILE: tensor2tensor/envs/env_problem_utils.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Utilities to deal with EnvProblem."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import functools
import time
import gym
import numpy as np
from tensor2tensor.envs import gym_env_problem
from tensor2tensor.envs import rendered_env_problem
from tensor2tensor.rl import gym_utils
def done_indices(dones):
"""Calculates the indices where dones has True."""
return np.argwhere(dones).squeeze(axis=1)
def play_env_problem_randomly(env_problem, num_steps):
"""Plays the env problem by randomly sampling actions for `num_steps`."""
# Reset all environments.
env_problem.reset()
# Play all environments, sampling random actions each time.
for _ in range(num_steps):
# Sample batch_size actions from the action space and stack them.
actions = np.stack([
env_problem.action_space.sample() for _ in range(env_problem.batch_size)
])
# Execute actions, observations are stored in `env_problem`.
_, _, dones, _ = env_problem.step(actions)
# Get the indices where we are done and reset those.
env_problem.reset(indices=done_indices(dones))
def get_completed_trajectories_from_env(env,
n_trajectories,
raw_trajectory=False):
"""Returns completed `n_trajectories` from `env`."""
# Just the raw trajectories.
if raw_trajectory:
return env.trajectories.completed_trajectories[:n_trajectories]
# The numpy version of the above.
completed_trajectories = []
for trajectory in env.trajectories.completed_trajectories[:n_trajectories]:
completed_trajectories.append(trajectory.as_numpy)
return completed_trajectories
def play_env_problem_with_policy(env,
policy_fun,
num_trajectories=1,
max_timestep=None,
reset=True,
state=None,
rng=None,
temperature=1.0,
boundary=32,
len_history_for_policy=32,
num_to_keep=1,
abort_fn=None,
raw_trajectory=False):
"""Plays the given env with the policy function to collect trajectories.
Args:
env: environment object, should be a subclass of env_problem.EnvProblem.
policy_fun: callable, taking in observations((B, RT) + OBS) and returning
back log-probabilities (B, AT, A).
num_trajectories: int, number of trajectories to collect.
max_timestep: int or None, if not None or a negative number, we cut any
trajectory that exceeds this time put it in the completed bin, and *dont*
reset the env.
reset: bool, true if we want to reset the envs. The envs are also reset if
max_max_timestep is None or < 0.
state: the state for `policy_fn`.
rng: jax rng, splittable.
temperature: float, temperature used in Gumbel sampling.
boundary: int, pad the sequences to the multiples of this number.
len_history_for_policy: int or None, the maximum history to keep for
applying the policy on. If None, use the whole history.
num_to_keep: int, while truncating trajectory how many time-steps to keep.
abort_fn: callable, If not None, then at every step call and abort the
trajectory collection if it returns True, if so reset the env and return
None.
raw_trajectory: bool, if True a list of trajectory.Trajectory objects is
returned, otherwise a list of numpy representations of
`trajectory.Trajectory` is returned.
Returns:
A tuple, (trajectories, number of completed trajectories). Where
trajectories is a list of triples of (observation, action, reward) ndarrays.
"""
def gumbel_sample(log_probs):
"""Gumbel sampling."""
u = np.random.uniform(low=1e-6, high=1.0 - 1e-6, size=log_probs.shape)
g = -np.log(-np.log(u))
return np.argmax((log_probs / temperature) + g, axis=-1)
# We need to reset all environments, if we're coming here the first time.
if reset or max_timestep is None or max_timestep <= 0:
env.reset()
else:
# Clear completed trajectories held internally.
env.trajectories.clear_completed_trajectories()
num_done_trajectories = 0
policy_application_total_time = 0
env_actions_total_time = 0
bare_env_run_time = 0
while env.trajectories.num_completed_trajectories < num_trajectories:
# Check if we should abort and return nothing.
if abort_fn and abort_fn():
# We should also reset the environment, since it will have some
# trajectories (complete and incomplete) that we want to discard.
env.reset()
return None, 0, {}, state
# Get all the observations for all the active trajectories.
# Shape is (B, RT) + OBS
# Bucket on whatever length is needed.
padded_observations, lengths = env.trajectories.observations_np(
boundary=boundary,
len_history_for_policy=len_history_for_policy)
B = padded_observations.shape[0] # pylint: disable=invalid-name
assert B == env.batch_size
assert (B,) == lengths.shape
t1 = time.time()
log_probs, value_preds, state, rng = policy_fun(
padded_observations, lengths, state=state, rng=rng)
policy_application_total_time += (time.time() - t1)
assert B == log_probs.shape[0]
actions = gumbel_sample(log_probs)
if isinstance(env.action_space, gym.spaces.Discrete):
actions = np.squeeze(actions, axis=1)
# Step through the env.
t1 = time.time()
_, _, dones, env_infos = env.step(
actions,
infos={
"log_prob_actions": log_probs,
"value_predictions": value_preds,
})
env_actions_total_time += (time.time() - t1)
bare_env_run_time += sum(
info["__bare_env_run_time__"] for info in env_infos)
# Count the number of done trajectories, the others could just have been
# truncated.
num_done_trajectories += np.sum(dones)
# Get the indices where we are done ...
done_idxs = done_indices(dones)
# ... and reset those.
t1 = time.time()
if done_idxs.size:
env.reset(indices=done_idxs)
env_actions_total_time += (time.time() - t1)
if max_timestep is None or max_timestep < 1:
continue
# Are there any trajectories that have exceeded the time-limit we want.
lengths = env.trajectories.trajectory_lengths
exceeded_time_limit_idxs = done_indices(lengths > max_timestep)
# If so, reset these as well.
t1 = time.time()
if exceeded_time_limit_idxs.size:
# This just cuts the trajectory, doesn't reset the env, so it continues
# from where it left off.
env.truncate(indices=exceeded_time_limit_idxs, num_to_keep=num_to_keep)
env_actions_total_time += (time.time() - t1)
# We have the trajectories we need, return a list of triples:
# (observations, actions, rewards)
completed_trajectories = get_completed_trajectories_from_env(
env, num_trajectories, raw_trajectory=raw_trajectory)
timing_info = {
"trajectory_collection/policy_application": policy_application_total_time,
"trajectory_collection/env_actions": env_actions_total_time,
"trajectory_collection/env_actions/bare_env": bare_env_run_time,
}
timing_info = {k: round(1000 * v, 2) for k, v in timing_info.items()}
return completed_trajectories, num_done_trajectories, timing_info, state
def make_env(batch_size=1,
env_problem_name="",
resize=True,
resize_dims=(105, 80),
max_timestep="None",
clip_rewards=True,
parallelism=1,
use_tpu=False,
num_actions=None,
rendered_env=True,
**env_kwargs):
"""Creates the env."""
if clip_rewards:
env_kwargs.update({"reward_range": (-1, 1), "discrete_rewards": True})
else:
env_kwargs.update({"discrete_rewards": False})
# TODO(henrykm) - below someone linked "resize" with "abnormality"
# Probably we need more nuanced concept of "abnormality"
# decoupled from "resize". Currently the resize flag implies
# that we switch from a generic env to a wrapped env.
# Overall this file and gym_utils.py look like good candidates
# for a refactor.
# No resizing needed, so let's be on the normal EnvProblem.
if not resize: # None or False
return gym_env_problem.GymEnvProblem(
base_env_name=env_problem_name,
batch_size=batch_size,
parallelism=parallelism,
**env_kwargs)
try:
max_timestep = int(max_timestep)
except Exception: # pylint: disable=broad-except
max_timestep = None
wrapper_fn = functools.partial(
gym_utils.gym_env_wrapper, **{
"rl_env_max_episode_steps": max_timestep,
"maxskip_env": True,
"rendered_env": rendered_env,
"rendered_env_resize_to": resize_dims,
"sticky_actions": False,
"output_dtype": np.int32 if use_tpu else None,
"num_actions": num_actions,
})
return rendered_env_problem.RenderedEnvProblem(
base_env_name=env_problem_name,
batch_size=batch_size,
parallelism=parallelism,
rendered_env=rendered_env,
env_wrapper_fn=wrapper_fn,
**env_kwargs)
================================================
FILE: tensor2tensor/envs/env_problem_utils_test.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Tests for env_problem_utils."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
from tensor2tensor.envs import env_problem_utils
from tensor2tensor.envs import gym_env_problem
from tensor2tensor.envs import tic_tac_toe_env # pylint: disable=unused-import
from tensor2tensor.envs import tic_tac_toe_env_problem
import tensorflow.compat.v1 as tf
class EnvProblemUtilsTest(tf.test.TestCase):
def test_play_env_problem_randomly(self):
batch_size = 5
num_steps = 100
ep = tic_tac_toe_env_problem.TicTacToeEnvProblem()
ep.initialize(batch_size=batch_size)
env_problem_utils.play_env_problem_randomly(ep, num_steps)
# We've played num_steps * batch_size steps + everytime we get 'done' we
# create another step + batch_size number of pending steps.
self.assertEqual(
num_steps * batch_size + len(ep.trajectories.completed_trajectories) +
batch_size, ep.trajectories.num_time_steps)
def test_play_env_problem_with_policy(self):
env = gym_env_problem.GymEnvProblem(
base_env_name="CartPole-v0", batch_size=2, reward_range=(-1, 1))
# Let's make sure that at-most 4 observations come to the policy function.
len_history_for_policy = 4
def policy_fun(observations, lengths, state=None, rng=None):
del lengths
b = observations.shape[0]
# Assert that observations from time-step len_history_for_policy onwards
# are zeros.
self.assertTrue(
np.all(observations[:, len_history_for_policy:, ...] == 0))
self.assertFalse(
np.all(observations[:, :len_history_for_policy, ...] == 0))
a = env.action_space.n
p = np.random.uniform(size=(b, 1, a))
p = np.exp(p)
p = p / np.sum(p, axis=-1, keepdims=True)
return np.log(p), np.mean(p, axis=-1), state, rng
max_timestep = 15
num_trajectories = 2
trajectories, _, _, _ = env_problem_utils.play_env_problem_with_policy(
env,
policy_fun,
num_trajectories=num_trajectories,
max_timestep=max_timestep,
len_history_for_policy=len_history_for_policy)
self.assertEqual(num_trajectories, len(trajectories))
# Check shapes within trajectories.
traj = trajectories[0]
T = traj[1].shape[0] # pylint: disable=invalid-name
self.assertEqual((T + 1, 4), traj[0].shape) # (4,) is OBS
self.assertEqual((T,), traj[2].shape)
self.assertEqual(T, len(traj[4]["log_prob_actions"]))
self.assertEqual(T, len(traj[4]["value_predictions"]))
self.assertLessEqual(T, max_timestep)
traj = trajectories[1]
T = traj[1].shape[0] # pylint: disable=invalid-name
self.assertEqual((T + 1, 4), traj[0].shape)
self.assertEqual((T,), traj[2].shape)
self.assertEqual(T, len(traj[4]["log_prob_actions"]))
self.assertEqual(T, len(traj[4]["value_predictions"]))
self.assertLessEqual(T, max_timestep)
if __name__ == "__main__":
tf.test.main()
================================================
FILE: tensor2tensor/envs/gym_env_problem.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Base class for envs that store their history.
EnvProblem subclasses Problem and also implements the Gym interface (step,
reset, render, close, seed)
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import copy
import multiprocessing.pool
import time
from absl import logging
import gym
import numpy as np
from tensor2tensor.envs import env_problem
from tensor2tensor.envs import trajectory
# This is a compatibility shim introduced to support NumPy 1.24. See:
# https://numpy.org/neps/nep-0034-infer-dtype-is-object.html
def _stack(xs):
try:
return np.stack(xs)
except ValueError:
return np.stack(np.asarray(xs, dtype=object))
class GymEnvProblem(env_problem.EnvProblem):
"""An EnvProblem implemented as a batch of gym envs.
This implementation should work well for cases where the env is not batched by
default ex: any gym env. In this case we create `batch_size` number of envs
and store them in a list. Any function then that interacts with the envs, like
reset, step or close goes over the env list to do the needful, ex: when reset
is called with specific indices we reset only those indices, etc.
The usage of this class will look like the following:
# 1. Creates and initializes the env_problem.
ep = env_problem.EnvProblem(...)
# 2. One needs to call reset() at the start, this resets all envs.
ep.reset()
# 3. Call step with actions for all envs, i.e. len(action) = batch_size
obs, rewards, dones, infos = ep.step(actions)
# 4. Figure out which envs got done and reset only those.
ep.reset(indices=env_problem_utils.done_indices(dones))
# 5. Go back to Step #3 to further interact with the env or just dump the
# generated data to disk by calling:
ep.generate_data(...)
# 6. If we now need to use this object again to play a few more iterations
# perhaps with a different batch size or maybe not recording the data, then
# we need to re-initialize environments and do some book-keeping, call:
ep.initialize_environments(batch_size)
# 7. Go back to Step #2, i.e. reset all envs.
NOTE: Look at `EnvProblemTest.test_interaction_with_env` and/or
`EnvProblemTest.test_generate_data`
NOTE: We rely heavily that the underlying environments expose a gym style
interface, i.e. in addition to reset(), step() and close() we have access to
the following properties: observation_space, action_space, reward_range.
"""
def __init__(self,
base_env_name=None,
env_wrapper_fn=None,
reward_range=None,
**kwargs):
"""Initializes this class by creating the envs and managing trajectories.
Args:
base_env_name: (string) passed to `gym.make` to make the underlying
environment.
env_wrapper_fn: (callable(env): env) Applies gym wrappers to the base
environment.
reward_range: (tuple(number, number) or None) the first element is the
minimum reward and the second is the maximum reward, used to clip and
process the raw reward in `process_rewards`. If None, this is inferred
from the inner environments.
**kwargs: (dict) Arguments passed to the base class.
"""
# Name for the base environment, will be used in `gym.make` in
# the default implementation of `initialize_environments`.
self._base_env_name = base_env_name
# An env generates data when it is given actions by an agent which is either
# a policy or a human -- this is supposed to be the `id` of the agent.
#
# In practice, this is used only to store (and possibly retrieve) history
# to an appropriate directory.
self._agent_id = "default"
# We clip rewards to this range before processing them further, as described
# in `process_rewards`.
self._reward_range = reward_range
# Initialize the environment(s).
# This can either be a list of environments of len `batch_size` or this can
# be a Neural Network, in which case it will be fed input with first
# dimension = `batch_size`.
self._envs = None
self._pool = None
self._env_wrapper_fn = env_wrapper_fn
# Call the super's ctor. It will use some of the member fields, so we call
# it in the end.
super(GymEnvProblem, self).__init__(**kwargs)
@property
def base_env_name(self):
return self._base_env_name
def _verify_same_spaces(self):
"""Verifies that all the envs have the same observation and action space."""
# Pre-conditions: self._envs is initialized.
if self._envs is None:
raise ValueError("Environments not initialized.")
if not isinstance(self._envs, list):
logging.warning("Not checking observation and action space "
"compatibility across envs, since there is just one.")
return
# NOTE: We compare string representations of observation_space and
# action_space because compositional classes like space.Tuple don't return
# true on object comparison.
if not all(
str(env.observation_space) == str(self.observation_space)
for env in self._envs):
err_str = ("All environments should have the same observation space, but "
"don't.")
logging.error(err_str)
# Log all observation spaces.
for i, env in enumerate(self._envs):
logging.error("Env[%d] has observation space [%s]", i,
env.observation_space)
raise ValueError(err_str)
if not all(
str(env.action_space) == str(self.action_space) for env in self._envs):
err_str = "All environments should have the same action space, but don't."
logging.error(err_str)
# Log all action spaces.
for i, env in enumerate(self._envs):
logging.error("Env[%d] has action space [%s]", i, env.action_space)
raise ValueError(err_str)
def initialize_environments(self,
batch_size=1,
parallelism=1,
per_env_kwargs=None,
**kwargs):
"""Initializes the environments.
Args:
batch_size: (int) Number of `self.base_env_name` envs to initialize.
parallelism: (int) If this is greater than one then we run the envs in
parallel using multi-threading.
per_env_kwargs: (list or None) An optional list of dictionaries to pass to
gym.make. If not None, length should match `batch_size`.
**kwargs: (dict) Kwargs to pass to gym.make.
"""
assert batch_size >= 1
if per_env_kwargs is not None:
assert batch_size == len(per_env_kwargs)
else:
per_env_kwargs = [{} for _ in range(batch_size)]
# By now `per_env_kwargs` is a list of dictionaries of size batch_size.
# The individual dictionaries maybe empty.
def union_dicts(dict1, dict2):
"""Union `dict1` and `dict2`."""
copy_dict1 = copy.copy(dict1)
copy_dict1.update(dict2)
return copy_dict1
self._envs = [
gym.make(self.base_env_name,
**union_dicts(kwargs, env_kwarg))
for env_kwarg in per_env_kwargs
]
self._parallelism = parallelism
self._pool = multiprocessing.pool.ThreadPool(self._parallelism)
if self._env_wrapper_fn is not None:
self._envs = list(map(self._env_wrapper_fn, self._envs))
self._verify_same_spaces()
# If self.reward_range is None, i.e. this means that we should take the
# reward range of the env.
if self.reward_range is None:
self._reward_range = self._envs[0].reward_range
# This data structure stores the history of each env.
#
# NOTE: Even if the env is a NN and can step in all batches concurrently, it
# is still valuable to store the trajectories separately.
self._trajectories = trajectory.BatchTrajectory(batch_size=batch_size)
def assert_common_preconditions(self):
# Asserts on the common pre-conditions of:
# - self._envs is initialized.
# - self._envs is a list.
assert self._envs
assert isinstance(self._envs, list)
@property
def observation_space(self):
return self._envs[0].observation_space
@property
def action_space(self):
return self._envs[0].action_space
@property
def reward_range(self):
return self._reward_range
def seed(self, seed=None):
if not self._envs:
logging.info("`seed` called on non-existent envs, doing nothing.")
return None
if not isinstance(self._envs, list):
logging.warning("`seed` called on non-list envs, doing nothing.")
return None
logging.warning(
"Called `seed` on EnvProblem, calling seed on the underlying envs.")
for env in self._envs:
env.seed(seed)
return super(GymEnvProblem, self).seed(seed=seed)
def close(self):
if not self._envs:
logging.info("`close` called on non-existent envs, doing nothing.")
return
if not isinstance(self._envs, list):
logging.warning("`close` called on non-list envs, doing nothing.")
return
# Call close on all the envs one by one.
for env in self._envs:
env.close()
def _reset(self, indices):
"""Resets environments at indices shouldn't pre-process or record.
Args:
indices: list of indices of underlying envs to call reset on.
Returns:
np.ndarray of stacked observations from the reset-ed envs.
"""
# This returns a numpy array with first dimension `len(indices)` and the
# rest being the dimensionality of the observation.
num_envs_to_reset = len(indices)
observations = [None] * num_envs_to_reset
def reset_at(idx):
observations[idx] = self._envs[indices[idx]].reset()
if self._parallelism > 1:
self._pool.map(reset_at, range(num_envs_to_reset))
else:
for i in range(num_envs_to_reset):
reset_at(i)
return _stack(observations)
def _step(self, actions):
"""Takes a step in all environments, shouldn't pre-process or record.
Args:
actions: (np.ndarray) with first dimension equal to the batch size.
Returns:
a tuple of stacked raw observations, raw rewards, dones and infos.
"""
assert len(actions) == len(self._envs)
observations = [None] * self.batch_size
rewards = [None] * self.batch_size
dones = [None] * self.batch_size
infos = [{} for _ in range(self.batch_size)]
def apply_step(i):
t1 = time.time()
observations[i], rewards[i], dones[i], infos[i] = self._envs[i].step(
actions[i])
t2 = time.time()
infos[i]["__bare_env_run_time__"] = t2 - t1
if self._parallelism > 1:
self._pool.map(apply_step, range(self.batch_size))
else:
for i in range(self.batch_size):
apply_step(i)
# Convert each list (observations, rewards, ...) into np.array and return a
# tuple.
return tuple(map(_stack, [observations, rewards, dones, infos]))
================================================
FILE: tensor2tensor/envs/gym_env_problem_test.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Tests for tensor2tensor.envs.gym_env_problem."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import gym
from gym.spaces import Box
from gym.spaces import Discrete
import numpy as np
from tensor2tensor.data_generators import generator_utils
from tensor2tensor.data_generators import problem
from tensor2tensor.envs import env_problem
from tensor2tensor.envs import env_problem_utils
from tensor2tensor.envs import gym_env_problem
from tensor2tensor.layers import modalities
import tensorflow.compat.v1 as tf
from tensorflow.compat.v1 import estimator as tf_estimator
class GymEnvProblemTest(tf.test.TestCase):
def setUp(self):
self.tmp_dir = os.path.join(tf.test.get_temp_dir(), "tmp_dir")
tf.gfile.MakeDirs(self.tmp_dir)
def tearDown(self):
tf.gfile.DeleteRecursively(self.tmp_dir)
def test_setup(self):
ep = gym_env_problem.GymEnvProblem(
base_env_name="CartPole-v0", batch_size=5)
# Checks that environments were created and they are `batch_size` in number.
ep.assert_common_preconditions()
# Expectations on the observation space.
observation_space = ep.observation_space
self.assertIsInstance(observation_space, Box)
self.assertEqual(observation_space.shape, (4,))
self.assertEqual(observation_space.dtype, np.float32)
# Expectations on the action space.
action_space = ep.action_space
self.assertTrue(isinstance(action_space, Discrete))
self.assertEqual(action_space.shape, ())
self.assertEqual(action_space.dtype, np.int64)
self.assertEqual(ep.num_actions, 2)
# Reward range is infinite here.
self.assertFalse(ep.is_reward_range_finite)
def test_reward_range(self):
# Passing reward_range=None means take the reward range of the underlying
# environment as the reward range.
ep = gym_env_problem.GymEnvProblem(
base_env_name="FrozenLake-v1", batch_size=5, reward_range=None)
ep.assert_common_preconditions()
# Assert reward range is finite here.
self.assertTrue(ep.is_reward_range_finite)
# Assert that it is as expected of the underlying environment, since reward_
self.assertEqual(0, ep.reward_range[0])
self.assertEqual(1, ep.reward_range[1])
def test_default_processed_rewards_discrete(self):
# This differs in the above because it has a Tuple observation space.
ep = gym_env_problem.GymEnvProblem(
base_env_name="KellyCoinflip-v0", batch_size=5, reward_range=None)
ep.assert_common_preconditions()
# Assert reward range is finite here.
self.assertTrue(ep.is_reward_range_finite)
# Assert that it is as expected of the underlying environment.
reward_range = ep.reward_range
self.assertEqual(0, reward_range[0])
# Google's version of Gym has maxWealth, vs max_wealth externally.
max_wealth = getattr(ep._envs[0], "maxWealth",
getattr(ep._envs[0], "max_wealth", None))
self.assertIsNotNone(max_wealth)
self.assertEqual(max_wealth, reward_range[1])
# Check that the processed rewards are discrete.
self.assertTrue(ep.is_processed_rewards_discrete)
# Assert on the number of rewards.
self.assertEqual(ep.num_rewards, reward_range[1] - reward_range[0] + 1)
def test_interaction_with_env(self):
batch_size = 5
reward_range = (-1, 1)
ep = gym_env_problem.GymEnvProblem(
base_env_name="KellyCoinflip-v0",
batch_size=batch_size,
reward_range=reward_range)
# Resets all environments.
ep.reset()
# Let's play a few steps.
nsteps = 100
num_trajectories_completed = 0
num_timesteps_completed = 0
# If batch_done_at_step[i] = j then it means that i^th env last got done at
# step = j.
batch_done_at_step = np.full(batch_size, -1)
for i in range(nsteps):
# Sample batch_size actions from the action space and stack them (since
# that is the expected type).
actions = np.stack([ep.action_space.sample() for _ in range(batch_size)])
_, _, dones, _ = ep.step(actions)
# Do the book-keeping on number of trajectories completed and expect that
# it matches ep's completed number.
num_done = sum(dones)
num_trajectories_completed += num_done
self.assertEqual(num_trajectories_completed,
len(ep.trajectories.completed_trajectories))
# Get the indices where we are done ...
done_indices = env_problem_utils.done_indices(dones)
# ... and reset those.
ep.reset(indices=done_indices)
# If nothing got done, go on to the next step.
if done_indices.size == 0:
# i.e. this is an empty array.
continue
# See when these indices were last done and calculate how many time-steps
# each one took to get done.
num_timesteps_completed += sum(i + 1 - batch_done_at_step[done_indices])
batch_done_at_step[done_indices] = i
# This should also match the number of time-steps completed given by ep.
num_timesteps_completed_ep = sum(
ct.num_time_steps for ct in ep.trajectories.completed_trajectories)
self.assertEqual(num_timesteps_completed, num_timesteps_completed_ep)
# Reset the trajectories.
ep.trajectories.reset_batch_trajectories()
self.assertEqual(0, len(ep.trajectories.completed_trajectories))
def read_tfrecord_dataset(self, filenames, ep):
# Read the dataset at `filenames` into a tf.data.Dataset and returns the
# number of time-steps (just the number of records in the dataset) and the
# number of trajectories.
last_timestep = -1
num_time_steps = 0
num_trajectories = 0
for ex in generator_utils.tfrecord_iterator(
filenames, example_spec=ep.example_reading_spec()[0]):
num_time_steps += 1
this_timestep = ex[env_problem.TIMESTEP_FIELD][0]
if 1 + last_timestep != this_timestep:
num_trajectories += 1
self.assertEqual(0, this_timestep)
last_timestep = this_timestep
num_trajectories += 1
return num_trajectories, num_time_steps
def play_env(self,
env=None,
nsteps=100,
base_env_name=None,
batch_size=5,
reward_range=None):
"""Creates `GymEnvProblem` with the given arguments and plays it randomly.
Args:
env: optional env.
nsteps: plays the env randomly for nsteps.
base_env_name: passed to GymEnvProblem's init.
batch_size: passed to GymEnvProblem's init.
reward_range: passed to GymEnvProblem's init.
Returns:
tuple of gym_env_problem, number of trajectories done,
number of trajectories done in the last step.
"""
if env is None:
env = gym_env_problem.GymEnvProblem(
base_env_name=base_env_name,
batch_size=batch_size,
reward_range=reward_range)
# Usually done by a registered subclass, we do this manually in the test.
env.name = base_env_name
# Reset all environments.
env.reset()
# Play for some steps to generate data.
num_dones = 0
num_dones_in_last_step = 0
for _ in range(nsteps):
# Sample actions.
actions = np.stack([env.action_space.sample() for _ in range(batch_size)])
# Step through it.
_, _, dones, _ = env.step(actions)
# Get the indices where we are done ...
done_indices = env_problem_utils.done_indices(dones)
# ... and reset those.
env.reset(indices=done_indices)
# count the number of dones we got, in this step and overall.
num_dones_in_last_step = sum(dones)
num_dones += num_dones_in_last_step
return env, num_dones, num_dones_in_last_step
def test_generate_data(self):
base_env_name = "CartPole-v0"
batch_size = 5
reward_range = (-1, 1)
nsteps = 100
ep, num_dones, num_dones_in_last_step = self.play_env(
base_env_name=base_env_name,
batch_size=batch_size,
reward_range=reward_range,
nsteps=nsteps)
# This is because every num_dones starts a new trajectory, and a further
# batch_size are active at the last step when we call generate_data, but
# the ones that got done in the last step (these have only one time-step in
# their trajectory) will be skipped.
expected_num_trajectories = num_dones + batch_size - num_dones_in_last_step
# Similar logic as above, nsteps * batch_size overall `step` calls are made.
expected_num_time_steps = (
nsteps * batch_size) + num_dones + batch_size - num_dones_in_last_step
# Dump the completed data to disk.
ep.generate_data(self.tmp_dir, self.tmp_dir)
# Read the written files and assert on the number of time steps.
training_filenames = ep.training_filepaths(
self.tmp_dir, ep.num_shards[problem.DatasetSplit.TRAIN], True)
dev_filenames = ep.dev_filepaths(
self.tmp_dir, ep.num_shards[problem.DatasetSplit.EVAL], True)
training_trajectories, training_timesteps = self.read_tfrecord_dataset(
training_filenames, ep)
dev_trajectories, dev_timesteps = self.read_tfrecord_dataset(
dev_filenames, ep)
# This tests what we wrote on disk matches with what we computed.
self.assertEqual(expected_num_time_steps,
training_timesteps + dev_timesteps)
self.assertEqual(expected_num_trajectories,
training_trajectories + dev_trajectories)
def test_problem_dataset_works(self):
# We need to derive this class to set the required methods.
class TestEnv(gym_env_problem.GymEnvProblem):
name = "TestEnv"
@property
def input_modality(self):
return modalities.ModalityType.REAL_L2_LOSS
@property
def input_vocab_size(self):
return None
@property
def target_modality(self):
return modalities.ModalityType.SYMBOL_WEIGHTS_ALL
@property
def target_vocab_size(self):
return 2
@property
def action_modality(self):
return modalities.ModalityType.SYMBOL_WEIGHTS_ALL
@property
def reward_modality(self):
return modalities.ModalityType.SYMBOL_WEIGHTS_ALL
base_env_name = "CartPole-v0"
batch_size = 5
reward_range = (-1, 1)
env = TestEnv(
base_env_name=base_env_name,
batch_size=batch_size,
reward_range=reward_range)
nsteps = 100
ep, _, _ = self.play_env(env=env, nsteps=nsteps)
# Dump the completed data to disk.
ep.generate_data(self.tmp_dir, self.tmp_dir)
# Read the actual files and count the trajectories and time-steps.
dev_filenames = ep.dev_filepaths(
self.tmp_dir, ep.num_shards[problem.DatasetSplit.EVAL], True)
dev_trajectories, dev_timesteps = self.read_tfrecord_dataset(
dev_filenames, ep)
# Count them using a tf.data.Dataset.
dev_dataset = ep.dataset(tf_estimator.ModeKeys.EVAL, data_dir=self.tmp_dir)
last_timestep = -1
dev_timesteps_ds = 0
dev_trajectories_ds = 0
iterator = dev_dataset.make_one_shot_iterator()
next_element = iterator.get_next()
with tf.Session() as session:
while True:
try:
tf_example_dict = session.run(next_element)
# We have a time-step.
dev_timesteps_ds += 1
this_timestep = tf_example_dict[env_problem.TIMESTEP_FIELD][
0] # [0] since every value in tf_example_dict is an array/list.
if 1 + last_timestep != this_timestep:
dev_trajectories_ds += 1
self.assertEqual(0, this_timestep)
last_timestep = this_timestep
except tf.errors.OutOfRangeError:
dev_trajectories_ds += 1
break
# Make sure that they agree.
self.assertEqual(dev_trajectories, dev_trajectories_ds)
self.assertEqual(dev_timesteps, dev_timesteps_ds)
def test_resets_properly(self):
base_env_name = "CartPole-v0"
batch_size = 5
reward_range = (-1, 1)
nsteps = 100
env = gym_env_problem.GymEnvProblem(
base_env_name=base_env_name,
batch_size=batch_size,
reward_range=reward_range)
env.name = base_env_name
num_dones = 0
while num_dones == 0:
env, num_dones, _ = self.play_env(env=env,
nsteps=nsteps,
batch_size=batch_size,
reward_range=reward_range)
# Some completed trajectories have been generated.
self.assertGreater(env.trajectories.num_completed_trajectories, 0)
# This should clear the env completely of any state.
env.reset()
# Assert that there aren't any completed trajectories in the env now.
self.assertEqual(env.trajectories.num_completed_trajectories, 0)
def test_per_env_kwargs(self):
# Creating a dummy class where we specify the action at which the env
# returns done.
class TestPerEnvKwargsEnv(gym.Env):
"""Test environment with the `done action` specified."""
action_space = Discrete(3)
observation_space = Box(low=-1.0, high=1.0, shape=())
def __init__(self, done_action=0):
self._done_action = done_action
def _generate_ob(self):
return self.observation_space.sample()
def step(self, action):
done = self._done_action == action
reward = 1 if done else 0
return (self._generate_ob(), reward, done, {})
def reset(self):
return self._generate_ob()
# Registering it with gym.
test_env_name = "TestPerEnvKwargsEnv-v0"
gym.envs.register(id=test_env_name, entry_point=TestPerEnvKwargsEnv)
# Creating a batch of those with different done actions.
base_env_name = test_env_name
batch_size = 2
reward_range = (-1, 1)
per_env_kwargs = [{"done_action": 1}, {"done_action": 2}]
env = gym_env_problem.GymEnvProblem(
base_env_name=base_env_name,
batch_size=batch_size,
reward_range=reward_range,
per_env_kwargs=per_env_kwargs)
_ = env.reset()
# Finally querying the done actions.
_, _, d, _ = env.step(np.array([0, 0]))
self.assertFalse(d[0])
self.assertFalse(d[1])
_, _, d, _ = env.step(np.array([1, 2]))
self.assertTrue(d[0])
self.assertTrue(d[1])
if __name__ == "__main__":
tf.test.main()
================================================
FILE: tensor2tensor/envs/gym_spaces_utils.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Few utility functions to deal with gym spaces.
gym.spaces.Box and gym.spaces.Discrete are easiest to support.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from gym.spaces import Box
from gym.spaces import Discrete
import numpy as np
import tensorflow.compat.v1 as tf
def box_space_spec(box_space, tf_dtype):
return tf.FixedLenFeature(box_space.shape, tf_dtype)
def discrete_space_spec(discrete_space, tf_dtype):
del discrete_space # this is not needed.
return tf.FixedLenFeature((1,), tf_dtype)
def gym_space_spec(gym_space):
"""Returns a reading spec of a gym space.
NOTE: Only implemented currently for Box and Discrete.
Args:
gym_space: instance of gym.spaces whose spec we want.
Returns:
Reading spec for that space.
Raises:
NotImplementedError: For spaces whose reading spec we haven't implemented.
"""
# First try to determine the type.
try:
tf_dtype = tf.as_dtype(gym_space.dtype)
except TypeError as e:
tf.logging.error("Cannot convert space's type [%s] to tf.dtype",
gym_space.dtype)
raise e
# Now hand it over to the specialized functions.
if isinstance(gym_space, Box):
return box_space_spec(gym_space, tf_dtype)
elif isinstance(gym_space, Discrete):
return discrete_space_spec(gym_space, tf_dtype)
else:
raise NotImplementedError
def gym_space_encode(gym_space, observation):
# We should return something that generator_utils.to_example can consume.
if isinstance(gym_space, Discrete):
return [observation]
if isinstance(gym_space, Box):
return observation.reshape(-1).tolist()
raise NotImplementedError
def cardinality(gym_space):
"""Number of elements that can be represented by the space.
Makes the most sense for Discrete or Box type with integral dtype, ex: number
of actions in an action space.
Args:
gym_space: The gym space.
Returns:
np.int64 number of observations that can be represented by this space, or
returns None when this doesn't make sense, i.e. float boxes etc.
Raises:
NotImplementedError when a space's cardinality makes sense but we haven't
implemented it.
"""
if (gym_space.dtype == np.float32) or (gym_space.dtype == np.float64):
tf.logging.warn("Returning None for a float gym space's cardinality: %s",
gym_space)
return None
if isinstance(gym_space, Discrete):
return gym_space.n
if isinstance(gym_space, Box):
# Construct a box with all possible values in this box and take a product.
return np.prod(gym_space.high - gym_space.low + 1)
raise NotImplementedError
================================================
FILE: tensor2tensor/envs/gym_spaces_utils_test.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Tests for gym_spaces_utils.py."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from gym.spaces import Box
from gym.spaces import Discrete
import numpy as np
from tensor2tensor.envs import gym_spaces_utils
import tensorflow.compat.v1 as tf
class GymSpacesUtilsTest(tf.test.TestCase):
def test_discrete_space_spec(self):
discrete_space = Discrete(100)
spec = gym_spaces_utils.gym_space_spec(discrete_space)
self.assertIsInstance(spec, tf.FixedLenFeature)
self.assertEqual(spec.dtype, tf.int64)
self.assertListEqual(list(spec.shape), [1])
def test_box_space_spec(self):
box_space = Box(low=0, high=10, shape=[5, 6], dtype=np.float32)
spec = gym_spaces_utils.gym_space_spec(box_space)
self.assertIsInstance(spec, tf.FixedLenFeature)
self.assertEqual(spec.dtype, tf.float32)
self.assertListEqual(list(spec.shape), [5, 6])
def test_discrete_space_encode(self):
discrete_space = Discrete(100)
value = discrete_space.sample()
encoded_value = gym_spaces_utils.gym_space_encode(discrete_space, value)
self.assertListEqual([value], encoded_value)
def test_box_space_encode(self):
box_space = Box(low=0, high=10, shape=[2], dtype=np.int64)
value = np.array([2, 3])
encoded_value = gym_spaces_utils.gym_space_encode(box_space, value)
self.assertListEqual([2, 3], encoded_value)
if __name__ == '__main__':
tf.test.main()
================================================
FILE: tensor2tensor/envs/mujoco_problems.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Mujoco Gym environments."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import functools
from tensor2tensor.envs import rendered_env_problem
from tensor2tensor.layers import modalities
from tensor2tensor.rl import gym_utils
from tensor2tensor.utils import registry
@registry.register_env_problem
class ReacherEnvProblem(rendered_env_problem.RenderedEnvProblem):
"""Mujoco's reacher environment."""
def __init__(self):
base_env_name = "Reacher-v2"
wrapper_fn = functools.partial(
gym_utils.gym_env_wrapper, **{
"rl_env_max_episode_steps": -1,
"maxskip_env": False,
"rendered_env": True,
"rendered_env_resize_to": None, # Do not resize frames
"sticky_actions": False,
"output_dtype": None,
"num_actions": None,
})
super(ReacherEnvProblem, self).__init__(
base_env_name=base_env_name, env_wrapper_fn=wrapper_fn)
@property
def input_modality(self):
return modalities.ModalityType.VIDEO
@property
def target_modality(self):
return modalities.ModalityType.VIDEO
@property
def action_modality(self):
return modalities.ModalityType.IDENTITY
@property
def reward_modality(self):
return modalities.ModalityType.IDENTITY
@property
def input_vocab_size(self):
return 256
@property
def target_vocab_size(self):
return 256
================================================
FILE: tensor2tensor/envs/mujoco_problems_test.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Tests for tensor2tensor.envs.mujoco_problems."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
from tensor2tensor.envs import env_problem_utils
from tensor2tensor.envs import mujoco_problems # pylint: disable=unused-import
from tensor2tensor.utils import registry
import tensorflow.compat.v1 as tf
class ReacherEnvProblemTest(tf.test.TestCase):
def test_registration_and_interaction_with_env_problem(self):
batch_size = 5
# This ensures that registration has occurred.
ep = registry.env_problem("reacher_env_problem", batch_size=batch_size)
ep.reset()
num_done = 0
nsteps = 100
for _ in range(nsteps):
actions = np.stack([ep.action_space.sample() for _ in range(batch_size)])
obs, rewards, dones, infos = ep.step(actions)
# Assert that things are happening batchwise.
self.assertEqual(batch_size, len(obs))
self.assertEqual(batch_size, len(rewards))
self.assertEqual(batch_size, len(dones))
self.assertEqual(batch_size, len(infos))
done_indices = env_problem_utils.done_indices(dones)
ep.reset(done_indices)
num_done += sum(dones)
# Assert that something got done atleast,
self.assertGreater(num_done, 0)
if __name__ == "__main__":
tf.test.main()
================================================
FILE: tensor2tensor/envs/rendered_env_problem.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Base class for env problems with RGB array as observation space."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
import png
import six
from tensor2tensor.data_generators import video_utils
from tensor2tensor.envs import env_problem
from tensor2tensor.envs import gym_env_problem
from tensor2tensor.utils import contrib
import tensorflow.compat.v1 as tf
_IMAGE_ENCODED_FIELD = "image/encoded"
_IMAGE_FORMAT_FIELD = "image/format"
_IMAGE_HEIGHT_FIELD = "image/height"
_IMAGE_WIDTH_FIELD = "image/width"
_FRAME_NUMBER_FIELD = "frame_number"
_FORMAT = "png"
class RenderedEnvProblem(gym_env_problem.GymEnvProblem,
video_utils.VideoProblem):
"""An `EnvProblem` when observations are RGB arrays.
This takes care of wrapping a rendered gym environment to behave like a
`VideoProblem`. This class assumes the underlying gym environment is either a
`gym_utils.RenderedEnv` or it natively returns rendered scene for
observations. i.e. the underlying gym environment should have a
`Box` observation space with the following shape: [frame_height, frame_width,
channels]
Note: The method resolution order for this class is:
`RenderedEnvProblem`, `EnvProblem`, `Env`, `VideoProblem`, `Problem`
"""
def __init__(self, *args, **kwargs):
"""Initialize by calling both parents' constructors."""
gym_env_problem.GymEnvProblem.__init__(self, *args, **kwargs)
video_utils.VideoProblem.__init__(self)
def initialize_environments(self,
batch_size=1,
parallelism=1,
rendered_env=True,
per_env_kwargs=None,
**kwargs):
gym_env_problem.GymEnvProblem.initialize_environments(
self, batch_size=batch_size, parallelism=parallelism,
per_env_kwargs=per_env_kwargs, **kwargs)
# Assert the underlying gym environment has correct observation space
if rendered_env:
assert len(self.observation_spec.shape) == 3
def example_reading_spec(self):
"""Return a mix of env and video data fields and decoders."""
slim = contrib.slim()
video_fields, video_decoders = (
video_utils.VideoProblem.example_reading_spec(self))
env_fields, env_decoders = (
gym_env_problem.GymEnvProblem.example_reading_spec(self))
# Remove raw observations field since we want to capture them as videos.
env_fields.pop(env_problem.OBSERVATION_FIELD)
env_decoders.pop(env_problem.OBSERVATION_FIELD)
# Add frame number spec and decoder.
env_fields[_FRAME_NUMBER_FIELD] = tf.FixedLenFeature((1,), tf.int64)
env_decoders[_FRAME_NUMBER_FIELD] = slim.tfexample_decoder.Tensor(
_FRAME_NUMBER_FIELD)
# Add video fields and decoders
env_fields.update(video_fields)
env_decoders.update(video_decoders)
return env_fields, env_decoders
def _generate_time_steps(self, trajectory_list):
"""Transforms time step observations to frames of a video."""
for time_step in gym_env_problem.GymEnvProblem._generate_time_steps(
self, trajectory_list):
# Convert the rendered observations from numpy to png format.
frame_np = np.array(time_step.pop(env_problem.OBSERVATION_FIELD))
frame_np = frame_np.reshape(
[self.frame_height, self.frame_width, self.num_channels])
# TODO(msaffar) Add support for non RGB rendered environments
frame = png.from_array(frame_np, "RGB", info={"bitdepth": 8})
frame_buffer = six.BytesIO()
frame.save(frame_buffer)
# Put the encoded frame back.
time_step[_IMAGE_ENCODED_FIELD] = [frame_buffer.getvalue()]
time_step[_IMAGE_FORMAT_FIELD] = [_FORMAT]
time_step[_IMAGE_HEIGHT_FIELD] = [self.frame_height]
time_step[_IMAGE_WIDTH_FIELD] = [self.frame_width]
# Add the frame number
time_step[_FRAME_NUMBER_FIELD] = time_step[env_problem.TIMESTEP_FIELD]
yield time_step
@property
def num_channels(self):
return self.observation_spec.shape[2]
@property
def frame_height(self):
return self.observation_spec.shape[0]
@property
def frame_width(self):
return self.observation_spec.shape[1]
@property
def total_number_of_frames(self):
"""Upper bound on the total number of frames across all environments.
This is used to decide sharding. See `VideoProblem.total_number_of_frames`
for more details.
Returns:
number of frames among all examples in the dataset.
"""
return self.trajectories.num_time_steps
================================================
FILE: tensor2tensor/envs/rendered_env_problem_test.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Tests for tensor2tensor.envs.rendered_env_problem."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from tensor2tensor.envs import env_problem
from tensor2tensor.envs import env_problem_utils
from tensor2tensor.envs import rendered_env_problem
from tensor2tensor.envs.mujoco_problems import ReacherEnvProblem
import tensorflow.compat.v1 as tf
class RenderedEnvProblemTest(tf.test.TestCase):
def test_generate_timesteps(self):
env = ReacherEnvProblem()
env.initialize(batch_size=2)
env_problem_utils.play_env_problem_randomly(env, num_steps=5)
env.trajectories.complete_all_trajectories()
frame_number = 0
for time_step in env._generate_time_steps(
env.trajectories.completed_trajectories):
# original observation should not be in time_step
self.assertNotIn(env_problem.OBSERVATION_FIELD, time_step)
# validate frame
self.assertIn(rendered_env_problem._IMAGE_ENCODED_FIELD, time_step)
self.assertIn(rendered_env_problem._IMAGE_HEIGHT_FIELD, time_step)
self.assertIn(rendered_env_problem._IMAGE_WIDTH_FIELD, time_step)
self.assertIn(rendered_env_problem._IMAGE_FORMAT_FIELD, time_step)
self.assertIn(rendered_env_problem._FRAME_NUMBER_FIELD, time_step)
decoded_frame = tf.image.decode_png(
time_step[rendered_env_problem._IMAGE_ENCODED_FIELD][0])
decoded_frame = self.evaluate(decoded_frame)
self.assertListEqual(
[env.frame_height, env.frame_width, env.num_channels],
list(decoded_frame.shape))
self.assertListEqual([rendered_env_problem._FORMAT],
time_step[rendered_env_problem._IMAGE_FORMAT_FIELD])
self.assertListEqual([frame_number],
time_step[rendered_env_problem._FRAME_NUMBER_FIELD])
self.assertListEqual([env.frame_width],
time_step[rendered_env_problem._IMAGE_WIDTH_FIELD])
self.assertListEqual([env.frame_height],
time_step[rendered_env_problem._IMAGE_HEIGHT_FIELD])
frame_number += 1
frame_number %= 6
if __name__ == "__main__":
tf.test.main()
================================================
FILE: tensor2tensor/envs/tic_tac_toe_env.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Gym Tic-Tac-Toe environment.
Environment acts like the second player and first player is either environment
or the agent. The environment follows a random policy for now.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import gym
from gym import spaces
from gym.utils import seeding
import numpy as np
from tensor2tensor.data_generators import problem
from tensor2tensor.layers import modalities
from tensor2tensor.rl import gym_utils
def encode_pos(i, j):
"""Encodes a pair (i, j) as a scalar position on the board."""
return 3 * i + j
def decode_pos(pos):
"""Decoes a scalar position on the board as a pair (i, j)."""
return pos // 3, pos % 3
def get_open_spaces(board):
"""Given a representation of the board, returns a list of open spaces."""
open_spaces = []
for i in range(3):
for j in range(3):
if board[i][j] == 0:
open_spaces.append(encode_pos(i, j))
return open_spaces
def get_reward_and_done(board):
"""Given a representation of the board, returns reward and done."""
# Returns (reward, done) where:
# reward: -1 means lost, +1 means win, 0 means draw or continuing.
# done: True if the game is over, i.e. someone won or it is a draw.
# Sum all rows ...
all_sums = [np.sum(board[i, :]) for i in range(3)]
# ... all columns
all_sums.extend([np.sum(board[:, i]) for i in range(3)])
# and both diagonals.
all_sums.append(np.sum([board[i, i] for i in range(3)]))
all_sums.append(np.sum([board[i, 2 - i] for i in range(3)]))
if -3 in all_sums:
return -1, True
if 3 in all_sums:
return 1, True
done = True
if get_open_spaces(board):
done = False
return 0, done
# TODO(afrozm): This should eventually subclass Problem.
class TicTacToeEnv(gym.Env):
"""Simple TicTacToe Env, starts the game randomly half of the time."""
def __init__(self, strict=False):
self.strict = strict
# What about metadata and spec?
self.reward_range = (-1.0, 1.0)
# Action space -- 9 positions that we can chose to mark.
self.action_space = spaces.Discrete(9)
# Observation space -- this hopefully does what we need.
self.observation_space = spaces.Box(
low=-1, high=1, shape=(3, 3), dtype=np.int64)
# Set the seed.
self.np_random = None
self.seed()
# Start the game.
self.board_state = None
self.done = False
self.reset()
def seed(self, seed=None):
self.np_random, seed = seeding.np_random(seed)
return [seed]
# TODO(afrozm): Parametrize by some policy so that the env plays in an optimal
# way.
def play_random_move(self):
# Select open spaces.
open_spaces = get_open_spaces(self.board_state)
if not open_spaces:
return False
# Choose a space and mark it.
pos = self.np_random.choice(open_spaces)
i, j = decode_pos(pos)
self.board_state[i, j] = -1
def reset(self):
self.board_state = np.zeros((3, 3), dtype=np.int64)
# We"ll start with a 50% chance.
if self.np_random.choice([0, 1]) == 0:
self.play_random_move()
# Return the observation.
return self.board_state
def render(self, mode="human"):
# Unused.
del mode
board_str = ""
for i in range(3):
for j in range(3):
pos = self.board_state[i, j]
if pos == -1:
board_str += "x"
elif pos == 0:
board_str += "-"
else:
board_str += "o"
board_str += "\n"
return board_str
def step(self, action):
# Are we already done?
if self.strict:
assert not self.done
# Action has to belong to the action state.
assert self.action_space.contains(action)
# Is it a legitimate move, i.e. is that position open to play?
is_legit_move = action in get_open_spaces(self.board_state)
# Shouldn"t be an illegal action -- is a noop if not strict.
if self.strict:
assert is_legit_move
# If strict mode is off, then let this be a noop and env not play either.
if not is_legit_move:
return self.board_state, 0, False, {}
# This is a legit move, perform the action and check if done, etc etc.
i, j = decode_pos(action)
self.board_state[i, j] = 1
reward, done = get_reward_and_done(self.board_state)
if done:
self.done = True
return self.board_state, reward, True, {}
# If not done already, play our move.
self.play_random_move()
reward, done = get_reward_and_done(self.board_state)
self.done = done
return self.board_state, reward, self.done, {}
def hparams(self, defaults, unused_model_hparams):
p = defaults
p.modality = {
"inputs": modalities.ModalityType.IDENTITY_SYMBOL,
"targets": modalities.ModalityType.IDENTITY_SYMBOL,
}
p.vocab_size = {
"inputs": 3, # since at each box, the input is either x, o or -.
# nevermind that we have a 3x3 box.
"targets": 3, # -1, 0, 1
}
p.input_space_id = 0 # problem.SpaceID.GENERIC
p.target_space_id = 0 # problem.SpaceID.GENERIC
# TODO(afrozm): Figure out how to get rid of this.
class DummyPolicyProblemTTT(problem.Problem):
"""Dummy Problem for running the policy."""
def __init__(self):
super(DummyPolicyProblemTTT, self).__init__()
self._ttt_env = TicTacToeEnv()
def hparams(self, defaults, model_hparams):
# Update the env's hparams.
self._ttt_env.hparams(defaults, model_hparams)
# Do these belong here?
defaults.modality.update({
"input_action": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
"input_reward": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
"target_action": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
"target_reward": modalities.ModalityType.SYMBOL_WEIGHTS_ALL,
"target_policy": modalities.ModalityType.IDENTITY,
"target_value": modalities.ModalityType.IDENTITY,
})
defaults.vocab_size.update({
"input_action": self.num_actions,
"input_reward": 3, # -1, 0, +1 ?
"target_action": self.num_actions,
"target_reward": 3, # -1, 0, +1 ?
"target_policy": None,
"target_value": None,
})
@property
def num_actions(self):
return self._ttt_env.action_space.n
def register():
# Register this with gym.
unused_tictactoe_id, unused_tictactoe_env = gym_utils.register_gym_env(
"tensor2tensor.envs.tic_tac_toe_env:TicTacToeEnv", version="v0")
# TODO(afrozm): Fix the registration and make it automatic.
register()
================================================
FILE: tensor2tensor/envs/tic_tac_toe_env_problem.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""TicTacToeEnvProblem wraps the TicTacToeEnv in an EnvProblem."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from tensor2tensor.envs import gym_env_problem
from tensor2tensor.layers import modalities
from tensor2tensor.utils import registry
@registry.register_env_problem
class TicTacToeEnvProblem(gym_env_problem.GymEnvProblem):
"""Plays `batch_size` games of tic-tac-toe."""
def __init__(self):
super(TicTacToeEnvProblem, self).__init__(
base_env_name="T2TEnv-TicTacToeEnv-v0",
reward_range=(-1, 1))
@property
def input_modality(self):
return modalities.ModalityType.IDENTITY_SYMBOL
@property
def input_vocab_size(self):
# Since a box can be either x or o or empty.
return 3
@property
def target_modality(self):
return modalities.ModalityType.IDENTITY_SYMBOL
@property
def target_vocab_size(self):
# Since reward is either -1 or 0 or +1.
return 3
@property
def action_modality(self):
return modalities.ModalityType.SYMBOL_WEIGHTS_ALL
================================================
FILE: tensor2tensor/envs/tic_tac_toe_env_problem_test.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Tests for tensor2tensor.envs.tic_tac_toe_env_problem."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
from tensor2tensor.envs import env_problem_utils
from tensor2tensor.envs import tic_tac_toe_env # pylint: disable=unused-import
from tensor2tensor.envs import tic_tac_toe_env_problem # pylint: disable=unused-import
from tensor2tensor.utils import registry
import tensorflow.compat.v1 as tf
class TicTacToeEnvProblemTest(tf.test.TestCase):
def test_registration_and_interaction_with_env_problem(self):
batch_size = 5
# This ensures that registration has occurred.
ep = registry.env_problem("tic_tac_toe_env_problem", batch_size=batch_size)
ep.reset()
num_done, num_lost, num_won, num_draw = 0, 0, 0, 0
nsteps = 100
for _ in range(nsteps):
actions = np.stack([ep.action_space.sample() for _ in range(batch_size)])
obs, rewards, dones, infos = ep.step(actions)
# Assert that things are happening batchwise.
self.assertEqual(batch_size, len(obs))
self.assertEqual(batch_size, len(rewards))
self.assertEqual(batch_size, len(dones))
self.assertEqual(batch_size, len(infos))
done_indices = env_problem_utils.done_indices(dones)
ep.reset(done_indices)
num_done += sum(dones)
for r, d in zip(rewards, dones):
if not d:
continue
if r == -1:
num_lost += 1
elif r == 0:
num_draw += 1
elif r == 1:
num_won += 1
else:
raise ValueError("reward should be -1, 0, 1 but is {}".format(r))
# Assert that something got done atleast, without that the next assert is
# meaningless.
self.assertGreater(num_done, 0)
# Assert that things are consistent.
self.assertEqual(num_done, num_won + num_lost + num_draw)
if __name__ == "__main__":
tf.test.main()
================================================
FILE: tensor2tensor/envs/tic_tac_toe_env_test.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Tests for tensor2tensor.envs.tic_tac_toe_env."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from tensor2tensor.envs import tic_tac_toe_env as ttt_env
import tensorflow.compat.v1 as tf
class TicTacToeEnvTest(tf.test.TestCase):
def test_start(self):
ttt = ttt_env.TicTacToeEnv(strict=True)
self.assertFalse(ttt.done)
# At max one move may have been played by the env.
spaces = ttt_env.get_open_spaces(ttt.board_state)
num_open_spaces = len(spaces)
# i.e. either 8 or 9
self.assertGreater(num_open_spaces, 7)
# Play a move
observation, reward, done, unused_info = ttt.step(spaces[0])
# The environment should also have played a move.
spaces = ttt_env.get_open_spaces(observation)
self.assertEqual(num_open_spaces - 2, len(spaces))
# Since at-max 3 moves have been played, the game can't end.
self.assertEqual(reward, 0)
self.assertFalse(done)
def test_env_actions(self):
# Environment keeps taking actions and not us, we should eventually lose.
ttt = ttt_env.TicTacToeEnv(strict=True)
for _ in range(9):
ttt.play_random_move()
if ttt.done:
break
reward, done = ttt_env.get_reward_and_done(ttt.board_state)
self.assertEqual(-1, reward)
self.assertTrue(done)
def test_keep_playing(self):
ttt = ttt_env.TicTacToeEnv(strict=False)
done = False
while not done:
# sample an action from the action space.
action = ttt.action_space.sample()
# play it -- could be a no-op since we don't see if positions are empty.
unused_observation, reward, done, unused_info = ttt.step(action)
# done is True, so either:
# we won
# env won or
# no space left
we_won = reward == 1
env_won = reward == -1
space = bool(ttt_env.get_open_spaces(ttt.board_state))
self.assertTrue(we_won or env_won or not space)
if __name__ == '__main__':
tf.test.main()
================================================
FILE: tensor2tensor/envs/time_step.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""TimeStep is a simple class that holds the information seen at a time-step.
Let:
r_t = Reward(s_{t-1}, a_{t-1}, s_t) - reward for getting into a state.
d_t = Done(s_t) - is this state terminal.
a_t = Action performed at state s_t
i_t = (optional) Dictionary of key, value pairs of miscellaneous data.
Then the sequence of states, actions and rewards looks like the following:
s0, a0/i0 s1/r1/d1, a1/i1 s2/r2/d2, a2/i2 s3/r3/d3, ...
TimeStep holds (s_t, d_t, r_t, a_t, i_t).
NOTE: When we call step on an environment at time-step t, we supply a_t and in
return the env gives us s_{t+1}, d_{t+1}, r_{t+1}
So, we'd have to add the actions a_t/i_t to the current time-step, but add the
observations, rewards and dones to a new time-step.
NOTE: wrt `info` - A good solution could be to have two additional fields in
TimeStep - structured algo_info (a namedtuple, possibly different for every
algorithm, or None if we don't use any) and unstructured env_info (a dict).))
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import collections
class TimeStep(
collections.namedtuple(
"TimeStep",
["observation", "done", "raw_reward", "processed_reward", "action",
"info"])):
"""This class represents the time-step as mentioned above."""
def replace(self, **kwargs):
"""Exposes the underlying namedtuple replace."""
# NOTE: This RETURNS a NEW time-step with the replacements, i.e. doesn't
# modify self, since namedtuple is immutable.
# This allows this to be called like ts.replace(action=a, raw_reward=r) etc.
return self._replace(**kwargs)
@classmethod
def create_time_step(cls,
observation=None,
done=False,
raw_reward=None,
processed_reward=None,
action=None,
info=None):
"""Creates a TimeStep with both rewards and actions as optional."""
return cls(observation, done, raw_reward, processed_reward, action,
info)
================================================
FILE: tensor2tensor/envs/time_step_test.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Tests for tensor2tensor.envs.time_step."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from tensor2tensor.envs import time_step
import tensorflow.compat.v1 as tf
class TimeStepTest(tf.test.TestCase):
def test_create_time_step(self):
ts = time_step.TimeStep.create_time_step(
observation=1, done=True, raw_reward=1.0, processed_reward=1, action=1,
info={1: 1, 2: 4})
self.assertEqual(1, ts.observation)
self.assertTrue(ts.done)
self.assertNear(1.0, ts.raw_reward, 1e-6)
self.assertEqual(1, ts.processed_reward)
self.assertEqual(1, ts.action)
self.assertEqual({1: 1, 2: 4}, ts.info)
def test_replace(self):
ts = time_step.TimeStep.create_time_step(observation=1, action=1)
self.assertFalse(ts.done)
tsr = ts.replace(action=2, done=True, info={1: 1, 2: 4})
# Asert that ts didn't change.
self.assertFalse(ts.done)
self.assertEqual(1, ts.observation)
self.assertEqual(1, ts.action)
# But tsr is as expected.
self.assertTrue(tsr.done)
self.assertEqual(1, tsr.observation) # unchanged
self.assertEqual(2, tsr.action) # changed
self.assertEqual({1: 1, 2: 4}, tsr.info)
if __name__ == '__main__':
tf.test.main()
================================================
FILE: tensor2tensor/envs/trajectory.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Trajectory manages a sequence of TimeSteps.
BatchTrajectory manages a batch of trajectories, also keeping account of
completed trajectories.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import pickle
import re
import sys
import time
from absl import logging
import cloudpickle
import numpy as np
from tensor2tensor.envs import time_step
import tensorflow.compat.v1 as tf
TRAJECTORY_FILE_FORMAT = r"trajectory_epoch_{epoch}_env_id_{env_id}_temperature_{temperature}_r_{r}.pkl"
def get_pickle_module():
if sys.version_info[0] < 3:
return cloudpickle
return pickle
class Trajectory(object):
"""Basically a list of TimeSteps with convenience methods."""
def __init__(self, time_steps=None):
# Contains a list of time steps.
if time_steps is None:
self._time_steps = []
else:
self._time_steps = time_steps
def __str__(self):
if not self.time_steps:
return "Trajectory[]"
return "Trajectory[{}]".format(", ".join(str(ts) for ts in self.time_steps))
def add_time_step(self, **create_time_step_kwargs):
"""Creates a time-step and appends it to the list.
Args:
**create_time_step_kwargs: Forwarded to
time_step.TimeStep.create_time_step.
"""
ts = time_step.TimeStep.create_time_step(**create_time_step_kwargs)
assert isinstance(ts, time_step.TimeStep)
self._time_steps.append(ts)
def change_last_time_step(self, **replace_time_step_kwargs):
"""Replace the last time-steps with the given kwargs."""
# Pre-conditions: self._time_steps shouldn't be empty.
assert self._time_steps
self._time_steps[-1] = self._time_steps[-1].replace(
**replace_time_step_kwargs)
def truncate(self, num_to_keep=1):
"""Truncate trajectories, keeping the last `num_to_keep` time-steps."""
# We return `ts_copy` back to the truncator.
ts_copy = self._time_steps[:]
# We keep the last few observations.
self._time_steps = self._time_steps[-num_to_keep:]
# NOTE: We will need to set the rewards to 0, to eliminate double counting.
for i in range(self.num_time_steps):
self._time_steps[i] = self._time_steps[i].replace(
raw_reward=0, processed_reward=0)
return Trajectory(time_steps=ts_copy)
@property
def last_time_step(self):
# Pre-conditions: self._time_steps shouldn't be empty.
assert self._time_steps
return self._time_steps[-1]
@property
def num_time_steps(self):
return len(self._time_steps)
@property
def is_active(self):
return bool(self.num_time_steps)
@property
def time_steps(self):
return self._time_steps
@property
def done(self):
return self.is_active and self.last_time_step.done
# TODO(afrozm): Add discounting and rewards-to-go when it makes sense.
@property
def reward(self):
"""Returns a tuple of sum of raw and processed rewards."""
raw_rewards, processed_rewards = 0, 0
for ts in self.time_steps:
# NOTE: raw_reward and processed_reward are None for the first time-step.
if ts.raw_reward is not None:
raw_rewards += ts.raw_reward
if ts.processed_reward is not None:
processed_rewards += ts.processed_reward
return raw_rewards, processed_rewards
@property
def observations_np(self):
return np.stack([ts.observation for ts in self.time_steps])
def last_n_observations_np(self, n=None):
if n is not None:
n = -n # pylint: disable=invalid-unary-operand-type
return np.stack([ts.observation for ts in self.time_steps[n:]])
@property
def actions_np(self):
# The last action is None, so let's skip it.
return np.stack([ts.action for ts in self.time_steps[:-1]])
@property
def info_np(self):
if not self.time_steps or not self.time_steps[0].info:
return None
info_np_dict = {}
for info_key in self.time_steps[0].info:
# Same as actions, the last info is missing, so we skip it.
info_np_dict[info_key] = np.stack(
[ts.info[info_key] for ts in self.time_steps[:-1]])
return info_np_dict
@property
def rewards_np(self):
# The first reward is None, so let's skip it.
return np.stack([ts.processed_reward for ts in self.time_steps[1:]])
@property
def raw_rewards_np(self):
return np.stack([ts.raw_reward for ts in self.time_steps[1:]])
@property
def as_numpy(self):
# TODO(afrozm): Return a named tuple here, ex: TrajectoryArrays
return (self.observations_np, self.actions_np, self.rewards_np,
self.raw_rewards_np, self.info_np)
class BatchTrajectory(object):
"""Basically a batch of active trajectories and a list of completed ones."""
def __init__(self,
batch_size=1,
trajectories=None,
completed_trajectories=None):
self.batch_size = batch_size
# Stores trajectories that are currently active, i.e. aren't done or reset.
self._trajectories = trajectories or [
Trajectory() for _ in range(self.batch_size)
]
# Stores trajectories that are completed.
# NOTE: We don't track the index this came from, as it's not needed, right?
self._completed_trajectories = completed_trajectories or []
def reset_batch_trajectories(self):
self.__init__(batch_size=self.batch_size)
def __str__(self):
string = "BatchTrajectory["
for i, t in enumerate(self.trajectories):
string += "Trajectory {} = {}\n".format(i, str(t))
for i, t in enumerate(self.completed_trajectories):
string += "Completed Trajectory {} = {}\n".format(i, str(t))
return string + "]"
@property
def trajectories(self):
return self._trajectories
@property
def completed_trajectories(self):
return self._completed_trajectories
def clear_completed_trajectories(self, num=None):
"""Clear the first `num` completed trajectories, or all if num is None."""
if num is None:
self._completed_trajectories = []
else:
self._completed_trajectories = self._completed_trajectories[num:]
def _complete_trajectory(self, trajectory, index):
"""Completes the given trajectory at the given index."""
assert isinstance(trajectory, Trajectory)
# This *should* be the case.
assert trajectory.last_time_step.action is None
# Add to completed trajectories.
self._completed_trajectories.append(trajectory)
# Make a new one to replace it.
self._trajectories[index] = Trajectory()
def truncate_trajectories(self, indices, num_to_keep=1):
"""Truncate trajectories at specified indices.
This puts the truncated trajectories in the completed list and makes new
trajectories with the observation from the trajectory that was truncated at
the same index.
Args:
indices: iterable with the indices to truncate.
num_to_keep: int, number of last time-steps to keep while truncating.
"""
for index in indices:
trajectory = self._trajectories[index]
assert trajectory.is_active, "Trajectory to truncate can't be inactive."
# Now `trajectory` just consists of the last `num_to_keep` observations
# and actions. Rewards are zeroed out.
# The old data is placed in `old_trajectory`.
old_trajectory = trajectory.truncate(num_to_keep=num_to_keep)
# We put the old data in _completed_trajectories.
self._completed_trajectories.append(old_trajectory)
def reset(self, indices, observations):
"""Resets trajectories at given indices and populates observations.
Reset can either be called right at the beginning, when there are no
time-steps, or to reset a currently active trajectory.
If resetting a currently active trajectory then we save it in
self._completed_trajectories.
Args:
indices: 1-D np.ndarray stating the indices to reset.
observations: np.ndarray of shape (indices len, obs.shape) of observations
"""
# Pre-conditions: indices, observations are np arrays.
# : indices is one-dimensional.
# : their first dimension (batch) is the same.
assert isinstance(indices, np.ndarray)
assert len(indices.shape) == 1
assert isinstance(observations, np.ndarray)
assert indices.shape[0] == observations.shape[0]
for index, observation in zip(indices, observations):
trajectory = self._trajectories[index]
# Are we starting a new trajectory at the given index?
if not trajectory.is_active:
# Then create a new time-step here with the given observation.
trajectory.add_time_step(observation=observation)
# That's all we need to do here.
continue
# If however we are resetting a currently active trajectory then we need
# to put that in self._completed_trajectories and make a new trajectory
# with the current observation.
# TODO(afrozm): Should we mark these are done? Or is the done=False and
# this being the last time-step in the trajectory good enough to recognize
# that this was reset?
# Mark trajectory as completed and move into completed_trajectories.
self._complete_trajectory(trajectory, index)
# Put the observation in the newly created trajectory.
# TODO(afrozm): Add 0 reward.
self._trajectories[index].add_time_step(observation=observation)
def complete_all_trajectories(self):
"""Essentially same as reset, but we don't have observations."""
for index in range(self.batch_size):
trajectory = self._trajectories[index]
# TODO(pkozakowski): This assertion breaks something in SimPLe trajectory
# collection code - we're probably doing something wrong there. Commenting
# out the assertion as a temporary measure.
# assert trajectory.is_active
if trajectory.is_active:
self._complete_trajectory(trajectory, index)
def step(self,
observations,
raw_rewards,
processed_rewards,
dones,
actions,
infos=None):
"""Record the information obtained from taking a step in all envs.
Records (observation, rewards, done) in a new time-step and actions in the
current time-step.
If any trajectory gets done, we move that trajectory to
completed_trajectories.
Args:
observations: ndarray of first dimension self.batch_size, which has the
observations after we've stepped, i.e. s_{t+1} where t is the current
state.
raw_rewards: ndarray of first dimension self.batch_size containing raw
rewards i.e. r_{t+1}.
processed_rewards: ndarray of first dimension self.batch_size containing
processed rewards. i.e. r_{t+1}
dones: ndarray of first dimension self.batch_size, containing true at an
index if that env is done, i.e. d_{t+1}
actions: ndarray of first dimension self.batch_size, containing actions
applied at the current time-step, which leads to the observations
rewards and done at the next time-step, i.e. a_t
infos: (optional) a dictionary of keys and values, where all the values
have the first dimension as self.batch_size.
"""
# Pre-conditions
assert isinstance(observations, np.ndarray)
assert isinstance(raw_rewards, np.ndarray)
assert isinstance(processed_rewards, np.ndarray)
assert isinstance(dones, np.ndarray)
assert isinstance(actions, np.ndarray)
if infos:
assert isinstance(infos, dict)
# We assume that we step in all envs, i.e. not like reset where we can reset
# some envs and not others.
assert self.batch_size == observations.shape[0]
assert self.batch_size == raw_rewards.shape[0]
assert self.batch_size == processed_rewards.shape[0]
assert self.batch_size == dones.shape[0]
assert self.batch_size == actions.shape[0]
if infos:
for _, v in infos.items():
assert self.batch_size == len(v)
def extract_info_at_index(infos, index):
if not infos:
return None
return {k: v[index] for k, v in infos.items()}
for index in range(self.batch_size):
trajectory = self._trajectories[index]
# NOTE: If the trajectory isn't active, that means it doesn't have any
# time-steps in it, but we are in step, so the assumption is that it has
# a prior observation from which we are stepping away from.
# TODO(afrozm): Let's re-visit this if it becomes too restrictive.
assert trajectory.is_active
# To this trajectory's last time-step, set actions.
trajectory.change_last_time_step(
action=actions[index], info=extract_info_at_index(infos, index))
# Create a new time-step to add observation, done & rewards (no actions).
trajectory.add_time_step(
observation=observations[index],
done=dones[index],
raw_reward=raw_rewards[index],
processed_reward=processed_rewards[index])
# If the trajectory is completed, i.e. dones[index] == True, then we
# account for it right-away.
if dones[index]:
self._complete_trajectory(trajectory, index)
# NOTE: The new trajectory at `index` is going to be in-active and
# `reset` should be called on it.
assert not self._trajectories[index].is_active
@staticmethod
def _trajectory_lengths(trajectories):
return np.array([t.num_time_steps for t in trajectories])
@property
def num_completed_time_steps(self):
"""Returns the number of time-steps in completed trajectories."""
return sum(BatchTrajectory._trajectory_lengths(self.completed_trajectories))
@property
def num_time_steps(self):
"""Returns the number of time-steps in completed and incomplete trajectories."""
num_time_steps = sum(BatchTrajectory._trajectory_lengths(self.trajectories))
return num_time_steps + self.num_completed_time_steps
@property
def trajectory_lengths(self):
return BatchTrajectory._trajectory_lengths(self.trajectories)
@property
def num_completed_trajectories(self):
"""Returns the number of completed trajectories."""
return len(self.completed_trajectories)
# TODO(afrozm): Take in an already padded observation ndarray and just append
# the last time-step and adding more padding if needed.
def observations_np(self, boundary=20, len_history_for_policy=20):
"""Pads the observations in all the trajectories and returns them.
Args:
boundary: integer, Observations will be padded to (n * boundary) + 1 where
n is an integer.
len_history_for_policy: int, For each trajectory return only the last
`len_history_for_policy` observations. Set to None for all the
observations.
Returns:
padded_observations: (self.batch_size, n * boundary + 1) + OBS
"""
list_observations_np_ts = [
t.last_n_observations_np(n=len_history_for_policy)
for t in self.trajectories
]
# Every element in `list_observations_np_ts` is shaped (t,) + OBS
OBS = list_observations_np_ts[0].shape[1:] # pylint: disable=invalid-name
trajectory_lengths = np.stack(
[obs.shape[0] for obs in list_observations_np_ts])
t_max = max(trajectory_lengths)
# t_max is rounded to the next multiple of `boundary`
boundary = int(boundary)
bucket_length = boundary * int(np.ceil(float(t_max) / boundary))
def padding_config(obs):
# We're padding the first axis only, since that is the time-step.
num_to_pad = bucket_length + 1 - obs.shape[0]
return [(0, num_to_pad)] + [(0, 0)] * len(OBS)
return np.stack([
np.pad(obs, padding_config(obs), "constant")
for obs in list_observations_np_ts
]), trajectory_lengths
@staticmethod
def parse_trajectory_file_name(trajectory_file_name):
"""Parse out the trajectory file's groups and return to caller."""
base_trajectory_file_name = os.path.basename(trajectory_file_name)
trajectory_file_regexp = TRAJECTORY_FILE_FORMAT.format(
epoch="(.*)",
env_id="(.*)",
temperature="(.*)",
r="(.*)",
)
compiled_regexp = re.compile(trajectory_file_regexp)
r = compiled_regexp.match(base_trajectory_file_name)
if not r:
return None
g = r.groups()
if len(g) is not compiled_regexp.groups:
return None
# epoch, env_id, temp, random string
try:
epoch = int(g[0])
env_id = int(g[1])
temperature = float(g[2])
random_string = g[3]
except ValueError:
logging.error("Trajectory file name isn't parseable: %s",
base_trajectory_file_name)
return None
return epoch, env_id, temperature, random_string
@staticmethod
def load_from_directory(trajectory_dir,
epoch=None,
temperature=None,
n_trajectories=None,
up_sample=False,
sleep_time_secs=0.1,
max_tries=100,
wait_forever=False):
"""Load trajectories from specified dir and epoch.
Args:
trajectory_dir: (string) directory to find trajectories.
epoch: (int) epoch for which to load trajectories, if None we don't filter
on an epoch.
temperature: (float) this is used to filter the trajectory files, if None
we don't filter on temperature.
n_trajectories: (int) This is the batch size of the returned
BatchTrajectory object if one is returned. If set to None, then the
number of trajectories becomes the batch size. If set to some number,
then we wait for those many trajectory files to be available.
up_sample: (bool) If there are fewer than required (n_trajectories) number
of incomplete trajectories, then we upsample to make up the numbers.
sleep_time_secs: (float) Sleep time, to wait for min_trajectories. We
exponentially back-off this up till a maximum of 10 seconds.
max_tries: (int) The number of tries to get min_trajectories trajectories.
wait_forever: (bool) If true, overrides max_tries and waits forever.
Returns:
A BatchTrajectory object with all the constraints satisfied or None.
"""
# Modify the format to get a glob with desired epoch and temperature.
trajectory_file_glob = TRAJECTORY_FILE_FORMAT.format(
epoch=epoch if epoch is not None else "*",
env_id="*",
temperature=temperature if temperature is not None else "*",
r="*",
)
trajectory_files = tf.io.gfile.glob(
os.path.join(trajectory_dir, trajectory_file_glob))
if n_trajectories:
# We need to get `n_trajectories` number of `trajectory_files`.
# This works out to a maximum ~3hr waiting period.
while ((max_tries > 0 or wait_forever) and
len(trajectory_files) < n_trajectories):
logging.info(
"Sleeping for %s seconds while waiting for %s trajectories, found "
"%s right now.", sleep_time_secs, n_trajectories,
len(trajectory_files))
time.sleep(sleep_time_secs)
max_tries -= 1
sleep_time_secs = min(10.0, sleep_time_secs * 2)
trajectory_files = tf.io.gfile.glob(
os.path.join(trajectory_dir, trajectory_file_glob))
# We can't get the required number of files and we can't up-sample either.
if (len(trajectory_files) < n_trajectories) and not up_sample:
return None
# Sample up or down as the case maybe.
trajectory_files = list(
np.random.choice(trajectory_files, n_trajectories))
# We read and load all the files, revisit if this becomes a problem.
trajectories_buffer = []
for trajectory_file in trajectory_files:
with tf.io.gfile.GFile(trajectory_file, "rb") as f:
trajectory = get_pickle_module().load(f)
assert isinstance(trajectory, Trajectory)
trajectories_buffer.append(trajectory)
if not trajectories_buffer:
return None
# If n_trajectories wasn't set, then set to the number of trajectories we're
# returning.
n_trajectories = n_trajectories or len(trajectories_buffer)
# Construct and return a new BatchTrajectory object.
return BatchTrajectory(
batch_size=n_trajectories,
trajectories=[Trajectory() for _ in range(n_trajectories)],
completed_trajectories=trajectories_buffer)
================================================
FILE: tensor2tensor/envs/trajectory_test.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Tests for tensor2tensor.envs.trajectory."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import numpy as np
from tensor2tensor.envs import time_step
from tensor2tensor.envs import trajectory
import tensorflow.compat.v1 as tf
from tensorflow.compat.v1.io import gfile
class TrajectoryTest(tf.test.TestCase):
def test_empty_trajectory(self):
t = trajectory.Trajectory()
self.assertFalse(t.is_active)
self.assertEqual(0, t.num_time_steps)
self.assertFalse(t.done)
def test_add_time_step(self):
t = trajectory.Trajectory()
t.add_time_step(observation=1, done=True)
# Test that the trajectory is now active.
self.assertTrue(t.is_active)
added_t = t.last_time_step
self.assertEqual(1, added_t.observation)
self.assertTrue(added_t.done)
self.assertIsNone(None, added_t.raw_reward)
self.assertIsNone(None, added_t.processed_reward)
self.assertIsNone(None, added_t.action)
self.assertEqual(1, t.num_time_steps)
def test_change_last_time_step(self):
t = trajectory.Trajectory()
t.add_time_step(observation=1, done=False)
t.add_time_step(observation=1, done=True)
self.assertTrue(t.is_active)
num_ts_old = t.num_time_steps
self.assertEqual(2, num_ts_old)
# Assert on what the last time-step is currently.
ts = t.last_time_step
self.assertEqual(1, ts.observation)
self.assertTrue(ts.done)
self.assertEqual(None, ts.action)
# Change the last time-step.
t.change_last_time_step(done=False, action=5)
# Assert that it changed.
ts = t.last_time_step
self.assertEqual(1, ts.observation) # unchanged, since we didn't change it.
self.assertFalse(ts.done) # was True earlier
self.assertEqual(5, ts.action) # was None earlier
# Assert on the number of steps remaining the same as before.
self.assertEqual(num_ts_old, t.num_time_steps)
def test_reward(self):
t = trajectory.Trajectory()
# first time-step doesn't have rewards, since they are on entering a state.
t.add_time_step(
observation=1, raw_reward=None, processed_reward=None, done=False)
t.add_time_step(
observation=2, raw_reward=2, processed_reward=200, done=False)
t.add_time_step(
observation=3, raw_reward=3, processed_reward=300, done=True)
raw_reward, processed_reward = t.reward
self.assertEqual(5, raw_reward)
self.assertEqual(500, processed_reward)
def test_observation_np(self):
t = trajectory.Trajectory()
ts = 5
shape = (3, 4)
for _ in range(ts):
t.add_time_step(observation=np.random.uniform(size=shape), done=False)
self.assertEqual((ts,) + shape, t.observations_np.shape)
def test_truncate_and_last_n_observations_np(self):
t = trajectory.Trajectory()
ts = 5
shape = (3, 4)
for _ in range(ts):
t.add_time_step(observation=np.random.uniform(size=shape), done=False)
original_obs = np.copy(t.observations_np)
self.assertEqual((ts,) + shape, original_obs.shape)
# Now let's just get the observations from the last 2 steps.
num_to_keep = 2
truncated_original_obs = original_obs[-num_to_keep:, ...]
# Let's get the last `num_to_keep` observations
last_n_observations_np = np.copy(t.last_n_observations_np(n=num_to_keep))
# Now truncate the trajectory and get the same.
_ = t.truncate(num_to_keep=num_to_keep)
truncated_np = np.copy(t.observations_np)
# These should be the expected length.
self.assertEqual((2,) + shape, last_n_observations_np.shape)
self.assertEqual((2,) + shape, truncated_np.shape)
# Test the last `num_to_keep` are the same.
self.assertAllEqual(truncated_np, truncated_original_obs)
self.assertAllEqual(last_n_observations_np, truncated_original_obs)
def test_as_numpy(self):
t = trajectory.Trajectory()
shape = (3, 4)
# We'll have `ts` observations and `ts-1` actions and rewards.
ts = 5
num_actions = 6
observations = np.random.uniform(size=(ts,) + shape)
actions = np.random.choice(range(num_actions), size=(ts - 1,))
rewards = np.random.choice([-1, 0, 1], size=(ts - 1,))
squares = np.arange(ts - 1)**2
cubes = np.arange(ts - 1)**3
def get_info(i):
return {"sq": squares[i], "cu": cubes[i]}
# First time-step has no reward.
t.add_time_step(
observation=observations[0],
done=False,
action=actions[0],
info=get_info(0))
for i in range(1, ts - 1):
t.add_time_step(
observation=observations[i],
done=False,
raw_reward=rewards[i - 1],
processed_reward=rewards[i - 1],
action=actions[i],
info=get_info(i))
# Last time-step has no action.
t.add_time_step(
observation=observations[-1],
done=False,
raw_reward=rewards[-1],
processed_reward=rewards[-1])
traj_np = t.as_numpy
self.assertAllEqual(observations, traj_np[0])
self.assertAllEqual(actions, traj_np[1])
self.assertAllEqual(rewards, traj_np[2])
self.assertAllEqual(squares, traj_np[4]["sq"])
self.assertAllEqual(cubes, traj_np[4]["cu"])
class BatchTrajectoryTest(tf.test.TestCase):
BATCH_SIZE = 10
OBSERVATION_SHAPE = (3, 4)
def get_random_observations_rewards_actions_dones(self, batch_size=None):
batch_size = batch_size or self.BATCH_SIZE
# Random observations, rewards, actions, done of the expected shape.
observations = np.random.rand(*((batch_size,) + self.OBSERVATION_SHAPE))
raw_rewards = np.random.randn(batch_size)
actions = np.random.randn(batch_size)
# 40% change of being done.
dones = np.random.random((batch_size,)) > 0.6
return observations, raw_rewards, actions, dones
def test_creation(self):
bt = trajectory.BatchTrajectory(batch_size=self.BATCH_SIZE)
self.assertEqual(self.BATCH_SIZE, len(bt.trajectories))
self.assertEqual(0, bt.num_completed_trajectories)
def test_reset_all(self):
bt = trajectory.BatchTrajectory(batch_size=self.BATCH_SIZE)
indices = np.arange(self.BATCH_SIZE)
observations, _, _, _ = self.get_random_observations_rewards_actions_dones()
# Call reset.
bt.reset(indices, observations)
# Assert that all trajectories are active and not done (reset never marks
# anything as done).
self.assertTrue(all(t.is_active for t in bt.trajectories))
self.assertEqual(0, bt.num_completed_trajectories)
def test_num_time_steps(self):
bt = trajectory.BatchTrajectory(batch_size=self.BATCH_SIZE)
self.assertEqual(0, bt.num_completed_time_steps)
self.assertEqual(0, bt.num_time_steps)
def test_reset_some(self):
bt = trajectory.BatchTrajectory(batch_size=self.BATCH_SIZE)
indices = np.arange(self.BATCH_SIZE // 2)
observations, _, _, _ = self.get_random_observations_rewards_actions_dones(
batch_size=self.BATCH_SIZE // 2)
# Just reset the first half.
bt.reset(indices, observations)
# So first half are active, rest aren't.
self.assertTrue(
all(t.is_active for t in bt.trajectories[:self.BATCH_SIZE // 2]))
self.assertTrue(
all(not t.is_active for t in bt.trajectories[self.BATCH_SIZE // 2:]))
# Nothing is done anyways.
self.assertEqual(0, bt.num_completed_trajectories)
def test_truncate(self):
batch_size = 1
bt = trajectory.BatchTrajectory(batch_size=batch_size)
indices = np.arange(batch_size)
observations, _, _, _ = (
self.get_random_observations_rewards_actions_dones(
batch_size=batch_size))
# Have to call reset first.
bt.reset(indices, observations)
# Take a few steps.
ts = 5
for _ in range(ts):
(observations, rewards, actions,
dones) = self.get_random_observations_rewards_actions_dones(
batch_size=batch_size)
dones[...] = False
bt.step(observations, rewards, rewards, dones, actions)
self.assertEqual(0, bt.num_completed_trajectories)
num_to_keep = 2
bt.truncate_trajectories(indices, num_to_keep=num_to_keep)
self.assertEqual(batch_size, bt.num_completed_trajectories)
# Assert they are all active.
# Since the last `num_to_keep` observations were duplicated.
self.assertTrue(all(t.is_active for t in bt.trajectories))
orig_obs = bt.completed_trajectories[0].observations_np
# + 1 because of the initial reset
self.assertEqual(ts + 1, orig_obs.shape[0])
trunc_obs = bt.trajectories[0].observations_np
self.assertEqual(num_to_keep, trunc_obs.shape[0])
self.assertEqual(num_to_keep, bt.trajectories[0].num_time_steps)
# Test that the observations are the same.
self.assertAllEqual(orig_obs[-num_to_keep:, ...], trunc_obs)
def test_step(self):
bt = trajectory.BatchTrajectory(batch_size=self.BATCH_SIZE)
indices = np.arange(self.BATCH_SIZE)
observations, _, _, _ = self.get_random_observations_rewards_actions_dones()
# Have to call reset first.
bt.reset(indices, observations)
# Create some fake data for calling step.
new_observations, raw_rewards, actions, dones = (
self.get_random_observations_rewards_actions_dones())
processed_rewards = raw_rewards.astype(np.int64)
# Force mark the first one as done anyways, so that there is something to
# test.
dones[0] = True
num_done = sum(dones)
self.assertLessEqual(1, num_done) # i.e. num_done is atleast 1.
num_not_done = len(dones) - num_done
# Finally call step.
bt.step(new_observations, raw_rewards, processed_rewards, dones, actions)
# Expect to see `num_done` number of completed trajectories.
self.assertEqual(num_done, bt.num_completed_trajectories)
# Expect to see that the rest are marked as active.
num_active = sum(t.is_active for t in bt.trajectories)
self.assertEqual(num_not_done, num_active)
def test_desired_placement_of_rewards_and_actions(self):
batch_size = 1
bt = trajectory.BatchTrajectory(batch_size=batch_size)
indices = np.arange(batch_size)
observations, _, _, _ = self.get_random_observations_rewards_actions_dones(
batch_size=batch_size)
# Have to call reset first.
bt.reset(indices, observations)
# Create some fake data for calling step.
new_observations, raw_rewards, actions, _ = (
self.get_random_observations_rewards_actions_dones(
batch_size=batch_size))
processed_rewards = raw_rewards.astype(np.int64)
dones = np.full(batch_size, False)
# Call step.
bt.step(new_observations, raw_rewards, processed_rewards, dones, actions)
# Assert that nothing is done, since dones is False
self.assertEqual(0, bt.num_completed_trajectories)
# The only trajectory is active.
self.assertEqual(batch_size, len(bt.trajectories))
t = bt.trajectories[0]
self.assertTrue(t.is_active)
self.assertEqual(2, t.num_time_steps)
ts = t.time_steps
# Now assert on placements
# i.e. the old observation/done is first and the new one comes later.
self.assertAllEqual(observations[0], ts[0].observation)
self.assertAllEqual(new_observations[0], ts[1].observation)
self.assertEqual(False, ts[0].done)
self.assertEqual(False, ts[1].done)
# Similarly actions went to the first time-step.
self.assertEqual(actions[0], ts[0].action)
self.assertIsNone(ts[1].action)
# However make sure reward went into the second time-step and not the first.
self.assertNear(raw_rewards[0], ts[1].raw_reward, 1e-6)
self.assertIsNone(ts[0].raw_reward)
# Similarly with processed_rewards.
self.assertEqual(processed_rewards[0], ts[1].processed_reward)
self.assertIsNone(ts[0].processed_reward)
def test_observations_np(self):
bt = trajectory.BatchTrajectory(batch_size=self.BATCH_SIZE)
indices = np.arange(self.BATCH_SIZE)
observations, _, _, _ = self.get_random_observations_rewards_actions_dones()
# Have to call reset first.
bt.reset(indices, observations)
# Number of time-steps now looks like the following:
# (1, 1, 1, 1, 1, 1, 1, 1, 1, 1)
lengths = np.full((self.BATCH_SIZE,), 1)
ts = 5
for _ in range(ts):
(observations, rewards, actions,
dones) = self.get_random_observations_rewards_actions_dones()
dones[...] = False
bt.step(observations, rewards, rewards, dones, actions)
# Number of time-steps now looks like the following:
# (6, 6, 6, 6, 6, 6, 6, 6, 6, 6)
lengths = lengths + ts
# Now let's mark the first two as done.
observations, _, _, _ = self.get_random_observations_rewards_actions_dones(
batch_size=2)
bt.reset(np.array([0, 1]), observations)
# Number of time-steps now looks like the following:
# (1, 1, 6, 6, 6, 6, 6, 6, 6, 6)
lengths[0] = lengths[1] = 1
for _ in range(ts):
(observations, rewards, actions,
dones) = self.get_random_observations_rewards_actions_dones()
dones[...] = False
bt.step(observations, rewards, rewards, dones, actions)
# Number of time-steps now looks like the following:
# (6, 6, 11, 11, 11, 11, 11, 11, 11, 11)
lengths = lengths + ts
boundary = 20
len_history_for_policy = 40
padded_obs_np, padded_lengths = bt.observations_np(
boundary=boundary, len_history_for_policy=len_history_for_policy)
# The lengths are what we expect them to be.
self.assertAllEqual(lengths, padded_lengths)
# The padded_observations are the shape we expect them to be.
self.assertEqual((self.BATCH_SIZE, boundary + 1) + self.OBSERVATION_SHAPE,
padded_obs_np.shape)
# Let's now request the last n = [1, 2 * boundary) steps for the history.
for len_history_for_policy in range(1, 2 * boundary):
# The expected lengths will now be:
truncated_lengths = [min(l, len_history_for_policy) for l in lengths]
padded_obs_np, padded_lengths = bt.observations_np(
boundary=boundary, len_history_for_policy=len_history_for_policy)
self.assertAllEqual(truncated_lengths, padded_lengths)
# This shouldn't change, since even if we request lengths > boundary + 1
# there are no trajectories that long.
self.assertEqual((self.BATCH_SIZE, boundary + 1) + self.OBSERVATION_SHAPE,
padded_obs_np.shape)
# Let's do 10 more steps (to go on the other side of the boundary.
ts = 10
for _ in range(ts):
(observations, rewards, actions,
dones) = self.get_random_observations_rewards_actions_dones()
dones[...] = False
bt.step(observations, rewards, rewards, dones, actions)
# Number of time-steps now looks like the following:
# (16, 16, 21, 21, 21, 21, 21, 21, 21, 21)
lengths = lengths + ts
len_history_for_policy = 40
padded_obs_np, padded_lengths = bt.observations_np(
boundary=boundary, len_history_for_policy=len_history_for_policy)
# The lengths are what we expect them to be.
self.assertAllEqual(lengths, padded_lengths)
# The padded_observations are the shape we expect them to be.
self.assertEqual(
(self.BATCH_SIZE, (2 * boundary) + 1) + self.OBSERVATION_SHAPE,
padded_obs_np.shape)
# Test that the padding is the only part that is all 0s.
# NOTE: There is almost 0 probability that the random observation is all 0s.
zero_obs = np.full(self.OBSERVATION_SHAPE, 0.)
for b in range(self.BATCH_SIZE):
# The first lengths[b] will be actual data, rest is 0s.
for ts in range(lengths[b]):
self.assertFalse(np.all(zero_obs == padded_obs_np[b][ts]))
for ts in range(lengths[b], len(padded_obs_np[b])):
self.assertAllEqual(zero_obs, padded_obs_np[b][ts])
def test_parse_trajectory_file_name(self):
self.assertEqual(
(12, 13, 1.0, "abc"),
trajectory.BatchTrajectory.parse_trajectory_file_name(
"/tmp/trajectory_epoch_000012_env_id_000013_temperature_1.0_r_abc.pkl"
))
self.assertIsNone(
trajectory.BatchTrajectory.parse_trajectory_file_name(
"/tmp/trajectory_epoch_000012_env_id_000013.pkl"))
def test_load_from_directory(self):
output_dir = self.get_temp_dir()
epochs = [0, 1, 2]
env_ids = [0, 1, 2]
temperatures = [0.5, 1.0]
random_strings = ["a", "b"]
# Write some trajectories.
# There are 3x3x2x2 (36) trajectories, and of them 3x2x2 (12) are done.
for epoch in epochs:
for env_id in env_ids:
for temperature in temperatures:
for random_string in random_strings:
traj = trajectory.Trajectory(time_steps=[
time_step.TimeStep(
observation=epoch,
done=(epoch == 0),
raw_reward=1.0,
processed_reward=1.0,
action=env_id,
info={})
])
trajectory_file_name = trajectory.TRAJECTORY_FILE_FORMAT.format(
epoch=epoch,
env_id=env_id,
temperature=temperature,
r=random_string)
with gfile.GFile(
os.path.join(output_dir, trajectory_file_name), "w") as f:
trajectory.get_pickle_module().dump(traj, f)
# Load everything and check.
bt = trajectory.BatchTrajectory.load_from_directory(output_dir)
self.assertIsInstance(bt, trajectory.BatchTrajectory)
self.assertEqual(36, bt.num_completed_trajectories)
self.assertEqual(36, bt.batch_size)
bt = trajectory.BatchTrajectory.load_from_directory(output_dir, epoch=0)
self.assertEqual(12, bt.num_completed_trajectories)
self.assertEqual(12, bt.batch_size)
# Get 100 trajectories, but there aren't any.
bt = trajectory.BatchTrajectory.load_from_directory(
output_dir, epoch=0, n_trajectories=100, max_tries=0)
self.assertIsNone(bt)
bt = trajectory.BatchTrajectory.load_from_directory(
output_dir, epoch=0, temperature=0.5)
self.assertEqual(6, bt.num_completed_trajectories)
self.assertEqual(6, bt.batch_size)
bt = trajectory.BatchTrajectory.load_from_directory(output_dir, epoch=1)
self.assertEqual(12, bt.num_completed_trajectories)
self.assertEqual(12, bt.batch_size)
# Constraints cannot be satisfied.
bt = trajectory.BatchTrajectory.load_from_directory(
output_dir, epoch=1, n_trajectories=100, up_sample=False, max_tries=0)
self.assertIsNone(bt)
# Constraints can be satisfied.
bt = trajectory.BatchTrajectory.load_from_directory(
output_dir, epoch=1, n_trajectories=100, up_sample=True, max_tries=0)
self.assertEqual(100, bt.num_completed_trajectories)
self.assertEqual(100, bt.batch_size)
bt = trajectory.BatchTrajectory.load_from_directory(
output_dir, epoch=1, n_trajectories=10)
self.assertEqual(10, bt.num_completed_trajectories)
self.assertEqual(10, bt.batch_size)
gfile.rmtree(output_dir)
if __name__ == "__main__":
tf.test.main()
================================================
FILE: tensor2tensor/insights/README.md
================================================
# Tensor2Tensor Insights
The Insights packages provides an interactive webservice for understanding the
inner workings of a Tensor2Tensor model. It will provide a series of
visualizations extracted from a requested T2T model that informs model developers
and model users on how to improve or best utilize a model.
## Dependencies
Before using the Insights server, you must install [Bower](https://bower.io/)
which we use to manage our web component dependencies. You can easily install
this with the [Node Package Manager](https://www.npmjs.com/).
## Setup Instructions
After training a model, such as according to the Quick Start guide, you can run
the `t2t-insights-server` binary and begin querying it.
First, prepare the bower dependencies by navigating into the
`tensor2tensor/insights/polymer` directory and running `bower install`:
```
pushd tensor2tensor/insights/polymer
bower install
popd
```
The models run by server is then configured by a JSON version of the
InsightsConfiguration protocol buffer. Using the model trained in the Quick
Start guide, a sample configuration would be:
```
{
"configuration": [{
"source_language": "en",
"target_language": "de",
"label": "transformers_wmt32k",
"transformer": {
"model": "transformer",
"model_dir": "/tmp/t2t/train",
"data_dir": "/tmp/t2t/data",
"hparams": "",
"hparams_set": "transformer_base_single_gpu",
"problem": "translate_ende_wmt32k"
}
}],
"language": [{
"code": "en",
"name": "English"
},{
"code": "de",
"name": "German"
}]
}
```
With that saved to `configuration.json`, run the following:
```
t2t-insights-server \
--configuration=configuration.json \
--static_path=`pwd`/tensor2tensor/insights/polymer
```
This will bring up a minimal [Flask](http://flask.pocoo.org/) REST service
served by a [GUnicorn](http://gunicorn.org/) HTTP Server.
## Features to be developed
This is a minimal web server. We are in the process of adding additional
exciting features that give insight into a model's behavior:
* Integrating a multi-head attention visualization.
* Registering multiple models to compare their behavior.
* Indexing training data to find examples related to a current query.
* Tracking interesting query + translation pairs for deeper analysis.
================================================
FILE: tensor2tensor/insights/__init__.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
================================================
FILE: tensor2tensor/insights/graph.py
================================================
# coding=utf-8
# Copyright 2023 The Tensor2Tensor 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.
"""Graph representation for building decoding graph visualizations."""
class Vertex(object):
"""Vertex stores in and out edge connections to other Vertex instances.
The Vertex class supports serialization to a JSON data format expected by the
client side representation. When serializing, it generates the following
fields:
in_edge_index: The list of directed edge indices into the Vertex.
out_edge_index: The list of directed edge indices from the Vertex.
"""
def __init__(self, idx):
"""Initialize the Vertex.
Args:
idx: The index of the vertex.
"""
self.idx = idx
self.in_edges = []
self.out_edges = []
def to_dict(self):
"""Returns a simplified dictionary representing the Vertex.
Returns:
A dictionary that can easily be serialized to JSON.
"""
return {
"in_edge_index": self.in_edges,
"out_edge_index": self.out_edges,
}
class Edge(object):
"""Edge stores edge details connecting two Vertex instances.
The Edge class supports serialization to a JSON data format expected by the
client side representation. When serializing, it generates the following
fields:
source_index: The source Vertex index for this Edge.
target_index: The target Vertex index for this Edge.
data: Arbitrary data for this Edge.
"""
def __init__(self, idx):
"""Initialize the Edge.
Args:
idx: The index of the Edge.
"""
self.idx = idx
self.source = -1
self.target = -1
self.data = {}
def to_dict(self):
"""Returns a simplified dictionary representing the Vertex.
Returns:
A dictionary that can easily be serialized to JSON.
"""
return {
"source_index": self.source,
"target_index": self.target,
"data": self.data,
}
def __str__(self):
return str(self.to_dict())
class Graph(object):
"""A directed graph that can easily be JSON serialized for visualization.
When serializing, it generates the following fields:
edge: The list of all serialized Edge instances.
node: The list of all serialized Vertex instances.
"""
def __init__(self):
self.vertices = []
self.edges = []
self.vertex_map = {}
def new_vertex(self):
"""Creates and returns a new vertex.
Returns:
A new Vertex instance with a unique index.
"""
vertex = Vertex(len(self.vertices))
self.vertices.append(vertex)
return vertex
def get_vertex(self, key):
"""Returns or Creates a Vertex mapped by key.
Args:
key: A string reference for a vertex. May refer to a new Vertex in which
case it will be created.
Returns:
A the Vertex mapped to by key.
"""
if key in self.vertex_map:
return self.vertex_map[key]
vertex = self.new_vertex()
self.vertex_map[key] = vertex
return vertex
def add_edge(self, source, target):
"""Returns a new edge connecting source and target vertices.
Args:
source: The source Vertex.
target: The target Vertex.
Returns:
A new Edge linking source to target.
"""
edge = Edge(len(self.edges))
self.edges.append(edge)
source.out_edges.append(edge.idx)
target.in_edges.append(edge.idx)
edge.source = source.idx
edge.target = target.idx
return edge
def to_dict(self):
"""Returns a simplified dictionary representing the Graph.
Returns:
A dictionary that can easily be serialized to JSON.
"""
return {
"node": [v.to_dict() for v in self.vertices],
"edge": [e.to_dict() for e in self.edges]
}
================================================
FILE: tensor2tensor/insights/insight_configuration.proto
================================================
syntax = "proto3";
package tensor2tensor;
// Configures the Neural Machine Translation Insight Frontend with a set of
// supported query processors and languages.
message InsightConfiguration {
// Specifies zero or more models to inspect.
repeated QueryProcessorConfiguration configuration = 1;
// Specifies language codes and display names.
repeated Language language = 2;
}
// A displayable language name.
message Language {
// The BCP-47 Language code.
string code = 1;
// The language's display name.
string name = 2;
}
// Configures a QueryProcessor and registers it with the Insight Frontend when
// responding to analysis queries.
message QueryProcessorConfiguration {
// The model's BCP-47 source language code.
string source_language = 1;
// The model's BCP-47 target language code.
string target_language = 2;
// A short label for the model.
string label = 3;
// The QueryProcessor to use. By default we just use the TransformerModel.
string query_processor = 4;
// Configuration for the TransformerModel.
TransformerConfiguration transformer = 5;
}
// Specifies the parameters for a trained Transformer model to inspect. These
// parameters match those in t2t-trainer and t2t-decoder.
message TransformerConfiguration {
// The model type.
string model = 1;
// The trained model directory.
string model_dir = 2;
// The data directory for the model.
string data_dir = 3;
// The hyperparameter set for running the model.
string hparams_set = 4;
// Overriding hyperparameters.
string hparams = 5;
// The problem sets over which this model was trained and configured.
string problems = 6;
}
================================================
FILE: tensor2tensor/insights/polymer/.bowerrc
================================================
{
"directory": "."
}
================================================
FILE: tensor2tensor/insights/polymer/attention_visualization/attention-visualization.html
================================================
<!--
@license
Copyright 2018 The Tensor2Tensor 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.
-->
<link rel="import" href="../polymer/polymer.html">
<link rel="import" href="../iron-flex-layout/iron-flex-layout-classes.html">
<link rel="import" href="../iron-icon/iron-icon.html">
<link rel="import" href="../iron-icons/iron-icons.html">
<link rel="import" href="../paper-icon-button/paper-icon-button.html">
<link rel="import" href="../paper-slider/paper-slider.html">
<dom-module id="attention-visualization">
<template>
<custom-style>
<style is="custom-style" include="iron-flex iron-flex-alignment"></style>
</custom-style>
<style>
.background {
fill: #eee;
}
svg *::selection {
background: transparent;
}
rect.selection {
fill: transparent;
stroke: #333;
stroke-dasharray: 4px;
stroke-opacity: 0.5;
}
rect.cell-border {
stroke: #eee;
stroke-width: 0.3px;
}
rect.cell-selected {
stroke: rgb(51, 102, 153);
stroke-width: 0.5px;
}
g.cell-group {
pointer-events: all;
}
g.cell-hover rect {
stroke: #f00;
stroke-width: 1px;
}
text.mono {
fill: #aaa;
}
text.text-highlight {
fill: #c00;
}
text.weight-label {
fill: #ffffff;
font-size: 16px;
stroke: #ffffff;
}
text.text-selected {
fill: #000;
}
.svg-container {
display: inline-block;
overflow: hidden;
padding-bottom: 100%;
position: relative;
vertical-align: top;
width: 100%;
}
.svg-content-responsive {
display: inline-block;
left: 0px;
position: absolute;
top: 10px;
}
#tooltip {
background-color: white;
border-radius: 10px;
box-shadow: 4px 4px 10px rgba(0, 0, 0, 0.4);
height: auto;
padding: 10px;
pointer-events: none;
position: absolute;
width: auto;
z-index: 10;
}
#tooltip.hidden {
display: none;
}
</style>
<div class="layout horizontal">
<paper-icon-button id="home" on-tap="reset_" icon=home></paper-icon-button>
<paper-slider id="slider" min=0 max=100 value="{{zoomDepth_}}"></paper-slider>
</div>
<div id="tooltip" class="hidden">
<span>{{selectedProbability}}<span>
</div>
<div id="chart">
</div>
</template>
<script src="attention-visualization.js"></script>
</dom-module>
================================================
FILE: tensor2tensor/insights/polymer/attention_visualization/attention-visualization.js
================================================
/**
* @license
* Copyright 2018 The Tensor2Tensor 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.
*/
/**
* `<attention-visualization>` presents a heatmap of input-output associations.
*
* The heat map association shows source to target word association strengths
* according to some method.
*
* ### Usage
*
* <attention-visualization data="[[data]]"></attention-visualization>
*/
class AttentionVisualization extends Polymer.Element {
constructor() {
super();
/**
* D3.js DOM element.
* @private
*/
this.container_ = undefined;
/**
* @private
*/
this.margin_ = {
top: 150,
bottom: 50,
right: 10,
left: 100
};
/**
* D3.js DOM element.
* @private
*/
this.svg_ = undefined;
/**
* D3.js DOM element.
* @private
*/
this.vis_ = undefined;
/**
* D3.js DOM element.
* @private
*/
this.zoom_ = undefined;
}
/**
* @return {string} The component name.
*/
static get is() {
return 'attention-visualization';
}
/**
* @return {!Object} The component properties.
*/
static get properties() {
return {
/**
* @type {AttentionData}
*/
data: {
type: Object,
observer: 'dataUpdated_',
},
/**
* @type {number}
*/
zoomDepth_: {
type: Number,
},
};
}
/**
* @return {!Array<string>} The component observers.
*/
static get observers() {
return [
'zoomDepthChanged_(zoomDepth_)',
];
}
/**
* Sets the default zoom depth.
* @override
*/
ready() {
super.ready();
this.set('zoomDepth_', 20);
}
/**
* Sets the zoom state based on the updated depth.
* @param {number} zoomDepth the zoom depth.
* @private
*/
zoomDepthChanged_(zoomDepth) {
if (!this.container_) { return; }
if (zoomDepth == 0) {
zoomDepth = 0.000001;
}
let transform = d3.zoomTransform(this.vis_.node()).scale(zoomDepth / 20.0);
this.container_.attr("transform", transform);
}
/**
* Updates the heatmap.
* @param {!AttentionData} newData the new alignment data.
* @private
*/
dataUpdated_(newData) {
// Create the bounding areas and margins for the heatmap.
let cellDimension = 40;
let sourceTokens = newData.source_tokens;
let targetTokens = newData.target_tokens;
// Convert the attention weights to cell objects which also give access to
// the row and column indices.
let mapCells = newData.weights.map(function(d, i) {
return {
value: d,
row: Math.floor(i / targetTokens.length),
col: i % targetTokens.length
};
});
// Create the color scale.
let colorScale = d3.scaleQuantile().domain([0.0, 1.0]).range([
'#cccccc', '#b2b2b2', '#999999', '#7f7f7f',
'#666666', '#4c4c4c', '#333333', '#191919'
]);
this.zoom_ = d3.zoom().scaleExtent([1, 10]).on('zoom', zoomed.bind(this));
d3.select(this.$.chart).selectAll("*").remove();
// Create the bounding div and svgs which will contain all details.
this.svg_ = d3.select(this.$.chart)
.append('div')
.classed('svg-container', true)
.append('svg')
.attr('width', '100%')
.attr('height', '100%')
.classed('svg-content-responsive', true);
this.vis_ = this.svg_.append('g')
.attr('transform',
'translate(' + this.margin_.left + ',' + this.margin_.top + ')')
.call(this.zoom_)
.on('dblclick.zoom', null)
.on('wheel.zoom', null);
// Create a bounding rectangle upon which zooming and panning will take
// place.
this.vis_.append('rect')
.attr('width', '100%')
.attr('height', '100%')
.style('fill', 'none')
.style('pointer-events', 'all');
this.container_ = this.vis_.append('g');
// Initiate the panning and/or zooming.
function zoomed() {
this.container_.attr("transform",
d3.event.transform.scale(this.zoomDepth_ / 20.0));
}
// Place the source tokens along the vertical axis. Each token has an id
// based on it's index.
var sourceLabels = this.container_.append('g');
sourceLabels.selectAll('.source-label')
.data(sourceTokens)
.enter()
.append('text')
.text(function(d) {
return d;
})
.style('text-anchor', 'end')
.attr(
'id',
function(d, i) {
return 'row-' + i;
})
.attr('class', 'source-label mono')
.attr('transform', 'translate(-6,' + cellDimension / 1.5 + ')')
.attr('x', 0)
.attr('y', function(d, i) {
return i * cellDimension;
});
var targetLabels = this.container_.append('g');
// Place the target tokens along the horizontal axis. Each token has an id
// based on it's index.
targetLabels.selectAll('.target-label')
.data(targetTokens)
.enter()
.append('text')
.text(function(d) {
return d;
})
.style('text-anchor', 'left')
.attr(
'id',
function(d, i) {
return 'col-' + i;
})
.attr('class', 'target-label mono')
.attr(
'transform', 'translate(' + cellDimension / 2 + ',-6) rotate(-90)')
.attr(
'y',
function(d, i) {
return i * cellDimension;
})
.attr('x', 0);
// Create the heat map and populate with cells. Each cell will
// highlight when hovered over. Additionally, the column and row tokens
// will highlight to make clear which tokens are being observed. Lastly,
// each cell will trigger a popup showing details of the alignment state.
var heatMap = this.container_.append('g');
// Group the rectangle and text elements and capture the mouse events from
// both so that the rectangle can be highlighted when it's in focus.
let cellGroup = heatMap.selectAll('.cell')
.data(mapCells)
.enter()
.append('g')
.attr('class', 'cell-group')
.on('mouseover', function(d, i) {
// Highlight the newly hovered over cell and it's row/column
// tokens.
d3.select(this).classed('cell-hover', true);
sourceLabels.select('#row-' + d.row)
.classed('text-highlight', true);
targetLabels.select('#col-' + d.col)
.classed('text-highlight', true);
})
.on('mouseout', function(d) {
// Clear all highlighting.
d3.select(this).classed('cell-hover', false);
sourceLabels.select('#row-' + d.row)
.classed('text-highlight', false);
targetLabels.select('#col-' + d.col)
.classed('text-highlight', false);
});
// Add the rectangles for each cell.
cellGroup
.append('rect')
.attr(
'id',
function(d, i) {
return 'cell-' + i;
})
.attr('class', 'cell cell-border')
.attr(
'x',
function(d) {
return d.col * cellDimension;
})
.attr(
'y',
function(d) {
return d.row * cellDimension;
})
.attr('width', cellDimension)
.attr('height', cellDimension)
.style(
'fill',
function(d) {
return colorScale(d.value);
});
// Add the text for each cell.
cellGroup
.append('text')
.text(function(d) { return d.value.toFixed(2); })
.attr('class', 'weight weight-label')
.attr('x', function(d) { return 5 + (d.col * cellDimension); })
.attr('y', function(d) { return 25 + (d.row * cellDimension); });
}
/**
* Resets the pan and zoom state.
* @private
*/
reset_() {
if (!this.svg_) { return; }
this.vis_.call(this.zoom_.transform, d3.zoomIdentity);
this.set('zoomDepth_', 20);
}
}
customElements.define(AttentionVisualization.is, AttentionVisualization);
================================================
FILE: tensor2tensor/insights/polymer/bower.json
================================================
{
"name": "tensor2tensor-insights",
"homepage": "https://github.com/tensorflow/tensor2tensor",
"description": "Components for analyzing tensor2tensor neural machine translation models.",
"main": "index.html",
"keywords": [
"neural",
"machine",
"translation"
],
"authors": [
"kstevens@google.com"
],
"license": "Apache 2.0",
"private": true,
"ignore": [
"**/.*",
"node_modules",
"bower_components",
"test",
"tests"
],
"dependencies": {
"app-layout": "PolymerElements/app-layout#2.0.4",
"app-route": "PolymerElements/app-route#2.0.3",
"d3": "d3#4.12.2",
"iron-a11y-keys": "PolymerElements/iron-a11y-keys#2.0.0",
"iron-ajax": "PolymerElements/iron-ajax#2.0.0",
"iron-flex-layout": "PolymerElements/iron-flex-layout#2.0.0",
"iron-icon": "PolymerElements/iron-icon#2.0.0",
"iron-icons": "PolymerElements/iron-icons#2.0.0",
"iron-list": "PolymerElements/iron-list#2.0.0",
"iron-pages": "PolymerElements/iron-pages#2.0.0",
"iron-selector": "PolymerElements/iron-selector#2.0.0",
"neon-animation": "PolymerElements/neon-animation#2.0.0",
"paper-button": "PolymerElements/paper-button#2.0.0",
"paper-card": "PolymerElements/paper-card#2.0.0",
"paper-dialog": "PolymerElements/paper-dialog#2.0.0",
"paper-dropdown-menu": "PolymerElements/paper-dropdown-menu#2.0.0",
"paper-icon-button": "PolymerElements/paper-icon-button#2.0.0",
"paper-input": "PolymerElements/paper-input#2.0.0",
"paper-item": "PolymerElements/paper-item#2.0.0",
"paper-listbox": "PolymerElements/paper-listbox#2.0.0",
"paper-slider": "PolymerElements/paper-slider#2.0.0",
"paper-tabs": "PolymerElements/paper-tabs#2.0.0",
"paper-toggle-button": "PolymerElements/paper-toggle-button#2.0.0",
"paper-tooltip": "PolymerElements/paper-tooltip#2.0.0",
"paper-progress": "PolymerElements/paper-progress#2.0.0",
"polymer": "polymer/polymer#v2.3.1"
},
"resolutions": {
"webcomponentsjs": "^v1.0.19",
"polymer": "^v2.3.1",
"app-route": "^2.0.3",
"app-layout": "^2.0.4",
"iron-location": "1 - 2",
"iron-selector": "^2.0.0",
"neon-animation": "^2.0.0",
"iron-icon": "^2.0.0",
"iron-pages": "^2.0.0",
"iron-icons": "^2.0.0",
"paper-icon-button": "^2.0.0",
"paper-item": "^2.0.0",
"iron-flex-layout": "^2.0.0",
"paper-listbox": "^2.0.0",
"iron-a11y-keys": "^2.0.0",
"paper-dialog": "^2.0.0",
"iron-ajax": "^2.0.0",
"paper-progress": "^2.0.0",
"paper-dropdown-menu": "^2.0.0",
"paper-tabs": "^2.0.0",
"paper-input": "^2.0.0",
"paper-toggle-button": "^2.0.0",
"paper-slider": "^2.0.0",
"iron-list": "^2.0.0",
"paper-card": "^2.0.0",
"paper-tooltip": "^2.0.0",
"iron-overlay-behavior": "^2.2.0"
}
}
================================================
FILE: tensor2tensor/insights/polymer/common-types.js
================================================
/**
* @license
* Copyright 2018 The Tensor2Tensor 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.
*/
/**
* @fileoverview A set of shared types that will be replaced by js proto types.
* @externs
*/
/**
* A typedef for a nlp.nmt.mt_debug_fe.LanguageConfiguration message.
* This can't be converted to javascript yet because it transitively depends on
* tensorflow protos that can't be converted to javascript.
* TODO(kstevens): Remove this typedef when we remove the dependency on
* non-convertible tensorflow protos.
* @typedef {{
* code: string,
* name: string,
* hidden: ?boolean,
* }}
*/
let Language;
/**
* A typedef for a nlp.nmt.mt_debug_fe.SerializedConfiguration message.
* This can't be converted to javascript yet because it transitively depends on
* tensorflow protos that can't be converted to javascript.
* TODO(kstevens): Remove this typedef when we remove the dependency on
* non-convertible tensorflow protos.
* @typedef {{
* id: string,
* target: string,
* source_language: Language,
* target_language: Language,
* }}
*/
let Model;
/**
* @typedef {{
* name: string,
* localProbability: number,
* cumalitiveProbability: number,
* attention: Array<number>,
* children: Array<TreeNode>,
* }}
*/
let TreeNode;
/**
* @typedef {{
* source_tokens: Array<string>,
* target_tokens: Array<string>,
* weights: !Array<number>
* }}
*/
let AttentionData;
/**
* @typedef {{
* label: string,
* label_id: number,
* log_probability: number,
* total_log_probability: number,
* score: number,
* parent_id: number,
* }}
*/
let Candidate;
/**
* @typedef {{
* id: number,
* stepIndex: number,
* candidate: !Candidate,
* children: !Array<InteractiveNode>,
* }}
*/
let InteractiveNode;
/**
* @typedef {{
* step_name: string,
* segment: !Array<!{
* text: string,
* }>
* }}
*/
let QueryProcessingRewriteStep;
/**
* @typedef {{
* source_processing: !Array<!QueryProcessingRewriteStep>,
* target_processing: !Array<!QueryProcessingRewriteStep>,
* }}
*/
let QueryProcessingVisualization;
/**
* @typedef {{
* in_edge_index: !Array<number>,
* out_edge_index: !Array<number>,
* }}
*/
let BeamSearchNode;
/**
* @typedef {{
* label_id: number,
* label: string,
* log_probability: number,
* total_log_probability: number,
* score: number,
* completed: boolean,
* }}
*/
let BeamSearchCandidate;
/**
* @typedef {{
* source_index: number,
* target_index: number,
* data: !BeamSearchCandidate,
* }}
*/
let BeamSearchEdge;
/**
/**
* @typedef {{
* node: !Array<!BeamSearchNode>,
* edge: !Array<!BeamSearchEdge>,
* }}
*/
let SearchGraphVisualization;
/**
* @typedef {{
* candidate_list: !Array<{
* candidate: !Array<!BeamSearchCandidate>,
* }>,
* }}
*/
let GenerateCandidateResponse;
/**
* @typedef {{
* session_id: number,
* }}
*/
let StartTranslationResponse;
================================================
FILE: tensor2tensor/insights/polymer/explore_view/explore-view.html
================================================
<!--
@license
Copyright 2018 The Tensor2Tensor 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.
-->
<link rel="import" href="../polymer/polymer.html">
<link rel="import" href="../app-route/app-location.html">
<link rel="import" href="../app-route/app-route.html">
<link rel="import" href="../iron-a11y-keys/iron-a11y-keys.html">
<link rel="import" href="../iron-ajax/iron-ajax.html">
<link rel="import" href="../iron-flex-layout/iron-flex-layout-classes.html">
<link rel="import" href="../iron-icon/iron-icon.html">
<link rel="import" href="../iron-icons/iron-icons.html">
<link rel="import" href="../iron-list/iron-list.html">
<link rel="import" href="../paper-icon-button/paper-icon-button.html">
<link rel="import" href="../paper-input/paper-input.html">
<link rel="import" href="../paper-toggle-button/paper-toggle-button.html">
<link rel="import" href="../paper-progress/paper-progress.html">
<link rel="import" href="../query_card/query-card.html">
<link rel="import" href="../translation_result/translation-result.html">
<dom-module id="explore-view">
<template>
<style include="iron-flex iron-flex-alignment iron-flex-reverse">
:host {
padding: 24px;
@apply --layout-vertical;
@apply --layout-center;
}
query-card {
margin: 0px;
width: 90%;
}
div.rule {
@apply --layout-vertical;
}
paper-progress {
--paper-progress-active-color: #4285f4;
--paper-progress-height: 10px;
width: 90%;
}
translation-result {
margin: 12px 0px;
}
paper-input {
padding: 0px 6px;
}
paper-icon-button#clear {
color: var(--paper-red-300);
--paper-icon-button-ink-color: var(--paper-red-a100);
height: 23px;
padding: 0px 4px;
width: 23px;
}
paper-icon-button#translate {
background-color: #4d90fe;
color: #fff;
border-radius: 50%;
}
#result-list {
margin: 24px 0px;
width: 90%;
}
</style>
<!--
Extract the query information from the url if it exists.
-->
<app-route
route="{{subroute}}"
pattern="/:query"
tail="{{tailRoute}}"
data="{{queryData}}">
</app-route>
<query-card
route="{{route}}"
url="/api/list_models"
sub-route="{{subroute}}"
model="{{model_}}">
<!--
Include a text area and actionable button for sending translations.
-->
<div id="search-bar" class="layout horizontal center-center">
<paper-input class="flex" value="{{query_}}" label="translate" id="input">
</paper-input>
<iron-a11y-keys target="{{input}}" keys="enter" on-keys-pressed="translate_">
</iron-a11y-keys>
<paper-icon-button id="translate" on-tap="translate_" icon="translate" title="translate">
</paper-icon-button>
</div>
<div id="extra">
<h4>Rapid Response</h4>
<template is="dom-repeat" items="{{rules_}}">
<div class="rule">
<span on-tap="deleteRule_">
<iron-icon icon="remove-circle-outline">
</iron-icon>
Rule
</span>
<paper-input label="Source" value="{{item.source}}" type="text"></paper-input>
<paper-input label="Bad Target" value="{{item.bad_translations}}" type="text"></paper-input>
<paper-input value="{{item.good_translations}}" label="Good Target" type="text"></paper-input>
<paper-input value="{{item.attention_threshold}}" label="Threshold" type="number"></paper-input>
</div>
</template>
<span on-tap="addRule_">
<iron-icon icon="add-circle-outline">
</iron-icon>
Rule
</span>
</div>
</query-card>
<paper-progress id="loading" indeterminate disabled="[[!fetchingResult]]">
</paper-progress>
<div id="result-list" class="layout vertical vertical-reverse">
<template is="dom-repeat" items="[[results]]" as="result">
<translation-result result="[[result]]"></translation-result>
</template>
</div>
<iron-ajax
id="translateAjax"
url="{{url}}"
handle-as="json"
on-response="handleTranslationResponse_">
</iron-ajax>
</template>
<script src="../d3/d3.js"></script>
<script src="explore-view.js"></script>
</dom-module>
================================================
FILE: tensor2tensor/insights/polymer/explore_view/explore-view.js
================================================
/**
* @license
* Copyright 2018 The Tensor2Tensor 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.
*/
/**
* `<explore-view>` Presents a view for debuging translations.
*
* This provides an interactive interface for querying a backend service to
* fetch detailed analysis of a translation process. Each result will be
* provided as a stack.
*
* ### Usage
*
* <explore-view></explore-view>
*/
class ExploreView extends Polymer.Element {
/**
* @return {string} The component name.
*/
static get is() {
return 'explore-view';
}
/**
* @return {!Object} The component properties.
*/
static get properties() {
return {
route: {
type: Object,
},
/**
* @type {!Array<!{
* source: string,
* bad_translations: string,
* good_translations: string,
* attention_threshold: number
* }>}
*/
rules_: {
type: Array,
},
/**
* @type {?Model}
*/
model_: {
type: Object
},
/**
* @type {string}
*/
query_: {
type: Object,
}
};
}
/**
* @return {!Array<string>} The component observers.
*/
static get observers() {
return [
'modelChanged_(queryData, model_)',
];
}
/**
* @override
*/
ready() {
super.ready();
this.set('rules_', []);
this.set('fetchingResult', false);
}
/**
* Noop
* @public
*/
refresh() {
// Noop
}
/**
* Resets the results when a model changes and triggers a query automatically
* if one exists.
* @param {?{query: string}} queryData The current route data.
* @param {?Model} model Unused, but needed for triggering.
* @private
*/
modelChanged_(queryData, model) {
if (queryData && queryData.query) {
// Compose the query from the querydata field and the path in the rest of
// the route. If the link includes an escaped "/" app-route splits the
// query and remaining path on that escaped "/". So query appears to not
// include the rest of the intended query.
let query = unescape(queryData.query) + this.get('tailRoute').path;
this.set('query_', query);
this.translate_();
}
this.set('results', []);
this.set('rules_', []);
}
/**
* Sends a translation request to the server.
* @private
*/
translate_() {
if (!this.model_ || !this.model_.id) {
return;
}
var params = {
'source': this.query_,
'id': this.model_.id,
'sl': this.model_.source_language.code,
'tl': this.model_.target_language.code,
};
var paramList = this.createBodyValue_(params);
this.set('url', '/debug?' + paramList);
this.set('fetchingResult', true);
this.$.translateAjax.generateRequest();
}
/**
* Returns a string with all the query parameters composed together. This
* also serializes the rapid response rules provided.
* @param {!Object} params The params to combine.
* @returns {string} The params collapsed together.
* @private
*/
createBodyValue_(params) {
// Add the key value body parts.
var bodyParts = [];
for (var param in params) {
var value = window.encodeURIComponent(params[param]);
bodyParts.push(param + "=" + value);
}
// Add the rapid response rules.
for (var i = 0; i < this.rules_.length; ++i) {
var rule = this.rules_[i];
var value =
'src_lang: "' + this.model_.source_language.code + '" ' +
'trg_lang: "' + this.model_.target_language.code + '" ' +
'source: "' + rule['source'] + '" ' +
'bad_translations: "' + rule.bad_translations + '" ' +
'good_translations: "' + rule.good_translations + '" ' +
'attention_threshold: ' + rule.attention_threshold;
bodyParts.push('rule=' + window.encodeURIComponent(value));
}
// Combine everything together.
return bodyParts.join('&');
}
/**
* Adds the translation response to the list of results.
* @param {!Event} event The event object from the `response` event. This is
* required to access the current response, as there are timing issues when
* accessing the latest response with iron-ajax's `last-response` attribute.
* @private
*/
handleTranslationResponse_(event) {
this.set('fetchingResult', false);
this.push('results', {
response: event.detail.response,
query: this.query_,
model: this.model_,
});
}
/**
* Adds a new rapid response rule to be filled out.
* @private
*/
addRule_() {
this.push('rules_', {
source: '',
bad_translations: '',
good_translations: '',
attention_threshold: 0.9,
});
}
/**
* Deletes a rapid response rule.
* @param {Event} e The event in the dom repeat template element.
* @private
*/
deleteRule_(e) {
let model = e.model;
this.splice('rules_', model.index, 1);
}
}
customElements.define(ExploreView.is, ExploreView);
================================================
FILE: tensor2tensor/insights/polymer/graph_visualization/graph-visualization.html
================================================
<!--
@license
Copyright 2018 The Tensor2Tensor 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.
-->
<link rel="import" href="../polymer/polymer.html">
<link rel="import" href="../iron-ajax/iron-ajax.html">
<link rel="import" href="../iron-flex-layout/iron-flex-layout-classes.html">
<link rel="import" href="../iron-icon/iron-icon.html">
<link rel="import" href="../iron-icons/av-icons.html">
<link rel="import" href="../iron-icons/iron-icons.html">
<link rel="import" href="../paper-button/paper-button.html">
<link rel="import" href="../paper-icon-button/paper-icon-button.html">
<link rel="import" href="../paper-slider/paper-slider.html">
<link rel="import" href="../paper-tooltip/paper-tooltip.html">
<dom-module id="graph-visualization">
<template>
<style include="iron-flex iron-flex-alignment">
#chart {
border: 1px #ccc solid;
position: relative;
}
#help {
position: absolute;
top: 10px;
right: 10px;
z-index: 1;
}
.background {
fill: #eee;
}
line {
stroke: #000;
stroke-width: 1.5px;
}
rect {
fill: transparent;
}
.node circle {
cursor: pointer;
fill: #fff;
stroke: steelblue;
stroke-width: 1.5px;
}
g.selected circle {
fill: lightsteelblue;
}
.node text {
font-size: 12px;
}
path.link {
fill: none;
stroke-width: 1.5px;
}
circle.terminal {
fill: lightsteelblue;
}
circle.nonterminal {
fill: #fff;
}
text {
fill: #222;
}
g.fade circle,
g.fade text {
opacity: 0.1;
}
.svg-container {
display: inline-block;
position: relative;
width: 100%;
padding-bottom: 100%;
vertical-align: top;
overflow: hidden;
}
.svg-content-responsive {
display: inline-block;
position: absolute;
top: 10px;
left: 0px;
}
#info {
background: #fff;
border: 1px solid #bce8f1;
border-radius: 5px;
position: absolute;
right: 12px;
top: 12px;
z-index: 10;
}
#info .header {
background-color: #bce8f1;
border-color: #bce8f1;
color: #31708f;
}
#info .header,
#info .details {
padding: 12px 6px;
}
</style>
<div class="layout horizontal">
<paper-icon-button id="home" on-tap="reset_" icon=home></paper-icon-button>
<paper-slider id="slider" min=0 max=100 value="{{zoomDepth_}}"></paper-slider>
<div class="flex"></div>
<iron-pages selected="[[stepMode]]" attr-for-selected="name">
<div name="view">
<paper-button raised on-tap="startStepMode_">Start Step Decoding</paper-button>
</div>
<div name="edit">
<paper-icon-button on-tap="step_" icon=av:play-arrow></paper-icon-button>
<paper-button raised on-tap="exitStepMode_">Exit Step Decoding</paper-button>
</div>
</iron-pages>
</div>
<div id="chart">
<div id="info">
<div class="header">
Node Details
</div>
<div class="details">
<div>Token: <span>[[currentName]]</span></div>
<div>Token Probability: <span>[[currentProbability]]</span></div>
<div>Total Probability: <span>[[currentTotalProbability]]</span></div>
<div>Score: <span>[[score]]</span></div>
</div>
</div>
</div>
<iron-ajax
id="startAjax"
url="/api/remote_decoder_start"
method="POST"
body="[[startBody]]"
handle-as="json"
on-error="handleStartError_"
on-response="handleStartResponse_"
last-response="{{startResponse_}}">
</iron-ajax>
<iron-ajax
id="generateAjax"
url="/api/remote_decoder_generate"
method="POST"
body="[[generateBody]]"
params="[[generateParams]]"
handle-as="json"
on-error="handleGenerateError_"
on-response="handleGenerateResponse_"
last-response="{{generateResponse_}}">
</iron-ajax>
</template>
<script src="graph-visualization.js"></script>
</dom-module>
================================================
FILE: tensor2tensor/insights/polymer/graph_visualization/graph-visualization.js
================================================
/**
* @license
* Copyright 2018 The Tensor2Tensor 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.
*/
/**
* `<graph-visualization>` Presents a beam search decoding graph.
*
* The Beam Search decoding graph visualizes the entire search space of a
* sequence generation model. Each layer in the graph displays a decoding step
* with nodes in that layer representing generated candidates. If supported by
* the backend server, the graph can enter interactive mode where candidates can
* be selected for each generation step.
*
*
* ### Usage
*
* <graph-visualization data="[[data]]"></graph-visualization>
*/
class GraphVisualization extends Polymer.Element {
constructor() {
super();
/**
* @private
*/
this.svg_ = undefined;
/**
* @private
*/
this.vis_ = undefined;
/**
* @type {!TreeNode}
* @private
*/
this.rootTree_ = {
name: '',
localProbability: 0,
cumalitiveProbability: 0,
score: 0,
attention: [],
children: [],
};
/**
* @type {!InteractiveNode}
* @private
*/
this.interactiveRoot_ = {
id: this.nodeId_,
stepIndex: 0,
candidate: {
label: '<s>',
label_id: 1,
log_probability: 0,
total_log_probability: 0,
score: 0,
parent_id: 0
},
children: [],
};
/**
* @type {Array<!InteractiveNode>}
* @private
*/
this.selectedNodes_ = [];
/**
* @private
*/
this.stepNodes_ = [];
/**
* Metadata for navigating nodes.
* @private
*/
this.nodeId_ = 0;
/**
* D3.js helper object.
* @private
*/
this.partition_ = undefined;
/**
* D3.js helper object.
* @private
*/
this.zoom_ = undefined;
/**
* D3.js DOM element.
* @private
*/
this.container_ = undefined;
}
/**
* @return {string} The component name.
*/
static get is() {
return 'graph-visualization';
}
/**
* @return {!Object} The component properties.
*/
static get properties() {
return {
/**
* @type {!SearchGraphVisualization}
*/
data: {
type: Object,
observer: 'dataUpdated_',
},
/**
* @type {!Model}
*/
model: {
type: Object,
},
/**
* @type {string}
*/
query: {
type: String,
},
/**
* @type {number}
*/
zoomDepth_: {
type: Number,
value: 20,
},
/**
* @type {!StartTranslationResponse}
*/
startResponse_: {
type: Object,
},
/**
* @type {!GenerateCandidateResponse}
*/
generateResponse_: {
type: Object,
},
};
}
/**
* @return {!Array<string>} The component observers.
*/
static get observers() {
return [
'zoomDepthChanged_(zoomDepth_)',
];
}
/**
* Sets the default zoom depth.
* @override
*/
ready() {
super.ready();
this.set('zoomDepth_', 20);
this.set('stepMode', 'view');
}
/**
* Sets the zoom state based on the updated depth.
* @param {number} zoomDepth the zoom depth.
* @private
*/
zoomDepthChanged_(zoomDepth) {
if (!this.svg_) {
return;
}
if (zoomDepth == 0) {
zoomDepth = 0.000001;
}
let transform = d3.zoomTransform(this.svg_.node()).scale(zoomDepth / 20.0);
this.vis_.attr("transform", transform);
}
/**
* Converts the NMT Graph JSON format to a nested tree heirachy and plots the
* tree as a collapsible tree visualization.
* @private
*/
dataUpdated_() {
// We need to determine two key nodes in the graph:
// Root: This is the node with no in links and some out links.
// Term: This is the terminal node with no out links and some in links.
//
// Our plot will associate token with actual nodes. For all nodes except
// the Term node, this will work fine since in the tree, each node is
// referenced only once as the head of an edge.
//
// The Term node however needs to be duplicated for each edge ending at it
// so that each instance can have a unique token associated with it.
// Step 1) Find Root and Term node indices so they can be refered to later.
var rootIndex = -1;
var nodes = this.data.node;
for (var i = 0; i < nodes.length && rootIndex == -1; ++i) {
var node = nodes[i];
if (node.in_edge_index.length == 0 && node.out_edge_index.length != 0) {
rootIndex = i;
}
}
// Step 2) Create the root node in the tree. The tree structure will have
// the following components:
// name: The display name of the node. This will be some token.
// localProbability: The per time step probability of this node.
// cumulativeProbability: The total probability of this path in the beam
// search.
// score: A final score for this path in the beam search. This is
// typically the cumulativeProbability with zero or more penalties.
// attention: The attention vector associated with this node transition.
// children: The list of children in the tree, which are themselves trees.
this.rootTree_ = {
name: '',
localProbability: 0,
cumalitiveProbability: 0,
score: 0,
attention: [],
children: [],
};
// Step3) Add each child and it's children recursively starting from the
// root node.
var rootNode = nodes[rootIndex];
var edges = this.data.edge;
for (var i = 0; i < rootNode.out_edge_index.length; ++i) {
// Get the edge.
var outEdge = edges[rootNode.out_edge_index[i]];
this.addChildToTree_(this.rootTree_, outEdge, nodes, edges);
}
this.propagateLabel_(this.rootTree_);
this.createSVG_();
this.plotTree_(this.rootTree_);
}
/**
* Forwards path labels from a node's child to the current node.
* @param {!TreeNode} node The node to annotate.
* @private
*/
propagateLabel_(node) {
var hasNBest = false;
var hasBeam = false;
var hasAlternative = false;
for (var i = 0; i < node.children.length; ++i) {
hasNBest = hasNBest || node.children[i].pathType == 'nbest';
hasBeam = hasBeam || node.children[i].pathType == 'beam';
hasAlternative = hasAlternative ||
node.children[i].pathType == 'alternative';
}
if (hasNBest) {
node.pathType = 'nbest';
} else if (hasBeam) {
node.pathType = 'beam';
} else if (hasAlternative) {
node.pathType = 'beam';
} else {
node.pathType = 'unknown';
}
}
/**
* Iterates through all the children in tree and adds them as children to the
* top level tree.
* @param {!TreeNode} tree The current node in the tree to update with
* children.
* @param {!BeamSearchEdge} currentEdge The edge going into tree.
* @param {!Array<!BeamSearchNode>} nodes The list of all node objects.
* @param {!Array<!BeamSearchEdge>} edges The list of all edges between nodes.
* @private
*/
addChildToTree_(tree, currentEdge, nodes, edges) {
// The real edge information is nested in wonderfully named proto
// extensions. Extract the extension information appropriately.
var candidate = currentEdge.data;
// When the label for the new child is empty, we're at a terminal sink. So
// we ignore that node and instead label the parent.
if (candidate.label == '') {
tree.pathType = 'alternative';
return;
}
var node = nodes[currentEdge.target_index];
/**
* @type {TreeNode}
*/
var childTree = {
name: candidate.label,
attention: [],
localProbability: Math.pow(Math.E, candidate.log_probability),
cumalitiveProbability: Math.pow(Math.E, candidate.total_log_probability),
score: Math.pow(Math.E, candidate.score),
finished: currentEdge.completed || false,
children: [],
node: node,
edge: currentEdge,
pathType: 'unknown',
};
tree.children.push(childTree);
if (node.out_edge_index.length == 0) {
if (childTree.name == '</s>') {
childTree.pathType = 'nbest';
} else if (childTree.name == '' || candidate.finished) {
childTree.pathType = 'alternative';
} else {
childTree.pathType = 'beam';
}
} else {
for (var i = 0; i < node.out_edge_index.length; ++i) {
// Get the edge.
var outEdge = edges[node.out_edge_index[i]];
this.addChildToTree_(childTree, outEdge, nodes, edges);
this.propagateLabel_(childTree);
}
}
}
/**
* Creates the initial SVG canvas and associated structures. This will remove
* all previous svg elements.
* @private
*/
createSVG_() {
// Create the margins, width, and height.
var maxWidth = 1600;
var maxHeight = 1600;
var margins = [20, 120, 20, 20];
var width = maxWidth - margins[1] - margins[3];
var height = maxHeight - margins[0] - margins[2];
// Use a d3 partition which will place each node based it's number of
// descendents with the highest ranked path along the top.
this.partition_ = d3.partition().size([height, width]).padding(1);
// Set the initial position of the root of the tree to be a half the height
// and on the left..
this.rootTree_.x0 = height / 2;
this.rootTree_.y0 = 0;
this.zoom_ = d3.zoom()
.scaleExtent([1, 10])
.on("zoom", zoomed.bind(this));
d3.select(this.$.chart).selectAll('.svg-container').remove();
// Embed the SVG to host the tree and rotate it so that horizontal matches
// the height of the canvas.
this.svg_ = d3.select(this.$.chart)
.append("div")
.classed("svg-container", true)
.append("svg")
.attr("height", "100%")
.attr("width", "100%")
.classed("svg-content-responsive", true)
.call(this.zoom_)
.on('dblclick.zoom', null)
.on('wheel.zoom', null);
/**
* Note: For reasons not understood, the javascript compiler can't figure
* out the type of _zoomDepth at this line, so we need to coerce it into
* being a number.
* @type {number}
*/
let zoomDepth = parseInt(this.zoomDepth_, 10);
let transform = d3.zoomTransform(this.svg_.node()).scale(zoomDepth / 20.0);
this.vis_ = this.svg_.append('g')
.attr("transform", transform);
// Ensure that the entire svg element can be used for panning.
this.vis_.append("rect")
.attr("width", maxWidth)
.attr("height", maxWidth)
.style("fill", "none")
.style("pointer-events", "all");
this.container_ = this.vis_.append("g");
// Apply the zoom transformation.
function zoomed() {
this.vis_.attr("transform",
d3.event.transform.scale(this.zoomDepth_ / 20.0));
}
}
/**
* Examines and plots all reachable nodes in the rootTree with respect to the
* given current root.
* @param {!TreeNode} root The current root node to focus on.
* @private
*/
plotTree_(root) {
// Create the hierarchy. We accumulate node values by just counting the
// number of elements, rather than placing a weight on each node..
var treeHierachy = d3.hierarchy(this.rootTree_)
.sum(function(d) {
return 1;
})
.sort(function(a, b) {
return a.data.score - b.data.score;
});
this.partition_(treeHierachy);
// Create an enter object where we can add both nodes and links.
var enter = this.container_.selectAll(".node")
.data(treeHierachy.descendants())
.enter();
// Add the nodes in four steps:
// 1) A general group element to hold all node portions.
// 2) A rectangle with no visible elements.
// 3) A circle for the node.
// 4) a text label.
var node = enter.append("g")
.attr("class", function(d) {
return "node" + (d.children ? " node--internal" : " node--leaf");
})
.attr("transform", function(d) {
return "translate(" + d.y0 + "," + d.x0 + ")";
})
.attr('id', function(d, i) { return "g-" + i; });
node.append("rect")
.attr("width", function(d) { return d.y1 - d.y0; })
.attr("height", 24);
node.append("circle")
.attr("r", 10)
.attr("transform", "translate(10, 10)");
node.append("text")
.attr("x", 24)
.attr("y", 13)
.text(function(d) { return d.data.name; });
// Add out links from each node to it's parent. We link two nodes using the
// bottom center of the circle so that the text label can be placed at
// approximately the vertical center of the circle. This gives a decent
// layout while also not hiding any text.
enter.append("path")
.attr("class", "link")
.attr("d", function(d) {
if (!d.parent) { return ""; }
// Pad the placement of the links just below the center. We have to
// use x0 and y0 for location due to partition, which doesn't create
// standard x/y fields.
var nodeX = d.x0 + 16;
var nodeY = d.y0 + 10;
var parentX = d.parent.x0 + 16;
var parentY = d.parent.y0 + 10;
return "M" + + nodeY + "," + nodeX +
"C" + (nodeY + parentY) / 2 + "," + nodeX + " " +
(nodeY + parentY) / 2 + "," + parentX + " " +
parentY + "," + parentX;
})
.style('stroke', function(d) {
// Associate a different path color depend on the path type for the
// node.
if (d.data.pathType == 'unknown')
return '#222';
if (d.data.pathType == 'nbest')
return '#66ff33';
if (d.data.pathType == 'beam')
return '#ccc';
if (d.data.pathType == 'alternative')
return '#ff3300';
});
// Setup hover events on each node to place focus and highligh on the node
// being hovered over. We do this by adding opacity to all other nodes.
var nodes = this.container_.selectAll(".node");
node.on('mouseover', function(d, i) {
nodes.classed('fade', function(d, j) {
return i != j;
});
d3.select(this).classed('hover', true);
this.set('currentName', d.data.name);
this.set(
'currentProbability', this.displayNumber(d.data.localProbability));
this.set(
'currentTotalProbability',
this.displayNumber(d.data.cumalitiveProbability));
this.set('score', this.displayNumber(d.data.score));
}.bind(this))
.on('mouseout', function(d, i) {
nodes.classed("fade", false);
d3.select(this).classed("hover", false);
});
}
/**
* Resets the pan and zoom state.
* @private
*/
reset_() {
if (!this.svg_) {
return;
}
this.svg_.call(this.zoom_.transform, d3.zoomIdentity);
this.set('zoomDepth_', 20);
}
/**
* Returns the number value with only 2 significant digits.
* @param {number} value The value to present.
* @return {string} value with just two significant digits.
*/
displayNumber(value) {
return value.toFixed(2);
}
/**
* Enters step by step decoding mode.
* @private
*/
startStepMode_() {
this.set('stepMode', 'edit');
this.startTranslation_();
}
/**
* Exits step by step decoding mode.
* @private
*/
exitStepMode_() {
this.set('stepMode', 'view');
this.dataUpdated_();
}
/**
* Begins step by step decoding with the current model and query.
* @private
*/
startTranslation_() {
this.set('startBody', JSON.stringify({
model_id: {
language_pair: {
source_language: this.model.source_language.code,
target_language: this.model.target_language.code,
},
name: this.model.id,
},
input: this.query,
}));
this.$.startAjax.generateRequest();
}
/**
* Handles a start error.
* @private
*/
handleStartError_() {
console.log("failed");
}
/**
* Initializes the step by step decoding graph with the root note and makes
* the first generation step.
* @private
*/
handleStartResponse_() {
// Reset the node state and create the root of the tree. Later candidates
// that are returned from the generation call will be added.
this.nodeId_ = 0;
this.interactiveRoot_ = {
id: this.nodeId_,
stepIndex: 0,
candidate: {
label: '<s>',
label_id: 1,
log_probability: 0,
total_log_probability: 0,
score: 0,
parent_id: 0
},
children: [],
};
this.nodeId_++;
// Track which nodes are active and available as inputs to the next
// generation step. These will be updated with the candidates they
// generate.
this.selectedNodes_ = [this.interactiveRoot_];
// Redraw the entire plot with an interactive version.
this.createSVG_();
this.drawInteractiveTree_(this.interactiveRoot_);
// Make the first generation request.
this.step_(true);
}
/**
* Handles a generate ajax error.
* @private
*/
handleGenerateError_() {
console.log("generate failed");
}
/**
* Processes the returned candidates and adds them to the graph.
* @private
*/
handleGenerateResponse_() {
// Add the candidates returned and tag them with unique identifiers so we
// can ensure later generation steps don't try to include candidates that
// can't be proccesed any more (we can only use candidates from the most
// recent generation step as input due to limitations in the remote
// decoder).
let stepIndex = 0;
let newlySelectedNodes = [];
this.stepNodes_ = [];
for (var i = 0; i < this.generateResponse_.candidate_list.length; ++i) {
let selectedNode = this.selectedNodes_[i];
let candidateList = this.generateResponse_.candidate_list[i];
for (var j = 0; j < candidateList.candidate.length && j < 5; ++j) {
let candidate = candidateList.candidate[j];
// Tag the parent id so that the next generate call knows what network
// states to maintain.
candidate.parent_id = i;
let newNode = {
id: this.nodeId_,
stepIndex: stepIndex,
candidate: candidate,
children: [],
};
this.nodeId_++;
stepIndex++;
this.stepNodes_.push(newNode);
selectedNode.children.push(newNode);
// Select the first candidate.
if (j === 0) {
newNode.selected = true;
newlySelectedNodes.push(newNode);
}
}
}
this.selectedNodes_ = newlySelectedNodes;
// Reset the graph.
this.createSVG_();
this.drawInteractiveTree_(this.interactiveRoot_);
}
/**
* Draws the interactive tree.
* @param {InteractiveNode} rootNode The root node to draw out.
* @private
*/
drawInteractiveTree_(rootNode) {
let treeHierachy = d3.hierarchy(rootNode)
.sum(function(d) {
return 1;
})
.sort(function(a, b) {
return b.data.candidate.total_log_probability -
a.data.candidate.total_log_probability;
});
this.partition_(treeHierachy);
// Create an enter object where we can add both nodes and links.
var enter = this.container_.selectAll(".node")
.data(treeHierachy.descendants())
.enter();
// Add the nodes in four steps:
// 1) A general group element to hold all node portions.
// 2) A rectangle with no visible elements.
// 3) A circle for the node.
// 4) a text label.
var node = enter.append("g")
.attr("class", function(d) {
return "node" +
(d.children ? " node--internal" : " node--leaf") +
(d.data.selected ? " selected" : "");
})
.attr("transform", function(d) {
return "translate(" + d.y0 + "," + d.x0 + ")";
})
.attr('id', function(d, i) { return "g-" + i; });
node.append("rect")
.attr("width", function(d) { return d.y1 - d.y0; })
.attr("height", 24);
node.append("circle")
.attr("r", 10)
.attr("transform", "translate(10, 10)");
node.append("text")
.attr("x", 24)
.attr("y", 13)
.text(function(d) { return d.data.candidate.label; });
// Add out links from each node to it's parent. We link two nodes using the
// bottom center of the circle so that the text label can be placed at
// approximately the vertical center of the circle. This gives a decent
// layout while also not hiding any text.
enter.append("path")
.attr("class", "link")
.attr("d", function(d) {
if (!d.parent) { return ""; }
// Pad the placement of the links just below the center. We have to
// use x0 and y0 for location due to partition, which doesn't create
// standard x/y fields.
var nodeX = d.x0 + 16;
var nodeY = d.y0 + 10;
var parentX = d.parent.x0 + 16;
var parentY = d.parent.y0 + 10;
return "M" + + nodeY + "," + nodeX +
"C" + (nodeY + parentY) / 2 + "," + nodeX + " " +
(nodeY + parentY) / 2 + "," + parentX + " " +
parentY + "," + parentX;
})
.style('stroke', '#ccc');
node.on('mouseover', function(d, i) {
this.set('currentName', d.data.candidate.label);
this.set(
'currentProbability',
this.displayNumber(Math.exp(d.data.candidate.log_probability)));
this.set(
'currentTotalProbability',
this.displayNumber(Math.exp(d.data.candidate.total_log_probability)));
this.set('score', this.displayNumber(Math.exp(d.data.candidate.score)));
}.bind(this));
// Store a local pointer to stepNodes and selectedNodes so that the click
// handler can access them without having to replace the 'this' pointer.
// The click handler needs the default 'this' handler to update the state of
// the clicked upon node.
let stepNodes = this.stepNodes_;
let selectedNodes = this.selectedNodes_;
node.on('click', function(d, i) {
// Ignore nodes that fall out of bounds.
let stepIndex = d.data.stepIndex;
if (stepIndex >= stepNodes.length) {
return;
}
// Ignore nodes that are from different steps.
let node = stepNodes[stepIndex];
if (node.id != d.data.id) {
return;
}
// Update the selected state of the node and either add it to the selected
// list or remove it.
if (!node.selected) {
node.selected = true;
selectedNodes.push(node);
} else {
node.selected = false;
selectedNodes.splice(selectedNodes.indexOf(node), 1);
}
d3.select(this).classed('selected', node.selected);
});
}
/**
* Make one generation step with the candidates in the current selectedNodes
* list. If no nodes are selected, this silently does nothing.
* @param {boolean=} opt_skipNext If true, skips the next step during
* generation.
* @private
*/
step_(opt_skipNext) {
// Running generate without any nodes can put the decoder into a bad state
// and make the session unusable, so for now, silently skip this case.
if (this.selectedNodes_.length == 0) {
console.log("Skipping empty step.");
return;
}
this.set('generateParams', {
skip_next: opt_skipNext || false,
});
this.set('generateBody', JSON.stringify({
model_id: {
language_pair: {
source_language: this.model.source_language.code,
target_language: this.model.target_language.code,
},
name: this.model.id,
},
session_id: this.startResponse_.session_id,
candidate: this.selectedNodes_.map(function(node) {
return node.candidate;
}),
}));
this.$.generateAjax.generateRequest();
}
}
customElements.define(GraphVisualization.is, GraphVisualization);
================================================
FILE: tensor2tensor/insights/polymer/index.html
================================================
<!doctype html>
<!--
@license
Copyright 2018 The Tensor2Tensor 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.
-->
<html>
<head>
<meta name="viewport" content="width=device-width, minimum-scale=1.0, initial-scale=1, user-scalable=no">
<meta name="mobile-web-app-capable" content="yes">
<meta name="apple-mobile-web-app-capable" content="yes">
<meta name="apple-touch-fullscreen" content="yes">
<meta name="apple-mobile-web-app-status-bar-style"
content="black-translucent" >
<meta name="format-detection" content="telephone=no">
<title>NMT Research Frontend