<commit_before># Copyright 2021 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

import sys
from iree import runtime as iree_rt

import numpy.random as npr
import numpy as np
from examples import datasets


def main(args):
  vmfb_file = args[0]
  config = iree_rt.system_api.Config("dylib")
  trainer_module = iree_rt.system_api.load_vm_flatbuffer_file(vmfb_file,
                                                              driver="dylib")
  print(trainer_module)
  print("Random initializing...")
  trainer_module.initialize(np.asarray([34, 66], dtype=np.int32))

  print("Stepping...")
  train_batch = get_examples()
  print(trainer_module.update)
  for i in range(1000):
    batch = next(train_batch)
    trainer_module.update(*batch)
    accuracy = compute_accuracy(batch, trainer_module)
    print(f"Step {i} accuracy = {accuracy}")


def compute_accuracy(batch, trainer_module):
  inputs, targets = batch
  prediction_posterior = trainer_module.predict(inputs)
  target_class = np.argmax(targets, axis=1)
  predicted_class = np.argmax(prediction_posterior, axis=1)
  return np.mean(predicted_class == target_class)


def get_examples():
  batch_size = 128
  train_images, train_labels, test_images, test_labels = datasets.mnist()
  num_train = train_images.shape[0]
  num_complete_batches, leftover = divmod(num_train, batch_size)
  num_batches = num_complete_batches + bool(leftover)

  def data_stream():
    rng = npr.RandomState(0)
    while True:<commit_msg>Add checkpointing.<commit_after># Copyright 2021 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

import os
import sys
from iree import runtime as iree_rt

import numpy.random as npr
import numpy as np
from examples import datasets


def main(args):
  vmfb_file = args[0]
  checkpoint_file = None
  if len(args) > 1:
    checkpoint_file = args[1]
  config = iree_rt.system_api.Config("dylib")
  trainer_module = iree_rt.system_api.load_vm_flatbuffer_file(vmfb_file,
                                                              driver="dylib")
  print(trainer_module)

  if checkpoint_file and os.path.exists(checkpoint_file):
    print(f"Loading checkpoint {checkpoint_file}")
    loaded_arrays = []
    with np.load(checkpoint_file) as checkpoint:
      for _, value in checkpoint.items():
        loaded_arrays.append(value)
    trainer_module.set_opt_state(*loaded_arrays)
  else:
    print("Random initializing...")
    trainer_module.initialize(np.asarray([232, 843], dtype=np.int32))

  print("Stepping...")
  train_batch = get_examples()
  print(trainer_module.update)
  for i in range(1000):
    batch = next(train_batch)
    trainer_module.update(*batch)
    accuracy = compute_accuracy(batch, trainer_module)
    if (i % 50) == 0:
      print(f"Step {i} accuracy = {accuracy}")
    if checkpoint_file:
      np.savez(checkpoint_file, *trainer_module.get_opt_state())


def compute_accuracy(batch, trainer_module):
  inputs, targets = batch
  prediction_posterior = trainer_module.predict(inputs)
  target_class = np.argmax(targets, axis=1)
  predicted_class = np.argmax(prediction_posterior, axis=1)
  return np.mean(predicted_class == target_class)


def get_examples():
  batch_size = 128
  train_images, train_labels, test_images, test_labels = datasets.mnist()
  num_train = train_images.shape[0]
  num_complete_batches, leftover = divmod(num_train, batch_size)
  num_batches = num_complete_batches + bool(leftover)

  def data_stream():
    rng = npr.RandomState(0)
    while True:
{'fraction_non_alphanumeric': 0.10047720042417815, 'fraction_numerical': 0.015641569459172854, 'mean_word_length': 4.057640750670242, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 4, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6694492', 'n_tokens_mistral': 1281, 'n_tokens_neox': 1225, 'n_words': 272}