add dataset scripts

.gitignore

@@ -1,3 +1,5 @@
 .vscode
 venv
 *.pyc
+segment_*
+dataset.csv

convert_files.py

@@ -0,0 +1,17 @@
+import json
+from tqdm import tqdm
+from transformers import AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained('gpt2')
+
+for i in tqdm(range(298)):
+
+    with open(f'wikipedia_json_64_filtered/wikipedia.segmented.nltk.split.seq64.{i}.json', 'r') as f:
+        rows = json.load(f)
+
+    tokens = [row['gpt2_token'] for row in rows]
+    texts = tokenizer.batch_decode(tokens)
+
+    with open(f'wikipedia/{i}.txt', 'w') as f:
+        for txt in texts:
+            f.write(txt.strip() + '\n')

get_data.sh

@@ -0,0 +1,23 @@
+
+
+wget --load-cookies /tmp/cookies.txt "https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id=13NnkYAhwszQxc1C5HHfThnF7c1cjzjAD' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\1\n/p')&id=13NnkYAhwszQxc1C5HHfThnF7c1cjzjAD" -O segment_1.zip && rm -rf /tmp/cookies.txt
+wget --load-cookies /tmp/cookies.txt "https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id=14p6FHip_hGTXC-_7SYaK32BpEhZRDJI4' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\1\n/p')&id=14p6FHip_hGTXC-_7SYaK32BpEhZRDJI4" -O segment_2.zip && rm -rf /tmp/cookies.txt
+wget --load-cookies /tmp/cookies.txt "https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id=1IaRfTFh51Wf_zPtK6tjE6xw-up_Z6EyN' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\1\n/p')&id=1IaRfTFh51Wf_zPtK6tjE6xw-up_Z6EyN" -O segment_3.zip && rm -rf /tmp/cookies.txt
+wget --load-cookies /tmp/cookies.txt "https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id=1KGhV397Xfej56uJ9H10xD7tfLdhWlg4q' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\1\n/p')&id=1KGhV397Xfej56uJ9H10xD7tfLdhWlg4q" -O segment_4.zip && rm -rf /tmp/cookies.txt
+wget --load-cookies /tmp/cookies.txt "https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id=1LfsQ1s9wr1mBG3I1bbvnbyrYmnsrXxZt' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\1\n/p')&id=1LfsQ1s9wr1mBG3I1bbvnbyrYmnsrXxZt" -O segment_5.zip && rm -rf /tmp/cookies.txt
+wget --load-cookies /tmp/cookies.txt "https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id=1OctFe_JPR0Ajh77FzWdfeYnWZinKl2sW' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\1\n/p')&id=1OctFe_JPR0Ajh77FzWdfeYnWZinKl2sW" -O segment_6.zip && rm -rf /tmp/cookies.txt
+wget --load-cookies /tmp/cookies.txt "https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id=1W-Yi8gHCcT8O5F4TcDHScH7pOb0GQZdu' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\1\n/p')&id=1W-Yi8gHCcT8O5F4TcDHScH7pOb0GQZdu" -O segment_7.zip && rm -rf /tmp/cookies.txt
+wget --load-cookies /tmp/cookies.txt "https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id=1jgHjnpe7Vk1pvRgfnH4S4KiRrpUQyqyp' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\1\n/p')&id=1jgHjnpe7Vk1pvRgfnH4S4KiRrpUQyqyp" -O segment_8.zip && rm -rf /tmp/cookies.txt
+wget --load-cookies /tmp/cookies.txt "https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id=1oVst8RG8G2d21DL6q4DwO7aJxE1vA2fc' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\1\n/p')&id=1oVst8RG8G2d21DL6q4DwO7aJxE1vA2fc" -O segment_9.zip && rm -rf /tmp/cookies.txt
+wget --load-cookies /tmp/cookies.txt "https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id=1qwckIM8YBbU9bnArB6bAoStY3e9I1kqU' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\1\n/p')&id=1qwckIM8YBbU9bnArB6bAoStY3e9I1kqU" -O segment_0.zip && rm -rf /tmp/cookies.txt
+
+unzip segment_1.zip
+unzip segment_2.zip
+unzip segment_3.zip
+unzip segment_4.zip
+unzip segment_5.zip
+unzip segment_6.zip
+unzip segment_7.zip
+unzip segment_8.zip
+unzip segment_9.zip
+

merge_datasets.py

