<commit_before>/* Copyright 2015 Google Inc. All Rights Reserved.

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

    http://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.
==============================================================================*/

// See docs in ../ops/parsing_ops.cc.

#include <unordered_set>

#include <vector>

#include "tensorflow/core/example/example.pb.h"
#include "tensorflow/core/example/feature.pb_text.h"
#include "tensorflow/core/framework/numeric_op.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/core/util/example_proto_helper.h"
#include "tensorflow/core/util/sparse/sparse_tensor.h"
#include "tensorflow/core/util/work_sharder.h"

namespace tensorflow {

class ExampleParserOp : public OpKernel {
 public:
  explicit ExampleParserOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
    // Grab the attributes into member variables.
    OP_REQUIRES_OK(ctx, ctx->GetAttr("sparse_types", &sparse_types_));
    OP_REQUIRES_OK(ctx, ctx->GetAttr("Ndense", &num_dense_));
    OP_REQUIRES_OK(ctx, ctx->GetAttr("Nsparse", &num_sparse_));
    OP_REQUIRES_OK(ctx, ctx->GetAttr("Tdense", &dense_types_));
    OP_REQUIRES_OK(ctx, ctx->GetAttr("dense_shapes", &dense_shapes_));

    OP_REQUIRES(
        ctx, static_cast<size_t>(num_sparse_) == sparse_types_.size(),
        errors::InvalidArgument("len(sparse_keys) != len(sparse_types"));
    OP_REQUIRES(ctx, static_cast<size_t>(num_dense_) == dense_types_.size(),
                errors::InvalidArgument("len(dense_keys) != len(dense_types"));
    OP_REQUIRES(ctx, static_cast<size_t>(num_dense_) == dense_shapes_.size(),
                errors::InvalidArgument("len(dense_keys) != len(dense_shapes"));
    OP_REQUIRES(ctx, num_dense_ <= std::numeric_limits<int32>::max(),
                errors::InvalidArgument("num_dense_ too large"));
    for (const DataType& type : dense_types_) {
      OP_REQUIRES_OK(ctx, CheckValidType(type));
    }
    for (const DataType& type : sparse_types_) {
      OP_REQUIRES_OK(ctx, CheckValidType(type));
    }
  }

