Add extra scripts

notes/flax_to_pytorch.py

@@ -0,0 +1,35 @@
+import torch
+import numpy as np
+import jax.numpy as jnp
+
+from transformers import AutoTokenizer
+from transformers import FlaxT5ForConditionalGeneration
+from transformers import T5ForConditionalGeneration
+
+tokenizer = AutoTokenizer.from_pretrained("../")
+model_fx = FlaxT5ForConditionalGeneration.from_pretrained("../")
+model_pt = T5ForConditionalGeneration.from_pretrained("../", from_flax=True)
+model_pt.save_pretrained("./")
+
+text = "Hello To You"
+e_input_ids_fx = tokenizer(text, return_tensors="np", padding=True, max_length=128, truncation=True)
+d_input_ids_fx = jnp.ones((e_input_ids_fx.input_ids.shape[0], 1), dtype="i4") * model_fx.config.decoder_start_token_id
+
+e_input_ids_pt = tokenizer(text, return_tensors="pt", padding=True, max_length=128, truncation=True)
+d_input_ids_pt = np.ones((e_input_ids_pt.input_ids.shape[0], 1), dtype="i4") * model_pt.config.decoder_start_token_id
+
+
+print(e_input_ids_fx)
+print(d_input_ids_fx)
+
+print()
+
+encoder_pt = model_fx.encode(**e_input_ids_pt)
+decoder_pt = model_fx.decode(d_input_ids_pt, encoder_pt)
+logits_pt = decoder_pt.logits
+print(logits_pt)
+
+encoder_fx = model_fx.encode(**e_input_ids_fx)
+decoder_fx = model_fx.decode(d_input_ids_fx, encoder_fx)
+logits_fx = decoder_fx.logits
+print(logits_fx)

notes/flax_to_tf.py

@@ -0,0 +1,35 @@
+import torch
+import numpy as np
+import jax.numpy as jnp
+
+from transformers import AutoTokenizer
+from transformers import FlaxT5ForConditionalGeneration
+from transformers import TFT5ForConditionalGeneration
+
+tokenizer = AutoTokenizer.from_pretrained("../")
+model_fx = FlaxT5ForConditionalGeneration.from_pretrained("../")
+model_tf = TFT5ForConditionalGeneration.from_pretrained("./", from_pt=True)
+model_tf.save_pretrained("./")
+
+text = "Hello To You"
+e_input_ids_fx = tokenizer(text, return_tensors="np", padding=True, max_length=128, truncation=True)
+d_input_ids_fx = jnp.ones((e_input_ids_fx.input_ids.shape[0], 1), dtype="i4") * model_fx.config.decoder_start_token_id
+
+e_input_ids_tf = tokenizer(text, return_tensors="tf", padding=True, max_length=128, truncation=True)
+d_input_ids_tf = np.ones((e_input_ids_tf.input_ids.shape[0], 1), dtype="i4") * model_tf.config.decoder_start_token_id
+
+
+print(e_input_ids_fx)
+print(d_input_ids_fx)
+
+print()
+
+encoder_tf = model_fx.encode(**e_input_ids_tf)
+decoder_tf = model_fx.decode(d_input_ids_tf, encoder_tf)
+logits_tf = decoder_tf.logits
+print(logits_tf)
+
+encoder_fx = model_fx.encode(**e_input_ids_fx)
+decoder_fx = model_fx.decode(d_input_ids_fx, encoder_fx)
+logits_fx = decoder_fx.logits
+print(logits_fx)

notes/generation.py