@@ -0,0 +1,12 @@
+import datasets
+import pandas as pd
+
+dfs = []
+
+for i in range(10):
+    dfs.append(
+        datasets.ArrowReader.read_table(f'segment_{i}/dataset.arrow').to_pandas()
+    )
+
+full_df = pd.concat(dfs, ignore_index=True)
+full_df.to_csv('dataset.csv')

train.py

@@ -17,7 +17,6 @@
             - [ ] Send the schedule time to the compute_loss method and calculate a coefficient based on that.
 '''
 import logging
-import math
 import os
 import sys
 import time
@@ -31,6 +30,7 @@ from tqdm import tqdm
 
 import jax
 import jax.numpy as jnp
+import numpy as onp
 import optax
 import transformers
 from flax import jax_utils, traverse_util
@@ -44,7 +44,6 @@ from transformers import (
     is_tensorboard_available,
 )
 from transformers.models.t5.modeling_flax_t5 import shift_tokens_right
-from transformers.testing_utils import CaptureLogger
 
 from t5_vae_flax.src.t5_vae import FlaxT5VaeForAutoencoding
 from t5_vae_flax.src.config import T5VaeConfig
@@ -113,7 +112,7 @@ class ModelArguments:
 @dataclass
 class DataTrainingArguments:
     """
-    Arguments pertaining to what data we are going to input our model for training and eval.
+    Arguments pertaining to what data we are going to input our model for training.
     """
 
     dataset_name: Optional[str] = field(
@@ -123,10 +122,6 @@ class DataTrainingArguments:
         default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
     )
     train_file: Optional[str] = field(default=None, metadata={"help": "The input training data file (a text file)."})
-    validation_file: Optional[str] = field(
-        default=None,
-        metadata={"help": "An optional input evaluation data file to evaluate the perplexity on (a text file)."},
-    )
     max_train_samples: Optional[int] = field(
         default=None,
         metadata={
@@ -134,21 +129,8 @@ class DataTrainingArguments:
             "value if set."
         },
     )
-    max_eval_samples: Optional[int] = field(
-        default=None,
-        metadata={
-            "help": "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
-            "value if set."
-        },
-    )
     overwrite_cache: bool = field(
-        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
-    )
-    validation_split_percentage: Optional[int] = field(
-        default=5,
-        metadata={
-            "help": "The percentage of the train set used as validation set in case there's no validation split"
-        },
+        default=False, metadata={"help": "Overwrite the cached training sets"}
     )
     block_size: Optional[int] = field(
         default=None,
@@ -162,7 +144,7 @@ class DataTrainingArguments:
         default=False, metadata={"help": "Stream the dataset."}
     )
     overwrite_cache: bool = field(
-        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
+        default=False, metadata={"help": "Overwrite the cached training sets"}
     )
     preprocessing_num_workers: Optional[int] = field(
         default=None,
@@ -170,15 +152,12 @@ class DataTrainingArguments:
     )
 
     def __post_init__(self):
-        if self.dataset_name is None and self.train_file is None and self.validation_file is None:
-            raise ValueError("Need either a dataset name or a training/validation file.")
+        if self.dataset_name is None and self.train_file is None:
+            raise ValueError("Need either a dataset name or a training file.")
         else:
             if self.train_file is not None:
                 extension = self.train_file.split(".")[-1]
                 assert extension in ["csv", "json", "txt"], "`train_file` should be a csv, a json or a txt file."
-            if self.validation_file is not None:
-                extension = self.validation_file.split(".")[-1]
-                assert extension in ["csv", "json", "txt"], "`validation_file` should be a csv, a json or a txt file."
 
 
 class TrainState(train_state.TrainState):
@@ -188,28 +167,19 @@ class TrainState(train_state.TrainState):
         return jax_utils.replicate(self).replace(dropout_rng=shard_prng_key(self.dropout_rng))
 
 
-def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuffle: bool = False):
+def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int):
     """
     Returns batches of size `batch_size` from truncated `dataset`, sharded over all local devices.
     Shuffle batches if `shuffle` is `True`.
     """
-    steps_per_epoch = len(dataset) // batch_size
-
-    if shuffle:
-        batch_idx = jax.random.permutation(rng, len(dataset))
-    else:
-        batch_idx = jnp.arange(len(dataset))
-
-    batch_idx = batch_idx[: steps_per_epoch * batch_size]  # Skip incomplete batch.
-    batch_idx = batch_idx.reshape((steps_per_epoch, batch_size))
-
-    for idx in batch_idx:
-        batch = dataset[idx]
-        batch = {k: jnp.array(v) for k, v in batch.items()}
-
-        batch = shard(batch)
-
-        yield batch
+    batch = []
+    for row in dataset:
+        batch.append(row)
+        if len(batch) >= batch_size:
+            batch = {k: jnp.stack([row[k] for row in batch]) for k in batch[0].keys()}
+            batch = shard(batch)
+            yield batch
+            batch = []
 
 
 def write_train_metric(summary_writer, train_metrics, train_time, step):
@@ -222,11 +192,6 @@ def write_train_metric(summary_writer, train_metrics, train_time, step):
             summary_writer.scalar(tag, val, step - len(vals) + i + 1)
 
 
-def write_eval_metric(summary_writer, eval_metrics, step):
-    for metric_name, value in eval_metrics.items():
-        summary_writer.scalar(f"eval_{metric_name}", value, step)
-
-
 def create_learning_rate_fn(
     train_ds_size: int, train_batch_size: int, num_train_epochs: int, num_warmup_steps: int, learning_rate: float
 ) -> Callable[[int], jnp.array]:
@@ -284,9 +249,9 @@ def main():
         transformers.utils.logging.set_verbosity_error()
 
     # Set the verbosity to info of the Transformers logger (on main process only):
-    logger.info(f"Training/evaluation parameters {training_args}")
+    logger.info(f"Training parameters {training_args}")
 
-    #  Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
+    # Get the datasets: you can either provide your own CSV/JSON/TXT training files (see below)
     # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
     # (the dataset will be downloaded automatically from the datasets Hub).
     #
@@ -295,35 +260,7 @@ def main():
     #
     # In distributed training, the load_dataset function guarantees that only one local process can concurrently
     # download the dataset.
-    if data_args.dataset_name is not None:
-        # Downloading and loading a dataset from the hub.
-        dataset = load_dataset(
-            data_args.dataset_name, data_args.dataset_config_name, cache_dir=model_args.cache_dir, streaming=data_args.streaming, keep_in_memory=False
-        )
-
-        if "validation" not in dataset.keys():
-            dataset["validation"] = load_dataset(
-                data_args.dataset_name,
-                data_args.dataset_config_name,
-                split=f"train[:{data_args.validation_split_percentage}%]",
-                cache_dir=model_args.cache_dir,
-            )
-            dataset["train"] = load_dataset(
-                data_args.dataset_name,
-                data_args.dataset_config_name,
-                split=f"train[{data_args.validation_split_percentage}%:]",
-                cache_dir=model_args.cache_dir,
-            )
-    else:
-        data_files = {}
-        if data_args.train_file is not None:
-            data_files["train"] = data_args.train_file
-        if data_args.validation_file is not None:
-            data_files["validation"] = data_args.validation_file
-        extension = data_args.train_file.split(".")[-1]
-        if extension == "txt":
-            extension = "text"
-        dataset = load_dataset(extension, data_files=data_files, cache_dir=model_args.cache_dir)
+    dataset = load_dataset('text', data_files=[f'wikipedia/{i}.txt' for i in range(298)], cache_dir=model_args.cache_dir, streaming=True)['train']
     # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
     # https://huggingface.co/docs/datasets/loading_datasets.html.
 
@@ -381,37 +318,6 @@ def main():
         assert tokenizer.pad_token == '<PAD>'
 
     # Preprocessing the datasets.
-    # First we tokenize all the texts.
-    if training_args.do_train:
-        column_names = dataset["train"].column_names
-    else:
-        column_names = dataset["validation"].column_names
-    text_column_name = "text" if "text" in column_names else column_names[0]
-
-    # since this will be pickled to avoid _LazyModule error in Hasher force logger loading before tokenize_function
-    tok_logger = transformers.utils.logging.get_logger("transformers.tokenization_utils_base")
-
-    def tokenize_function(examples):
-        with CaptureLogger(tok_logger) as cl:
-            output = tokenizer(examples[text_column_name])
-        # clm input could be much much longer than block_size
-        if "Token indices sequence length is longer than the" in cl.out:
-            tok_logger.warning(
-                "^^^^^^^^^^^^^^^^ Please ignore the warning above - this long input will be chunked into smaller bits before being passed to the model."
-            )
-        return output
-
-    # remove dataset tasks
-    for k in dataset.keys():
-        dataset[k].info.task_templates = []
-
-    tokenized_datasets = dataset.map(
-        tokenize_function,
-        batched=True,
-        num_proc=data_args.preprocessing_num_workers,
-        remove_columns=column_names,
-        load_from_cache_file=not data_args.overwrite_cache,
-    )
 
     if data_args.block_size > tokenizer.model_max_length:
         logger.warning(
@@ -422,65 +328,27 @@ def main():
 
     pad_token_id, start_token_id = tokenizer.pad_token_id, config.decoder_start_token_id
 
-    def clip_texts(examples):
-        examples["labels"] = examples["input_ids"].copy()
-
-        for i, input_ids in enumerate(examples["input_ids"]):
-            if len(input_ids) > block_size:
-                for k in examples.keys():
-                    examples[k][i] = examples[k][i][:block_size]
-            elif len(input_ids) < block_size:
-                delta = block_size - len(input_ids)
-                examples['input_ids'][i] = examples['input_ids'][i] + [pad_token_id] * delta
-                examples['attention_mask'][i] = examples['attention_mask'][i] + [0] * delta
-                examples['labels'][i] = examples['labels'][i] + [-100] * delta
-
-        return examples
-
-    logger.info('clip_texts...')
-    clipped_lm_datasets = tokenized_datasets.map(
-        clip_texts,
-        batched=True,
-        num_proc=data_args.preprocessing_num_workers,
-        load_from_cache_file=not data_args.overwrite_cache,
-    )
-
-    def add_decoder_input_ids(examples):
-        arr_input_ids = jnp.array(examples["input_ids"])
-        pad = pad_token_id * jnp.ones((arr_input_ids.shape[0], 1), dtype=jnp.int32)
-        arr_pad_input_ids = jnp.concatenate((arr_input_ids, pad), axis=1)
-        examples['decoder_input_ids'] = shift_tokens_right(arr_pad_input_ids, pad_token_id, start_token_id)
+    def tokenize_function(examples):
+        output = tokenizer(examples["text"], return_tensors='jax', padding='max_length', max_length=block_size, truncation=True)
 
-        arr_attention_mask = jnp.array(examples['attention_mask'])
-        ones = jnp.ones((arr_attention_mask.shape[0], 1), dtype=jnp.int32)
-        examples['decoder_attention_mask'] = jnp.concatenate((ones, arr_attention_mask), axis=1)
+        output['labels'] = onp.array(output['input_ids'].copy())
+        output['labels'][output['labels'] == pad_token_id] = -100
+        output['labels'] = jnp.array(output['labels'])
 
-        for k in ['decoder_input_ids', 'decoder_attention_mask']:
-            examples[k] = examples[k].tolist()
+        pad = pad_token_id * jnp.ones((output['input_ids'].shape[0], 1), dtype=jnp.int32)
+        arr_pad_input_ids = jnp.concatenate((output['input_ids'], pad), axis=1)
+        output['decoder_input_ids'] = shift_tokens_right(arr_pad_input_ids, pad_token_id, start_token_id)
 
-        return examples
+        ones = jnp.ones((output['attention_mask'].shape[0], 1), dtype=jnp.int32)
+        output['decoder_attention_mask'] = jnp.concatenate((ones, output['attention_mask']), axis=1)
 
-    logger.info('add_decoder_input_ids...')
-    lm_datasets = clipped_lm_datasets.map(
-        add_decoder_input_ids,
-        batched=True,
-        num_proc=data_args.preprocessing_num_workers,
-        load_from_cache_file=not data_args.overwrite_cache,
-    )
+        return output
 
-    if training_args.do_train:
-        if "train" not in tokenized_datasets:
-            raise ValueError("--do_train requires a train dataset")
-        train_dataset = lm_datasets["train"]
-        if data_args.max_train_samples is not None:
-            train_dataset = train_dataset.select(range(data_args.max_train_samples))
+    tokenized_datasets = dataset.map(tokenize_function, batched=True)
 
-    if training_args.do_eval:
-        if "validation" not in tokenized_datasets:
-            raise ValueError("--do_eval requires a validation dataset")
-        eval_dataset = lm_datasets["validation"]
-        if data_args.max_eval_samples is not None:
-            eval_dataset = eval_dataset.select(range(data_args.max_eval_samples))
+    train_dataset = tokenized_datasets
+    if data_args.max_train_samples is not None:
+        train_dataset = train_dataset.select(range(data_args.max_train_samples))
 
     # Enable tensorboard only on the master node
     has_tensorboard = is_tensorboard_available()
@@ -507,13 +375,13 @@ def main():
     # Store some constant
     num_epochs = int(training_args.num_train_epochs)
     train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count()
-    eval_batch_size = int(training_args.per_device_eval_batch_size) * jax.device_count()
-    steps_per_epoch = len(train_dataset) // train_batch_size
+    train_dataset_len = 97876602
+    steps_per_epoch = train_dataset_len // train_batch_size
     total_train_steps = steps_per_epoch * num_epochs
 
     # Create learning rate schedule
     linear_decay_lr_schedule_fn = create_learning_rate_fn(
-        len(train_dataset),
+        train_dataset_len,
         train_batch_size,
         training_args.num_train_epochs,
         training_args.warmup_steps,
@@ -602,26 +470,14 @@ def main():
 
         return new_state, metrics
 
-    # Define eval fn
-    def eval_step(params, rng, batch):
-        labels = batch.pop("labels")
-        logits, latent_codes = model(**batch, params=params, train=False)[:2]
-        loss = loss_fn(logits, labels, latent_codes, rng)
-
-        # summarize metrics
-        metrics = {"loss": loss}
-        metrics = jax.lax.pmean(metrics, axis_name="batch")
-        return metrics
-
-    # Create parallel version of the train and eval step
+    # Create parallel version of the train step
     p_train_step = jax.pmap(train_step, "batch", donate_argnums=(0,))
-    p_eval_step = jax.pmap(eval_step, "batch")
 
     # Replicate the train state on each device
     state = state.replicate()
 
     logger.info("***** Running training *****")
-    logger.info(f"  Num examples = {len(train_dataset)}")
+    logger.info(f"  Num examples = {train_dataset_len}")
     logger.info(f"  Num Epochs = {num_epochs}")
     logger.info(f"  Instantaneous batch size per device = {training_args.per_device_train_batch_size}")
     logger.info(f"  Total train batch size (w. parallel & distributed) = {train_batch_size}")
@@ -638,15 +494,15 @@ def main():
         rng, input_rng = jax.random.split(rng)
 
         # Generate an epoch by shuffling sampling indices from the train dataset
-        train_loader = data_loader(input_rng, train_dataset, train_batch_size, shuffle=True)
-        steps_per_epoch = len(train_dataset) // train_batch_size
+        train_loader = data_loader(input_rng, train_dataset, train_batch_size)
+        steps_per_epoch = train_dataset_len // train_batch_size
         # train
         for step in tqdm(range(steps_per_epoch), desc="Training...", position=1, leave=False):
             batch = next(train_loader)
             state, train_metric = p_train_step(state, batch)
             train_metrics.append(train_metric)
 
-            cur_step = epoch * (len(train_dataset) // train_batch_size) + step
+            cur_step = epoch * (train_dataset_len // train_batch_size) + step
 
             if cur_step % training_args.logging_steps == 0 and cur_step > 0:
                 # Save metrics
@@ -661,36 +517,6 @@ def main():
 
                 train_metrics = []
 
-            if cur_step % training_args.eval_steps == 0 and cur_step > 0:
-                # ======================== Evaluating ==============================
-                eval_metrics = []
-                eval_loader = data_loader(input_rng, eval_dataset, eval_batch_size)
-                eval_steps = len(eval_dataset) // eval_batch_size
-                for _ in tqdm(range(eval_steps), desc="Evaluating...", position=2, leave=False):
-                    # Model forward
-                    batch = next(eval_loader)
-                    metrics = p_eval_step(state.params, state.dropout_rng, batch)
-                    eval_metrics.append(metrics)
-
-                # normalize eval metrics
-                eval_metrics = get_metrics(eval_metrics)
-                eval_metrics = jax.tree_map(jnp.mean, eval_metrics)
-
-                try:
-                    eval_metrics["perplexity"] = math.exp(eval_metrics["loss"])
-                except OverflowError:
-                    eval_metrics["perplexity"] = float("inf")
-
-                # Print metrics and update progress bar
-                desc = f"Step... ({cur_step} | Eval Loss: {eval_metrics['loss']} | Eval Perplexity: {eval_metrics['perplexity']})"
-                epochs.write(desc)
-                epochs.desc = desc
-
-                # Save metrics
-                if has_tensorboard and jax.process_index() == 0:
-                    cur_step = epoch * (len(train_dataset) // train_batch_size)
-                    write_eval_metric(summary_writer, eval_metrics, cur_step)
-
             if cur_step % training_args.save_steps == 0 and cur_step > 0:
                 # save checkpoint after each epoch and push checkpoint to the hub
                 if jax.process_index() == 0:

train.sh

@@ -0,0 +1,22 @@
+export RUN_NAME=single_latent
+
+# TODO update to not use tokenizer, instead use gpt2 one
+./venv/bin/python train.py \
+--t5_model_name_or_path="t5-base" \
+--output_dir="output/${RUN_NAME}" \
+--overwrite_output_dir \
+--do_train \
+--n_latent_tokens 1 \
+--latent_token_size 32 \
+--save_steps="2000" \
+--block_size="128" \
+--per_device_train_batch_size="100" \
+--train_file="INVALID.txt" \
+--overwrite_output_dir \
+--num_train_epochs="1" \
+
+# 200 batch size, 128 sequence len: ? (breaks)
+# 100 batch size, 128 sequence len: 252:38:58
+# 10 batch size, 128 sequence len: 281:32:53
+
+# Got ~12 hours to train, want 3 saves so one save every 4 hours

wiki_sentences.py

@@ -0,0 +1,46 @@
+# unused
+"""Wikipedia Sentences"""
+
+from __future__ import absolute_import, division, print_function
+
+import os
+import json
+
+import datasets
+
+
+_DESCRIPTION = """\
+Dataset of sentences from Wikipedia (from the [Optimus paper](https://arxiv.org/abs/2004.04092)).
+Each is of mex 64 words & <=256 GPT2 tokens.
+Each row is a tokenised sentence.
+{'token_ids': '{gpt2 token ids}'}
+This is to test the semantics of a Transformer-VAEs latent space by interpolating on sentences.
+"""
+
+NUM_SEGMENTS = 5
+DOWNLOAD_URLS = 'https://drive.google.com/file/d/13NnkYAhwszQxc1C5HHfThnF7c1cjzjAD/view?usp=sharing, https://drive.google.com/file/d/14p6FHip_hGTXC-_7SYaK32BpEhZRDJI4/view?usp=sharing, https://drive.google.com/file/d/1IaRfTFh51Wf_zPtK6tjE6xw-up_Z6EyN/view?usp=sharing, https://drive.google.com/file/d/1KGhV397Xfej56uJ9H10xD7tfLdhWlg4q/view?usp=sharing, https://drive.google.com/file/d/1LfsQ1s9wr1mBG3I1bbvnbyrYmnsrXxZt/view?usp=sharing, https://drive.google.com/file/d/1OctFe_JPR0Ajh77FzWdfeYnWZinKl2sW/view?usp=sharing, https://drive.google.com/file/d/1W-Yi8gHCcT8O5F4TcDHScH7pOb0GQZdu/view?usp=sharing, https://drive.google.com/file/d/1jgHjnpe7Vk1pvRgfnH4S4KiRrpUQyqyp/view?usp=sharing, https://drive.google.com/file/d/1oVst8RG8G2d21DL6q4DwO7aJxE1vA2fc/view?usp=sharing, https://drive.google.com/file/d/1qwckIM8YBbU9bnArB6bAoStY3e9I1kqU/view?usp=sharing'.split(', ')
+
+
+class WikiSentences(datasets.GeneratorBasedBuilder):
+    """Sentences from Wikipedia."""
+
+    BUILDER_CONFIGS = [datasets.BuilderConfig(name="main", description="Run through json files one by one.",)]
+
+    def _info(self):
+        return datasets.DatasetInfo(
+            description=_DESCRIPTION,
+            features=datasets.Features(
+                {
+                    'token_ids': [datasets.Value("int32")],
+                }
+            ),
+            homepage="https://github.com/Fraser-Greenlee/transformer-vae",
+        )
+
+    def _generate_examples(self, filepath):
+        """Generate examples."""
+        with open(filepath, encoding="utf-8") as json_lines_file:
+            for id_, line in enumerate(json_lines_file):
+                yield id_, json.loads(line)
+                if id_ >= self.config.max_num_samples:
+                    break