1 /* Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_SAVED_MODEL_API_H_ 17 #define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_SAVED_MODEL_API_H_ 18 19 #include <memory> 20 #include <string> 21 #include <unordered_set> 22 #include <vector> 23 24 #include "absl/container/flat_hash_map.h" 25 #include "tensorflow/c/experimental/saved_model/core/concrete_function.h" 26 #include "tensorflow/c/experimental/saved_model/core/signature_def_function.h" 27 #include "tensorflow/cc/saved_model/bundle_v2.h" 28 #include "tensorflow/core/platform/status.h" 29 30 namespace tensorflow { 31 32 // Note(bmzhao): This class is only TEMPORARILY virtual, as a way to unblock 33 // TFRT integration with TF Serving. Do not add more virtual implementations of 34 // this class. Eventually we want to remove this virtual base class indirection 35 // and have only a single implementation. 36 class SavedModelAPI { 37 public: 38 // Retrieve a function from the TF2 SavedModel, using the "path" to a function 39 // in a TF2 savedmodel. 40 // 41 // Note: `function` is a double pointer, so that implementations are 42 // able to return a pointer to an internal member. 43 virtual Status GetFunction(const std::string& function_path, 44 ConcreteFunction** function) = 0; 45 46 // Retrieve a list of child functions from a SavedModel given a starting node. 47 // 0 is the root node. 48 virtual Status GetFunctions( 49 int node_id, 50 absl::flat_hash_map<std::string, ConcreteFunction*>* functions) = 0; 51 52 // Retrieve a SignatureDefFunction from a SavedModel, using the key of the 53 // SignatureDef map: 54 // https://github.com/tensorflow/tensorflow/blob/69b08900b1e991d84bce31f3b404f5ed768f339f/tensorflow/core/protobuf/meta_graph.proto#L89 55 virtual Status GetSignatureDefFunction(const std::string& signature_def_key, 56 SignatureDefFunction** function) = 0; 57 58 virtual SavedModelV2Bundle* GetBundle() = 0; 59 60 virtual ~SavedModelAPI() = default; 61 }; 62 63 } // namespace tensorflow 64 65 #endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_SAVED_MODEL_API_H_ 66