# Path: mesh_transformer/build_model.py
# def build_model(params, tpu_name, region, preemptible, version=1):
#     gradient_accumulation_steps = params.get("gradient_accumulation_steps", 1)
#     cores_per_replica = params["cores_per_replica"]
#     tpu_size = params["tpu_size"]
# 
#     warmup_steps = params["warmup_steps"]
#     anneal_steps = params["anneal_steps"]
#     lr = params["lr"]
#     end_lr = params["end_lr"]
#     weight_decay = params["weight_decay"]
# 
#     assert tpu_size in [8, 32, 128, 256, 512]
# 
#     create_tpu(tpu_name, region, f"v3-{tpu_size}", preemptible)
#     assert wait_til(tpu_name, region, {'state': 'READY', 'health': 'HEALTHY'})
# 
#     conns = get_connection(tpu_name, region)
# 
#     assert len(conns) * 8 == tpu_size, "wrong size TPU for config"
# 
#     head_info = ray.init(include_dashboard=False, object_store_memory=10**9)
#     address = head_info['redis_address']
# 
#     with multiprocessing.pool.ThreadPool(processes=len(conns)) as p:
#         p.map(functools.partial(start_ray, address=address, version=version), conns)
# 
#     opt = optax.chain(
#         optax.scale(1 / gradient_accumulation_steps),
#         clip_by_global_norm(1, use_psum=(version == 1)),
#         optax.scale_by_adam(),
#         additive_weight_decay(weight_decay),
#         optax.scale(-1),
#         optax.scale_by_schedule(util.gpt3_schedule(warmup_steps, anneal_steps, lr, end_lr))
#     )
# 
#     params["optimizer"] = opt
# 
#     if version == 2:
#         model_fn = functools.partial(CausalTransformerV2, params)
#     elif version == 1:
#         model_fn = functools.partial(CausalTransformer, params)
#     else:
#         raise Exception(f"Version {version} does not exist")
# 
#     t = TPUCluster((tpu_size // cores_per_replica, cores_per_replica), len(conns), model_fn, version=version)
#     return t
# 
# Path: tasks/eval_harness.py
# class EvalHarnessAdaptor(LM):
#     def greedy_until(self, requests):
#         raise Exception("unimplemented")
# 
#     def loglikelihood_rolling(self, requests):
#         raise Exception("unimplemented")
# 
#     def __init__(self, tpu_cluster, seq, batch, shrink, min_seq=None):
#         super().__init__()
#         self.tpu = tpu_cluster
#         self.seq = seq
#         self.batch = batch
#         self.shrink = shrink
#         self.min_seq = min_seq
# 
#         self.pool = multiprocessing.Pool(initializer=process_init)
#         process_init()
# 
#     def convert_requests(self, requests):
#         return self.pool.imap(partial(process_request, seq=self.seq), requests)
# 
#     def loglikelihood(self, requests):
#         output = []
# 
#         r = self.convert_requests(requests)
#         zero_example = process_request(requests[0], self.seq)
# 
#         for b in tqdm(sample_batch(r, self.batch, zero_example),
#                       desc="LM eval harness",
#                       total=len(requests) // self.batch):
# 
#             if self.shrink:
#                 b = shrink_seq(b, min_seq=self.min_seq)
# 
#             out = self.tpu.eval(b)
# 
#             for loss, correct in zip(out["mask_loss"], out["each_correct"]):
#                 output.append((float(-loss), bool(correct)))
# 
#         return output
# 
# Path: tfrecord_loader.py
# class TFRecordNewInputs(TFRecordLoader):
#     def __init__(self, index_fname, batch_size, sample_size, restore_state=None):
#         def tf_parse(example_proto):
#             features = {
#                 "text": tf.io.VarLenFeature(tf.int64)
#             }
#             parsed_features = tf.io.parse_single_example(example_proto, features)
# 
#             return tf.cast(tf.sparse.to_dense(tf.sparse.reorder(parsed_features["text"])), tf.uint32)
# 
#         super().__init__(index_fname, batch_size, tf_parse, restore_state=restore_state)