@@ -0,0 +1,196 @@
+import logging
+import os
+import pandas as pd
+import random
+import re
+import sys
+import time
+from dataclasses import dataclass, field
+from functools import partial
+from pathlib import Path
+from typing import Callable, Optional
+
+import jax
+import jax.numpy as jnp
+
+from filelock import FileLock
+from flax import jax_utils, traverse_util
+from flax.jax_utils import unreplicate
+from flax.training import train_state
+from flax.training.common_utils import get_metrics, onehot, shard, shard_prng_key
+
+from transformers import FlaxAutoModelForSeq2SeqLM
+from transformers import AutoTokenizer
+
+from datasets import Dataset, load_dataset, load_metric
+from tqdm import tqdm
+import pandas as pd
+
+
+print(jax.devices())
+
+MODEL_NAME_OR_PATH = "../"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME_OR_PATH, use_fast=True)
+model = FlaxAutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME_OR_PATH)
+
+prefix = "items: "
+text_column = "inputs"
+target_column = "targets"
+max_source_length = 256
+max_target_length = 1024
+seed = 42
+eval_batch_size = 32
+# generation_kwargs = {
+#     "max_length": 1024,
+#     "min_length": 128,
+#     "no_repeat_ngram_size": 3,
+#     "do_sample": True,
+#     "top_k": 60,
+#     "top_p": 0.95
+# }
+generation_kwargs = {
+    "max_length": 1024,
+    "min_length": 128,
+    "no_repeat_ngram_size": 3,
+    "early_stopping": True,
+    "num_beams": 5,
+    "length_penalty": 1.5,
+}
+
+special_tokens = tokenizer.all_special_tokens
+tokens_map = {
+    "<sep>": "--",
+    "<section>": "\n"
+}
+def skip_special_tokens(text, special_tokens):
+    for token in special_tokens:
+        text = text.replace(token, '')
+
+    return text
+
+def target_postprocessing(texts, special_tokens):
+    if not isinstance(texts, list):
+        texts = [texts]
+    
+    new_texts = []
+    for text in texts:
+        text = skip_special_tokens(text, special_tokens)
+
+        for k, v in tokens_map.items():
+            text = text.replace(k, v)
+
+        new_texts.append(text)
+
+    return new_texts
+
+
+predict_dataset = load_dataset("csv", data_files={"test": "/home/m3hrdadfi/code/data/test.csv"}, delimiter="\t")["test"]
+print(predict_dataset)
+# predict_dataset = predict_dataset.select(range(10))
+# print(predict_dataset)
+column_names = predict_dataset.column_names
+print(column_names)
+
+
+# Setting padding="max_length" as we need fixed length inputs for jitted functions
+def preprocess_function(examples):
+    inputs = examples[text_column]
+    targets = examples[target_column]
+    inputs = [prefix + inp for inp in inputs]
+    model_inputs = tokenizer(
+        inputs, 
+        max_length=max_source_length, 
+        padding="max_length", 
+        truncation=True, 
+        return_tensors="np"
+    )
+
+    # Setup the tokenizer for targets
+    with tokenizer.as_target_tokenizer():
+        labels = tokenizer(
+            targets, 
+            max_length=max_target_length, 
+            padding="max_length", 
+            truncation=True, 
+            return_tensors="np"
+        )
+
+    model_inputs["labels"] = labels["input_ids"]
+
+    return model_inputs
+
+predict_dataset = predict_dataset.map(
+    preprocess_function,
+    batched=True,
+    num_proc=None,
+    remove_columns=column_names,
+    desc="Running tokenizer on prediction dataset",
+)
+
+def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuffle: bool = False):
+    """
+    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
+
+rng = jax.random.PRNGKey(seed)
+rng, dropout_rng = jax.random.split(rng)
+rng, input_rng = jax.random.split(rng)
+
+def generate_step(batch):
+    output_ids = model.generate(batch["input_ids"], attention_mask=batch["attention_mask"], **generation_kwargs)
+    return output_ids.sequences
+
+p_generate_step = jax.pmap(generate_step, "batch")
+
+pred_generations = []
+pred_labels = []
+pred_inputs = []
+pred_loader = data_loader(input_rng, predict_dataset, eval_batch_size)
+pred_steps = len(predict_dataset) // eval_batch_size
+
+for _ in tqdm(range(pred_steps), desc="Predicting...", position=2, leave=False):
+    # Model forward
+    batch = next(pred_loader)
+    inputs = batch["input_ids"]
+    labels = batch["labels"]
+
+    generated_ids = p_generate_step(batch)
+    pred_generations.extend(jax.device_get(generated_ids.reshape(-1, generation_kwargs["max_length"])))
+    pred_labels.extend(jax.device_get(labels.reshape(-1, labels.shape[-1])))
+    pred_inputs.extend(jax.device_get(inputs.reshape(-1, inputs.shape[-1])))
+
+inputs = tokenizer.batch_decode(pred_inputs, skip_special_tokens=True)
+true_recipe = target_postprocessing(
+    tokenizer.batch_decode(pred_labels, skip_special_tokens=False), 
+    special_tokens
+)
+generated_recipe = target_postprocessing(
+    tokenizer.batch_decode(pred_generations, skip_special_tokens=False),
+    special_tokens
+)
+test_output = {
+    "inputs": inputs,
+    "true_recipe": true_recipe,
+    "generated_recipe": generated_recipe
+}
+test_output = pd.DataFrame.from_dict(test_output)
+test_output.to_csv("./generated_recipes_b.csv", sep="\t", index=False, encoding="utf-8")