  void Compute(OpKernelContext* ctx) override {
    const Tensor* names;
    const Tensor* serialized;
    OpInputList dense_keys;
    OpInputList sparse_keys;
    OpInputList dense_defaults;

    // Grab the input list arguments.
    OP_REQUIRES_OK(ctx, ctx->input("names", &names));
    OP_REQUIRES_OK(ctx, ctx->input("serialized", &serialized));
    OP_REQUIRES_OK(ctx, ctx->input_list("dense_keys", &dense_keys));
    OP_REQUIRES_OK(ctx, ctx->input_list("sparse_keys", &sparse_keys));
    OP_REQUIRES_OK(ctx, ctx->input_list("dense_defaults", &dense_defaults));

    std::vector<string> dense_keys_t(num_dense_);
    std::vector<string> sparse_keys_t(num_sparse_);

    // Check that the input list sizes match the attribute declared sizes.
    CHECK_EQ(dense_keys.size(), num_dense_);
    CHECK_EQ(sparse_keys.size(), num_sparse_);

    // Copy from OpInputList to std::vector<string>.
    for (int di = 0; di < num_dense_; ++di) {
      dense_keys_t[di] = dense_keys[di].scalar<string>()();
    }
    for (int di = 0; di < num_sparse_; ++di) {
      sparse_keys_t[di] = sparse_keys[di].scalar<string>()();
    }

    bool has_names = (names->NumElements() > 0);
    if (has_names) {
      OP_REQUIRES(
          ctx, TensorShapeUtils::IsVector(names->shape()),
          errors::InvalidArgument("Expected names to be a vector, got shape: ",
                                  names->shape().DebugString()));
      OP_REQUIRES(
          ctx, names->NumElements() == serialized->NumElements(),
          errors::InvalidArgument(
              "Expected len(names) == len(serialized), but got: ",
              names->NumElements(), " vs. ", serialized->NumElements()));
    }
    auto names_t = names->flat<string>();

    OP_REQUIRES(ctx, TensorShapeUtils::IsVector(serialized->shape()),
                errors::InvalidArgument(
                    "Expected serialized to be a vector, got shape: ",
                    serialized->shape().DebugString()));
    OP_REQUIRES(ctx, dense_defaults.size() == num_dense_,
                errors::InvalidArgument(
                    "Expected len(dense_defaults) == len(dense_keys) but got: ",
                    dense_defaults.size(), " vs. ", num_dense_));

    for (int d = 0; d < static_cast<int>(num_dense_); ++d) {
      const Tensor& def_value = dense_defaults[d];
      if (def_value.NumElements() > 0) {
        OP_REQUIRES(ctx, def_value.shape() == dense_shapes_[d],
                    errors::InvalidArgument("def_value[", d, "].shape() == ",
                                            def_value.shape().DebugString(),
                                            " != dense_shapes_[", d, "] == ",
                                            dense_shapes_[d].DebugString()));
        OP_REQUIRES(ctx, def_value.dtype() == dense_types_[d],
                    errors::InvalidArgument(
                        "dense_defaults[", d, "].dtype() == ",
                        DataTypeString(def_value.dtype()), " != dense_types_[",
                        d, "] == ", DataTypeString(dense_types_[d])));
      }
    }

    auto serialized_t = serialized->vec<string>();

    const int64 batch_size = serialized_t.size();

    OpOutputList sparse_indices;
    OpOutputList sparse_values;
    OpOutputList sparse_shapes;
    OpOutputList dense_values;

    OP_REQUIRES_OK(ctx, ctx->output_list("sparse_indices", &sparse_indices));
    OP_REQUIRES_OK(ctx, ctx->output_list("sparse_values", &sparse_values));
    OP_REQUIRES_OK(ctx, ctx->output_list("sparse_shapes", &sparse_shapes));
    OP_REQUIRES_OK(ctx, ctx->output_list("dense_values", &dense_values));

    // Setup Dense features and the output_dense_values Tensor* vector.
    std::vector<FixedLenFeature> fixed_len_features(num_dense_);
    std::vector<Tensor*> output_dense_values(num_dense_);

    for (int d = 0; d < num_dense_; ++d) {
      // Preallocate dense_values, since we know their sizes
      TensorShape out_shape;
      out_shape.AddDim(batch_size);
      for (const int64 dim : dense_shapes_[d].dim_sizes())
        out_shape.AddDim(dim);
      Tensor* out = nullptr;
      dense_values.allocate(d, out_shape, &out);

      FixedLenFeature config;
      config.key = dense_keys_t[d];
      config.dtype = dense_types_[d];
      config.shape = dense_shapes_[d];
      config.default_value = dense_defaults[d];
      fixed_len_features[d] = config;
      output_dense_values[d] = dense_values[d];
    }

    // sparse_values_tmp will be num_sparse_ size map of batch_size length
    // tensor vector's, containing the sparse values from the input layer.
    // After these are all stored, we can allocate properly sized outputs
    // and copy data over. Doing it this way saves us the trouble of either
    // performing deserialization twice, or alternatively storing all copies of
    // the full Example protos.
    std::vector<std::vector<Tensor>> sparse_values_tmp(
        num_sparse_, std::vector<Tensor>(batch_size));

    // Setup Sparse features.
    std::vector<VarLenFeature> var_len_features(num_sparse_);
    for (int d = 0; d < num_sparse_; ++d) {
      VarLenFeature config;
      config.key = sparse_keys_t[d];
      config.dtype = sparse_types_[d];
      var_len_features[d] = config;
    }

    auto worker_threads = *(ctx->device()->tensorflow_cpu_worker_threads());

    // Estimate the cost of parsing each batch element.
    int64 work_unit_size = 100 + 100 * num_sparse_;
    for (int d = 0; d < num_dense_; ++d) {
      work_unit_size += dense_shapes_[d].num_elements();
    }

    mutex mu;

    auto DoWork = [&ctx, &mu, &serialized_t, has_names, &names_t,
                   &fixed_len_features, &var_len_features, &output_dense_values,
                   &sparse_values_tmp](int64 start, int64 limit) {
      Example ex;
      // Processing each Example in the batch starts here.
      for (std::size_t b = static_cast<size_t>(start);
           b < static_cast<size_t>(limit); ++b) {
        bool parse_success = ParseProtoUnlimited(&ex, serialized_t(b));
        if (!TF_PREDICT_TRUE(parse_success)) {
          mutex_lock l(mu);
          ctx->CtxFailure(errors::InvalidArgument(
              "Could not parse example input, value: '", serialized_t(b), "'"));
          return;
        }
        const string& example_name = (has_names) ? names_t(b) : "<unknown>";
        Status s = SingleExampleProtoToTensors(
            ex, example_name, b, fixed_len_features, var_len_features,
            &output_dense_values, &sparse_values_tmp);
        if (!TF_PREDICT_TRUE(s.ok())) {
          mutex_lock l(mu);
          ctx->CtxFailureWithWarning(s);
        }
      }
    };
<commit_msg>Adjust unit cost estimate for example_parsing_ops. The estimate was (mainly, I think) underestimating the cost of deserialization and dictionary lookup per dense key.<commit_after>/* Copyright 2015 Google Inc. All Rights Reserved.

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

    http://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.
==============================================================================*/

// See docs in ../ops/parsing_ops.cc.

#include <unordered_set>

#include <vector>

#include "tensorflow/core/example/example.pb.h"
#include "tensorflow/core/example/feature.pb_text.h"
#include "tensorflow/core/framework/numeric_op.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/core/util/example_proto_helper.h"
#include "tensorflow/core/util/sparse/sparse_tensor.h"
#include "tensorflow/core/util/work_sharder.h"

namespace tensorflow {

class ExampleParserOp : public OpKernel {
 public:
  explicit ExampleParserOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
    // Grab the attributes into member variables.
    OP_REQUIRES_OK(ctx, ctx->GetAttr("sparse_types", &sparse_types_));
    OP_REQUIRES_OK(ctx, ctx->GetAttr("Ndense", &num_dense_));
    OP_REQUIRES_OK(ctx, ctx->GetAttr("Nsparse", &num_sparse_));
    OP_REQUIRES_OK(ctx, ctx->GetAttr("Tdense", &dense_types_));
    OP_REQUIRES_OK(ctx, ctx->GetAttr("dense_shapes", &dense_shapes_));

    OP_REQUIRES(
        ctx, static_cast<size_t>(num_sparse_) == sparse_types_.size(),
        errors::InvalidArgument("len(sparse_keys) != len(sparse_types"));
    OP_REQUIRES(ctx, static_cast<size_t>(num_dense_) == dense_types_.size(),
                errors::InvalidArgument("len(dense_keys) != len(dense_types"));
    OP_REQUIRES(ctx, static_cast<size_t>(num_dense_) == dense_shapes_.size(),
                errors::InvalidArgument("len(dense_keys) != len(dense_shapes"));
    OP_REQUIRES(ctx, num_dense_ <= std::numeric_limits<int32>::max(),
                errors::InvalidArgument("num_dense_ too large"));
    for (const DataType& type : dense_types_) {
      OP_REQUIRES_OK(ctx, CheckValidType(type));
    }
    for (const DataType& type : sparse_types_) {
      OP_REQUIRES_OK(ctx, CheckValidType(type));
    }
  }

  void Compute(OpKernelContext* ctx) override {
    const Tensor* names;
    const Tensor* serialized;
    OpInputList dense_keys;
    OpInputList sparse_keys;
    OpInputList dense_defaults;

    // Grab the input list arguments.
    OP_REQUIRES_OK(ctx, ctx->input("names", &names));
    OP_REQUIRES_OK(ctx, ctx->input("serialized", &serialized));
    OP_REQUIRES_OK(ctx, ctx->input_list("dense_keys", &dense_keys));
    OP_REQUIRES_OK(ctx, ctx->input_list("sparse_keys", &sparse_keys));
    OP_REQUIRES_OK(ctx, ctx->input_list("dense_defaults", &dense_defaults));

    std::vector<string> dense_keys_t(num_dense_);
    std::vector<string> sparse_keys_t(num_sparse_);

    // Check that the input list sizes match the attribute declared sizes.
    CHECK_EQ(dense_keys.size(), num_dense_);
    CHECK_EQ(sparse_keys.size(), num_sparse_);

    // Copy from OpInputList to std::vector<string>.
    for (int di = 0; di < num_dense_; ++di) {
      dense_keys_t[di] = dense_keys[di].scalar<string>()();
    }
    for (int di = 0; di < num_sparse_; ++di) {
      sparse_keys_t[di] = sparse_keys[di].scalar<string>()();
    }

    bool has_names = (names->NumElements() > 0);
    if (has_names) {
      OP_REQUIRES(
          ctx, TensorShapeUtils::IsVector(names->shape()),
          errors::InvalidArgument("Expected names to be a vector, got shape: ",
                                  names->shape().DebugString()));
      OP_REQUIRES(
          ctx, names->NumElements() == serialized->NumElements(),
          errors::InvalidArgument(
              "Expected len(names) == len(serialized), but got: ",
              names->NumElements(), " vs. ", serialized->NumElements()));
    }
    auto names_t = names->flat<string>();

    OP_REQUIRES(ctx, TensorShapeUtils::IsVector(serialized->shape()),
                errors::InvalidArgument(
                    "Expected serialized to be a vector, got shape: ",
                    serialized->shape().DebugString()));
    OP_REQUIRES(ctx, dense_defaults.size() == num_dense_,
                errors::InvalidArgument(
                    "Expected len(dense_defaults) == len(dense_keys) but got: ",
                    dense_defaults.size(), " vs. ", num_dense_));

    for (int d = 0; d < static_cast<int>(num_dense_); ++d) {
      const Tensor& def_value = dense_defaults[d];
      if (def_value.NumElements() > 0) {
        OP_REQUIRES(ctx, def_value.shape() == dense_shapes_[d],
                    errors::InvalidArgument("def_value[", d, "].shape() == ",
                                            def_value.shape().DebugString(),
                                            " != dense_shapes_[", d, "] == ",
                                            dense_shapes_[d].DebugString()));
        OP_REQUIRES(ctx, def_value.dtype() == dense_types_[d],
                    errors::InvalidArgument(
                        "dense_defaults[", d, "].dtype() == ",
                        DataTypeString(def_value.dtype()), " != dense_types_[",
                        d, "] == ", DataTypeString(dense_types_[d])));
      }
    }

    auto serialized_t = serialized->vec<string>();

    const int64 batch_size = serialized_t.size();

    OpOutputList sparse_indices;
    OpOutputList sparse_values;
    OpOutputList sparse_shapes;
    OpOutputList dense_values;

    OP_REQUIRES_OK(ctx, ctx->output_list("sparse_indices", &sparse_indices));
    OP_REQUIRES_OK(ctx, ctx->output_list("sparse_values", &sparse_values));
    OP_REQUIRES_OK(ctx, ctx->output_list("sparse_shapes", &sparse_shapes));
    OP_REQUIRES_OK(ctx, ctx->output_list("dense_values", &dense_values));

    // Setup Dense features and the output_dense_values Tensor* vector.
    std::vector<FixedLenFeature> fixed_len_features(num_dense_);
    std::vector<Tensor*> output_dense_values(num_dense_);

    for (int d = 0; d < num_dense_; ++d) {
      // Preallocate dense_values, since we know their sizes
      TensorShape out_shape;
      out_shape.AddDim(batch_size);
      for (const int64 dim : dense_shapes_[d].dim_sizes())
        out_shape.AddDim(dim);
      Tensor* out = nullptr;
      dense_values.allocate(d, out_shape, &out);

      FixedLenFeature config;
      config.key = dense_keys_t[d];
      config.dtype = dense_types_[d];
      config.shape = dense_shapes_[d];
      config.default_value = dense_defaults[d];
      fixed_len_features[d] = config;
      output_dense_values[d] = dense_values[d];
    }

    // sparse_values_tmp will be num_sparse_ size map of batch_size length
    // tensor vector's, containing the sparse values from the input layer.
    // After these are all stored, we can allocate properly sized outputs
    // and copy data over. Doing it this way saves us the trouble of either
    // performing deserialization twice, or alternatively storing all copies of
    // the full Example protos.
    std::vector<std::vector<Tensor>> sparse_values_tmp(
        num_sparse_, std::vector<Tensor>(batch_size));

    // Setup Sparse features.
    std::vector<VarLenFeature> var_len_features(num_sparse_);
    for (int d = 0; d < num_sparse_; ++d) {
      VarLenFeature config;
      config.key = sparse_keys_t[d];
      config.dtype = sparse_types_[d];
      var_len_features[d] = config;
    }

    auto worker_threads = *(ctx->device()->tensorflow_cpu_worker_threads());

    // Estimate the cost of parsing each batch element.
    int64 work_unit_size = 1000 + 100 * num_sparse_;
    for (int d = 0; d < num_dense_; ++d) {
      work_unit_size += 100 + dense_shapes_[d].num_elements();
    }

    mutex mu;

    auto DoWork = [&ctx, &mu, &serialized_t, has_names, &names_t,
                   &fixed_len_features, &var_len_features, &output_dense_values,
                   &sparse_values_tmp](int64 start, int64 limit) {
      Example ex;
      // Processing each Example in the batch starts here.
      for (std::size_t b = static_cast<size_t>(start);
           b < static_cast<size_t>(limit); ++b) {
        bool parse_success = ParseProtoUnlimited(&ex, serialized_t(b));
        if (!TF_PREDICT_TRUE(parse_success)) {
          mutex_lock l(mu);
          ctx->CtxFailure(errors::InvalidArgument(
              "Could not parse example input, value: '", serialized_t(b), "'"));
          return;
        }
        const string& example_name = (has_names) ? names_t(b) : "<unknown>";
        Status s = SingleExampleProtoToTensors(
            ex, example_name, b, fixed_len_features, var_len_features,
            &output_dense_values, &sparse_values_tmp);
        if (!TF_PREDICT_TRUE(s.ok())) {
          mutex_lock l(mu);
          ctx->CtxFailureWithWarning(s);
        }
      }
    };

{'fraction_non_alphanumeric': 0.13278812009039062, 'fraction_numerical': 0.0038738835682771977, 'mean_word_length': 3.2042524315765664, 'pattern_counts': {'":': 0, '<': 65, '<?xml version=': 0, '>': 111, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 14, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2642897', 'n_tokens_mistral': 5828, 'n_tokens_neox': 5660, 'n_words': 1243}