<commit_before># Lint as: python3
# Copyright 2019 Google LLC
#
# 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
#
#      https://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.
"""Test utilities interop with TensorFlow."""

# pylint: disable=not-callable
# pylint: disable=invalid-name
# pylint: disable=protected-access

import collections
import os
import re
import tempfile

from .. import binding
from .. import compiler
from absl import flags
from absl import logging
import numpy as np
import tensorflow.compat.v2 as tf

flags.DEFINE_string(
    "target_backends", None,
    "Explicit comma-delimited list of target backends. "
    "(Overrides environment variables and auto detection)")
FLAGS = flags.FLAGS


def save_and_compile_tf_module(tf_module, exported_names=(),
                               target_backends=()):
  with tempfile.TemporaryDirectory() as sm_path:
    options = tf.saved_model.SaveOptions(save_debug_info=True)
    tf.saved_model.save(tf_module, sm_path, options=options)
    return compiler.tf_compile_saved_model(
        sm_path, exported_names=exported_names, target_backends=target_backends)


def dump_iree_module(m):
  print("Loaded module:", m.name)
  i = 0
  while True:<commit_msg>Adds ability to load a saved model from disk for testing.<commit_after># Lint as: python3
# Copyright 2019 Google LLC
#
# 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
#
#      https://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.
"""Test utilities interop with TensorFlow."""

# pylint: disable=not-callable
# pylint: disable=invalid-name
# pylint: disable=protected-access

import collections
import os
import re
import tempfile

from .. import binding
from .. import compiler
from absl import flags
from absl import logging
import numpy as np
import tensorflow.compat.v2 as tf

flags.DEFINE_string(
    "target_backends", None,
    "Explicit comma-delimited list of target backends. "
    "(Overrides environment variables and auto detection)")
FLAGS = flags.FLAGS

ORIGNAL_SAVED_MODEL_PATH_ATTR = "_ORIGINAL_SAVED_MODEL_PATH"


def save_and_compile_tf_module(tf_module, exported_names=(),
                               target_backends=()):
  """Saves and compiles a TensorFlow tf.Module.

  Note that if the module has the special _ORIGINAL_SAVED_MODEL_PATH attribute,
  then it will be compiled directly from that path instead of saved and then
  loaded.

  Args:
    tf_module: A tf.Module.
    exported_names: Iterable of dotted function names to consider for
      compilation.
    target_backends: Iterable of string backend names to compile for.

  Returns:
    An _IreeCompiledModule.
  """

  def compile_from_path(sm_path):
    return compiler.tf_compile_saved_model(
        sm_path, exported_names=exported_names, target_backends=target_backends)

  if hasattr(tf_module, ORIGNAL_SAVED_MODEL_PATH_ATTR):
    # Compile directly from the original path.
    sm_path = getattr(tf_module, ORIGNAL_SAVED_MODEL_PATH_ATTR)
    logging.info(
        "Compiling from original saved_model path (not round-tripping): %s",
        sm_path)
    return compile_from_path(sm_path)
  else:
    # Round-trip through a temporary director.
    with tempfile.TemporaryDirectory() as sm_path:
      options = tf.saved_model.SaveOptions(save_debug_info=True)
      tf.saved_model.save(tf_module, sm_path, options=options)
      return compile_from_path(sm_path)


def load_tf_module(path):
  """Wrapper around tf.saved_model.load which preserves the path.

  Args:
    path: The path to load from.

  Returns:
    The loaded module with an extra property _ORIGINAL_SAVED_MODEL_PATH added.
    This is used on subsequent compiles to load directly from the original
    path, which gives us unmolested access to the original debug information,
    which TensorFlow tends to lose on round-trip.
  """
  tf_module = tf.saved_model.load(path)
  assert not hasattr(tf_module, ORIGNAL_SAVED_MODEL_PATH_ATTR), (
      "Saved model (%s) already has attribute %s" %
      (path, ORIGNAL_SAVED_MODEL_PATH_ATTR))
  setattr(tf_module, ORIGNAL_SAVED_MODEL_PATH_ATTR, path)
  return tf_module


def dump_iree_module(m):
  print("Loaded module:", m.name)
  i = 0
  while True:
{'fraction_non_alphanumeric': 0.07219357397857992, 'fraction_numerical': 0.00436334787782626, 'mean_word_length': 4.156441717791411, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2643184', 'n_tokens_mistral': 1512, 'n_tokens_neox': 1477, 'n_words': 520}