Add an API in model_task_graph to create or use cached model resources.
PiperOrigin-RevId: 559174528
This commit is contained in:
parent
bcb83302bf
commit
8c4b971c14
|
@ -186,6 +186,21 @@ absl::StatusOr<const ModelResources*> ModelTaskGraph::CreateModelResources(
|
||||||
return model_resources_cache_service.GetObject().GetModelResources(tag);
|
return model_resources_cache_service.GetObject().GetModelResources(tag);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
absl::StatusOr<const ModelResources*> ModelTaskGraph::GetOrCreateModelResources(
|
||||||
|
SubgraphContext* sc, std::unique_ptr<proto::ExternalFile> external_file,
|
||||||
|
std::string tag_suffix) {
|
||||||
|
auto model_resources_cache_service = sc->Service(kModelResourcesCacheService);
|
||||||
|
if (model_resources_cache_service.IsAvailable()) {
|
||||||
|
std::string tag =
|
||||||
|
absl::StrCat(CreateModelResourcesTag(sc->OriginalNode()), tag_suffix);
|
||||||
|
if (model_resources_cache_service.GetObject().Exists(tag)) {
|
||||||
|
return model_resources_cache_service.GetObject().GetModelResources(tag);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return ModelTaskGraph::CreateModelResources(sc, std::move(external_file),
|
||||||
|
tag_suffix);
|
||||||
|
}
|
||||||
|
|
||||||
absl::StatusOr<const ModelAssetBundleResources*>
|
absl::StatusOr<const ModelAssetBundleResources*>
|
||||||
ModelTaskGraph::CreateModelAssetBundleResources(
|
ModelTaskGraph::CreateModelAssetBundleResources(
|
||||||
SubgraphContext* sc, std::unique_ptr<proto::ExternalFile> external_file,
|
SubgraphContext* sc, std::unique_ptr<proto::ExternalFile> external_file,
|
||||||
|
|
|
@ -87,6 +87,20 @@ class ModelTaskGraph : public Subgraph {
|
||||||
SubgraphContext* sc, std::unique_ptr<proto::ExternalFile> external_file,
|
SubgraphContext* sc, std::unique_ptr<proto::ExternalFile> external_file,
|
||||||
std::string tag_suffix = "");
|
std::string tag_suffix = "");
|
||||||
|
|
||||||
|
template <typename Options>
|
||||||
|
absl::StatusOr<const ModelResources*> GetOrCreateModelResources(
|
||||||
|
SubgraphContext* sc, std::string tag_suffix = "") {
|
||||||
|
auto external_file = std::make_unique<proto::ExternalFile>();
|
||||||
|
external_file->Swap(sc->MutableOptions<Options>()
|
||||||
|
->mutable_base_options()
|
||||||
|
->mutable_model_asset());
|
||||||
|
return GetOrCreateModelResources(sc, std::move(external_file), tag_suffix);
|
||||||
|
}
|
||||||
|
|
||||||
|
absl::StatusOr<const ModelResources*> GetOrCreateModelResources(
|
||||||
|
SubgraphContext* sc, std::unique_ptr<proto::ExternalFile> external_file,
|
||||||
|
std::string tag_suffix = "");
|
||||||
|
|
||||||
// If the model resources graph service is available, creates a model asset
|
// If the model resources graph service is available, creates a model asset
|
||||||
// bundle resources object from the subgraph context, and caches the created
|
// bundle resources object from the subgraph context, and caches the created
|
||||||
// model asset bundle resources into the model resources graph service on
|
// model asset bundle resources into the model resources graph service on
|
||||||
|
|
Loading…
Reference in New Issue
Block a user