import argparse
import json
import time
import numpy as np
import wandb
import multiprocessing
from tqdm import tqdm
from mesh_transformer.build_model import build_model
from lm_eval import evaluator, tasks
from tasks.eval_harness import EvalHarnessAdaptor
from tfrecord_loader import TFRecordNewInputs

    pe = params["pe"]
    assert pe in ["fixed", "rotary", "t5"]

    t = build_model(params, tpu_name, region, preemptible, version=args.version)

    try:
        t.save(0, bucket, model_dir, init=True, overwrite=clean_start)
        step = 0
        train_load_restore = None
    except Exception as e:
        print(f"Save failed with error {e}, trying to load instead...", e)
        step, aux = t.load(bucket, model_dir)
        train_load_restore = aux.get("train_loader", None)

        if train_load_restore is None:
            print("Failed to restore train loader state")

    train_dataset = TFRecordNewInputs(f"data/{params['train_set']}",
                                      batch_size=(
                                          gradient_accumulation_steps,
                                          per_replica_batch * tpu_size // cores_per_replica),
                                      sample_size=params['seq'],
                                      restore_state=train_load_restore)

    global_val_batch = int(per_replica_batch * tpu_size // cores_per_replica * params.get("val_batch_multiplier", 1))

    val_sets = {}

    for k, v in params['val_set'].items():

# Path: mesh_transformer/build_model.py
# def build_model(params, tpu_name, region, preemptible, version=1):
#     gradient_accumulation_steps = params.get("gradient_accumulation_steps", 1)
#     cores_per_replica = params["cores_per_replica"]
#     tpu_size = params["tpu_size"]
# 
#     warmup_steps = params["warmup_steps"]
#     anneal_steps = params["anneal_steps"]
#     lr = params["lr"]
#     end_lr = params["end_lr"]
#     weight_decay = params["weight_decay"]
# 
#     assert tpu_size in [8, 32, 128, 256, 512]
# 
#     create_tpu(tpu_name, region, f"v3-{tpu_size}", preemptible)
#     assert wait_til(tpu_name, region, {'state': 'READY', 'health': 'HEALTHY'})
# 
#     conns = get_connection(tpu_name, region)
# 
#     assert len(conns) * 8 == tpu_size, "wrong size TPU for config"
# 
#     head_info = ray.init(include_dashboard=False, object_store_memory=10**9)
#     address = head_info['redis_address']
# 
#     with multiprocessing.pool.ThreadPool(processes=len(conns)) as p:
#         p.map(functools.partial(start_ray, address=address, version=version), conns)
# 
#     opt = optax.chain(
#         optax.scale(1 / gradient_accumulation_steps),
#         clip_by_global_norm(1, use_psum=(version == 1)),
#         optax.scale_by_adam(),
#         additive_weight_decay(weight_decay),
#         optax.scale(-1),
#         optax.scale_by_schedule(util.gpt3_schedule(warmup_steps, anneal_steps, lr, end_lr))
#     )
# 
#     params["optimizer"] = opt
# 
#     if version == 2:
#         model_fn = functools.partial(CausalTransformerV2, params)
#     elif version == 1:
#         model_fn = functools.partial(CausalTransformer, params)
#     else:
#         raise Exception(f"Version {version} does not exist")
# 
#     t = TPUCluster((tpu_size // cores_per_replica, cores_per_replica), len(conns), model_fn, version=version)
#     return t
# 
# Path: tasks/eval_harness.py
# class EvalHarnessAdaptor(LM):
#     def greedy_until(self, requests):
#         raise Exception("unimplemented")
# 
#     def loglikelihood_rolling(self, requests):
#         raise Exception("unimplemented")
# 
#     def __init__(self, tpu_cluster, seq, batch, shrink, min_seq=None):
#         super().__init__()
#         self.tpu = tpu_cluster
#         self.seq = seq
#         self.batch = batch
#         self.shrink = shrink
#         self.min_seq = min_seq
# 
#         self.pool = multiprocessing.Pool(initializer=process_init)
#         process_init()
# 
#     def convert_requests(self, requests):
#         return self.pool.imap(partial(process_request, seq=self.seq), requests)
# 
#     def loglikelihood(self, requests):
#         output = []
# 
#         r = self.convert_requests(requests)
#         zero_example = process_request(requests[0], self.seq)
# 
#         for b in tqdm(sample_batch(r, self.batch, zero_example),
#                       desc="LM eval harness",
#                       total=len(requests) // self.batch):
# 
#             if self.shrink:
#                 b = shrink_seq(b, min_seq=self.min_seq)
# 
#             out = self.tpu.eval(b)
# 
#             for loss, correct in zip(out["mask_loss"], out["each_correct"]):
#                 output.append((float(-loss), bool(correct)))
# 
#         return output
# 
# Path: tfrecord_loader.py
# class TFRecordNewInputs(TFRecordLoader):
#     def __init__(self, index_fname, batch_size, sample_size, restore_state=None):
#         def tf_parse(example_proto):
#             features = {
#                 "text": tf.io.VarLenFeature(tf.int64)
#             }
#             parsed_features = tf.io.parse_single_example(example_proto, features)
# 
#             return tf.cast(tf.sparse.to_dense(tf.sparse.reorder(parsed_features["text"])), tf.uint32)
# 
#         super().__init__(index_fname, batch_size, tf_parse, restore_state=restore_state)


# Path: train.py
import argparse
import json
import time
import numpy as np
import wandb
import multiprocessing
from tqdm import tqdm
from mesh_transformer.build_model import build_model
from lm_eval import evaluator, tasks
from tasks.eval_harness import EvalHarnessAdaptor
from tfrecord_loader import TFRecordNewInputs


    pe = params["pe"]
    assert pe in ["fixed", "rotary", "t5"]

    t = build_model(params, tpu_name, region, preemptible, version=args.version)

    try:
        t.save(0, bucket, model_dir, init=True, overwrite=clean_start)
        step = 0
        train_load_restore = None
    except Exception as e:
        print(f"Save failed with error {e}, trying to load instead...", e)
        step, aux = t.load(bucket, model_dir)
        train_load_restore = aux.get("train_loader", None)

        if train_load_restore is None:
            print("Failed to restore train loader state")

    train_dataset = TFRecordNewInputs(f"data/{params['train_set']}",
                                      batch_size=(
                                          gradient_accumulation_steps,
                                          per_replica_batch * tpu_size // cores_per_replica),
                                      sample_size=params['seq'],
                                      restore_state=train_load_restore)

    global_val_batch = int(per_replica_batch * tpu_size // cores_per_replica * params.get("val_batch_multiplier", 1))

    val_sets = {}

    for k, v in params['val_set'].items():

        val_sets[k] = TFRecordNewInputs(f"data/{v}",