[PyOV] Expose api to Model class (#18141)

* [PyOV] Expose api to Model class

* assign op

* Update src/bindings/python/src/pyopenvino/graph/ops/assign.cpp

* add remove_ methids

* improve test

* codestyle

* assign operation

* codestyle

* test size

---------

Co-authored-by: Michal Lukaszewski <michal.lukaszewski@intel.com>

src/bindings/python/src/openvino/runtime/op/__init__.py

@@ -20,6 +20,7 @@ from openvino._pyopenvino.op import Constant
 """
 Constant.get_data = lambda self: np.array(self, copy=True)
 
+from openvino._pyopenvino.op import assign
 from openvino._pyopenvino.op import Parameter
 from openvino._pyopenvino.op import if_op
 from openvino._pyopenvino.op import loop

src/bindings/python/src/openvino/runtime/opset10/__init__.py

@@ -11,7 +11,7 @@ from openvino.runtime.opset8.ops import adaptive_max_pool
 from openvino.runtime.opset1.ops import add
 from openvino.runtime.opset1.ops import asin
 from openvino.runtime.opset4.ops import asinh
-from openvino.runtime.opset3.ops import assign
+from openvino.runtime.opset6.ops import assign
 from openvino.runtime.opset1.ops import atan
 from openvino.runtime.opset4.ops import atanh
 from openvino.runtime.opset1.ops import avg_pool

src/bindings/python/src/openvino/runtime/opset11/__init__.py

@@ -11,7 +11,7 @@ from openvino.runtime.opset8.ops import adaptive_max_pool
 from openvino.runtime.opset1.ops import add
 from openvino.runtime.opset1.ops import asin
 from openvino.runtime.opset4.ops import asinh
-from openvino.runtime.opset3.ops import assign
+from openvino.runtime.opset6.ops import assign
 from openvino.runtime.opset1.ops import atan
 from openvino.runtime.opset4.ops import atanh
 from openvino.runtime.opset1.ops import avg_pool

src/bindings/python/src/openvino/runtime/opset6/ops.py

@@ -9,7 +9,7 @@ import numpy as np
 from functools import partial
 
 from openvino.runtime import Node, Shape
-from openvino.runtime.op import Constant, Parameter
+from openvino.runtime.op import assign, Constant, Parameter
 from openvino.runtime.opset_utils import _get_node_factory
 from openvino.runtime.utils.decorators import binary_op, nameable_op, unary_op
 from openvino.runtime.utils.input_validation import (
@@ -124,22 +124,6 @@ def mvn(
     return _get_node_factory_opset6().create("MVN", inputs, attributes)
 
 
-@nameable_op
-def assign(new_value: NodeInput, variable_id: str, name: Optional[str] = None) -> Node:
-    """Return a node which produces the Assign operation.
-
-    :param new_value:    Node producing a value to be assigned to a variable.
-    :param variable_id:  Id of a variable to be updated.
-    :param name:         Optional name for output node.
-    :return: Assign node
-    """
-    return _get_node_factory_opset6().create(
-        "Assign",
-        [as_node(new_value)],
-        {"variable_id": variable_id},
-    )
-
-
 @nameable_op
 def read_value(init_value: NodeInput, variable_id: str, name: Optional[str] = None) -> Node:
     """Return a node which produces the Assign operation.

src/bindings/python/src/openvino/runtime/opset7/__init__.py

@@ -9,7 +9,7 @@ from openvino.runtime.opset4.ops import acosh
 from openvino.runtime.opset1.ops import add
 from openvino.runtime.opset1.ops import asin
 from openvino.runtime.opset4.ops import asinh
-from openvino.runtime.opset3.ops import assign
+from openvino.runtime.opset6.ops import assign
 from openvino.runtime.opset1.ops import atan
 from openvino.runtime.opset4.ops import atanh
 from openvino.runtime.opset1.ops import avg_pool

src/bindings/python/src/openvino/runtime/opset8/__init__.py

@@ -11,7 +11,7 @@ from openvino.runtime.opset8.ops import adaptive_max_pool
 from openvino.runtime.opset1.ops import add
 from openvino.runtime.opset1.ops import asin
 from openvino.runtime.opset4.ops import asinh
-from openvino.runtime.opset3.ops import assign
+from openvino.runtime.opset6.ops import assign
 from openvino.runtime.opset1.ops import atan
 from openvino.runtime.opset4.ops import atanh
 from openvino.runtime.opset1.ops import avg_pool

src/bindings/python/src/openvino/runtime/opset9/__init__.py

@@ -11,7 +11,7 @@ from openvino.runtime.opset8.ops import adaptive_max_pool
 from openvino.runtime.opset1.ops import add
 from openvino.runtime.opset1.ops import asin
 from openvino.runtime.opset4.ops import asinh
-from openvino.runtime.opset3.ops import assign
+from openvino.runtime.opset6.ops import assign
 from openvino.runtime.opset1.ops import atan
 from openvino.runtime.opset4.ops import atanh
 from openvino.runtime.opset1.ops import avg_pool

src/bindings/python/src/pyopenvino/graph/model.cpp

@@ -14,6 +14,7 @@
 #include "openvino/core/graph_util.hpp"
 #include "openvino/core/model.hpp"  // ov::Model
 #include "openvino/core/partial_shape.hpp"
+#include "openvino/op/assign.hpp"
 #include "openvino/op/parameter.hpp"  // ov::op::v0::Parameter
 #include "openvino/op/sink.hpp"
 #include "pyopenvino/core/common.hpp"
@@ -49,6 +50,16 @@ static ov::SinkVector cast_to_sink_vector(const std::vector<std::shared_ptr<ov::
     return sinks;
 }
 
+static std::vector<std::shared_ptr<ov::Node>> cast_to_node_vector(const ov::SinkVector& sinks) {
+    std::vector<std::shared_ptr<ov::Node>> nodes;
+    for (const auto& sink : sinks) {
+        auto node = std::dynamic_pointer_cast<ov::Node>(sink);
+        NGRAPH_CHECK(node != nullptr, "Sink {} is not instance of Node");
+        nodes.push_back(node);
+    }
+    return nodes;
+}
+
 void regclass_graph_Model(py::module m) {
     py::class_<ov::Model, std::shared_ptr<ov::Model>> model(m, "Model", py::module_local());
     model.doc() = "openvino.runtime.Model wraps ov::Model";
@@ -696,6 +707,130 @@ void regclass_graph_Model(py::module m) {
                     :rtype: int
                  )");
 
+    model.def("remove_result",
+              &ov::Model::remove_result,
+              py::arg("result"),
+              R"(
+                Delete Result node from the list of results. Method will not delete node from graph.
+
+                :param result: Result node to delete.
+            )");
+
+    model.def("remove_parameter",
+              &ov::Model::remove_parameter,
+              py::arg("parameter"),
+              R"(
+            Delete Parameter node from the list of parameters. Method will not delete node from graph. 
+            You need to replace Parameter with other operation manually.
+
+            Attention: Indexing of parameters can be changed.
+
+            Possible use of method is to replace input by variable. For it the following steps should be done:
+            * `Parameter` node should be replaced by `ReadValue`
+            * call remove_parameter(param) to remove input from the list
+            * check if any parameter indexes are saved/used somewhere, update it for all inputs because indexes can be changed
+            * call graph validation to check all changes
+
+            :param parameter: Parameter node to delete.
+        )");
+
+    model.def(
+        "remove_sink",
+        [](ov::Model& self, const py::object& node) {
+            if (py::isinstance<ov::op::v6::Assign>(node)) {
+                auto sink = std::dynamic_pointer_cast<ov::op::Sink>(node.cast<std::shared_ptr<ov::op::v6::Assign>>());
+                self.remove_sink(sink);
+            } else if (py::isinstance<ov::Node>(node)) {
+                auto sink = std::dynamic_pointer_cast<ov::op::Sink>(node.cast<std::shared_ptr<ov::Node>>());
+                self.remove_sink(sink);
+            } else {
+                throw py::type_error("Incorrect argument type. Sink node is expected as an argument.");
+            }
+        },
+        py::arg("sink"),
+        R"(
+                Delete sink node from the list of sinks. Method doesn't delete node from graph.
+
+                :param sink: Sink to delete.
+        )");
+
+    model.def("add_parameters",
+              &ov::Model::add_parameters,
+              py::arg("parameters"),
+              R"(
+                    Add new Parameter nodes to the list.
+
+                    Method doesn't change or validate graph, it should be done manually.
+                    For example, if you want to replace `ReadValue` node by `Parameter`, you should do the
+                    following steps:
+                    * replace node `ReadValue` by `Parameter` in graph
+                    * call add_parameter() to add new input to the list
+                    * call graph validation to check correctness of changes
+
+                    :param parameter: new Parameter nodes.
+                    :type parameter: List[op.Parameter]
+                 )");
+
+    model.def("add_results",
+              &ov::Model::add_results,
+              py::arg("results"),
+              R"(
+                    Add new Result nodes to the list.
+                    
+                    Method doesn't validate graph, it should be done manually after all changes.
+
+                    :param results: new Result nodes.
+                    :type results: List[op.Result]
+                 )");
+
+    model.def(
+        "add_sinks",
+        [](ov::Model& self, py::list& sinks) {
+            ov::SinkVector sinks_cpp;
+            for (py::handle sink : sinks) {
+                auto sink_cpp =
+                    std::dynamic_pointer_cast<ov::op::Sink>(sink.cast<std::shared_ptr<ov::op::v6::Assign>>());
+                NGRAPH_CHECK(sink_cpp != nullptr, "Assign {} is not instance of Sink");
+                sinks_cpp.push_back(sink_cpp);
+            }
+            self.add_sinks(sinks_cpp);
+        },
+        py::arg("sinks"),
+        R"(
+                    Add new sink nodes to the list.
+                    
+                    Method doesn't validate graph, it should be done manually after all changes.
+
+                    :param sinks: new sink nodes.
+                    :type sinks: List[openvino.runtime.Node]
+                 )");
+
+    model.def(
+        "get_sinks",
+        [](ov::Model& self) {
+            auto sinks = self.get_sinks();
+            return cast_to_node_vector(sinks);
+        },
+        R"(
+            Return a list of model's sinks.
+
+            :return: a list of model's sinks.
+            :rtype: List[openvino.runtime.Node]
+        )");
+
+    model.def_property_readonly(
+        "sinks",
+        [](ov::Model& self) {
+            auto sinks = self.get_sinks();
+            return cast_to_node_vector(sinks);
+        },
+        R"(
+            Return a list of model outputs.
+
+            :return: ResultVector containing model parameters.
+            :rtype: ResultVector
+        )");
+
     model.def(
         "evaluate",
         [](ov::Model& self,

src/bindings/python/src/pyopenvino/graph/ops/assign.cpp

@@ -0,0 +1,43 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "pyopenvino/graph/ops/assign.hpp"
+
+#include <pybind11/pybind11.h>
+
+#include "openvino/op/sink.hpp"
+#include "openvino/op/util/variable.hpp"
+#include "pyopenvino/core/common.hpp"
+
+namespace py = pybind11;
+
+void regclass_graph_op_Assign(py::module m) {
+    py::class_<ov::op::v6::Assign, std::shared_ptr<ov::op::v6::Assign>, ov::Node> assign(m, "assign");
+
+    assign.doc() = "openvino.runtime.op.assign wraps ov::op::v6::Assign";
+
+    assign.def(py::init<>());
+
+    assign.def(py::init([](py::object& new_value, const std::string& variable_id, const std::string& name) {
+                   auto node = new_value.cast<std::shared_ptr<ov::Node>>();
+
+                   auto variable = std::make_shared<ov::op::util::Variable>(
+                       ov::op::util::VariableInfo{ov::PartialShape::dynamic(), ov::element::dynamic, variable_id});
+                   return std::make_shared<ov::op::v6::Assign>(node, variable);
+               }),
+               py::arg("new_value"),
+               py::arg("variable_id"),
+               py::arg("name") = "");
+
+    assign.def("__repr__", [](ov::op::v6::Assign& self) {
+        std::stringstream shapes_ss;
+        for (size_t i = 0; i < self.get_output_size(); ++i) {
+            if (i > 0) {
+                shapes_ss << ", ";
+            }
+            shapes_ss << self.get_output_partial_shape(i);
+        }
+        return "<" + Common::get_class_name(self) + ": '" + self.get_friendly_name() + "' (" + shapes_ss.str() + ")>";
+    });
+}

src/bindings/python/src/pyopenvino/graph/ops/assign.hpp

@@ -0,0 +1,12 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include <pybind11/pybind11.h>
+#include "openvino/op/assign.hpp"
+
+namespace py = pybind11;
+
+void regclass_graph_op_Assign(py::module m);

src/bindings/python/src/pyopenvino/pyopenvino.cpp

@@ -46,6 +46,7 @@
 #include "pyopenvino/graph/discrete_type_info.hpp"
 #include "pyopenvino/graph/layout.hpp"
 #include "pyopenvino/graph/layout_helpers.hpp"
+#include "pyopenvino/graph/ops/assign.hpp"
 #include "pyopenvino/graph/ops/constant.hpp"
 #include "pyopenvino/graph/ops/if.hpp"
 #include "pyopenvino/graph/ops/loop.hpp"
@@ -185,6 +186,7 @@ PYBIND11_MODULE(_pyopenvino, m) {
     regclass_graph_descriptor_Tensor(m);
     regclass_graph_DiscreteTypeInfo(m);
     py::module m_op = m.def_submodule("op", "Package ngraph.impl.op that wraps ov::op");  // TODO(!)
+    regclass_graph_op_Assign(m_op);
     regclass_graph_op_Constant(m_op);
     regclass_graph_op_Parameter(m_op);
     regclass_graph_op_Result(m_op);

src/bindings/python/tests/test_graph/test_basic.py

@@ -515,26 +515,29 @@ def test_multiple_outputs():
     assert list(relu.get_output_shape(0)) == [4, 2]
 
 
-def test_sink_function_ctor():
+def test_sink_model_ctor():
     input_data = ops.parameter([2, 2], name="input_data", dtype=np.float32)
     rv = ops.read_value(input_data, "var_id_667")
     add = ops.add(rv, input_data, name="MemoryAdd")
     node = ops.assign(add, "var_id_667")
     res = ops.result(add, "res")
-    function = Model(results=[res], sinks=[node], parameters=[input_data], name="TestModel")
+    model = Model(results=[res], sinks=[node], parameters=[input_data], name="TestModel")
 
-    ordered_ops = function.get_ordered_ops()
+    ordered_ops = model.get_ordered_ops()
     op_types = [op.get_type_name() for op in ordered_ops]
+    sinks = model.get_sinks()
+    assert ["Assign"] == [sink.get_type_name() for sink in sinks]
+    assert model.sinks[0].get_output_shape(0) == Shape([2, 2])
     assert op_types == ["Parameter", "ReadValue", "Add", "Assign", "Result"]
-    assert len(function.get_ops()) == 5
-    assert function.get_output_size() == 1
-    assert function.get_output_op(0).get_type_name() == "Result"
-    assert function.get_output_element_type(0) == input_data.get_element_type()
-    assert list(function.get_output_shape(0)) == [2, 2]
-    assert (function.get_parameters()[0].get_partial_shape()) == PartialShape([2, 2])
-    assert len(function.get_parameters()) == 1
-    assert len(function.get_results()) == 1
-    assert function.get_friendly_name() == "TestModel"
+    assert len(model.get_ops()) == 5
+    assert model.get_output_size() == 1
+    assert model.get_output_op(0).get_type_name() == "Result"
+    assert model.get_output_element_type(0) == input_data.get_element_type()
+    assert list(model.get_output_shape(0)) == [2, 2]
+    assert (model.get_parameters()[0].get_partial_shape()) == PartialShape([2, 2])
+    assert len(model.get_parameters()) == 1
+    assert len(model.get_results()) == 1
+    assert model.get_friendly_name() == "TestModel"
 
 
 def test_strides_iteration_methods():

src/bindings/python/tests/test_runtime/test_model.py

@@ -7,7 +7,7 @@ import numpy as np
 import pytest
 import math
 
-import openvino.runtime.opset8 as ops
+import openvino.runtime.opset11 as ops
 from openvino.runtime import (
     Core,
     Model,
@@ -620,3 +620,40 @@ def test_serialize_complex_rt_info(request, tmp_path):
 
     os.remove(xml_path)
     os.remove(bin_path)
+
+
+def test_model_add_remove_result_parameter_sink():
+    param = ops.parameter(PartialShape([1]), dtype=np.float32, name="param")
+    relu1 = ops.relu(param, name="relu1")
+    relu2 = ops.relu(relu1, name="relu2")
+    result = ops.result(relu2, "res")
+    model = Model([result], [param], "TestModel")
+
+    result2 = ops.result(relu2, "res2")
+    model.add_results([result2])
+
+    results = model.get_results()
+    assert len(results) == 2
+    assert results[0].get_output_element_type(0) == Type.f32
+    assert results[0].get_output_partial_shape(0) == PartialShape([1])
+
+    model.remove_result(result)
+    assert len(model.results) == 1
+
+    param1 = ops.parameter(PartialShape([1]), name="param1")
+    model.add_parameters([param1])
+
+    params = model.parameters
+    assert (params[0].get_partial_shape()) == PartialShape([1])
+    assert len(params) == 2
+
+    model.remove_parameter(param)
+    assert len(model.parameters) == 1
+
+    assign = ops.assign()
+    model.add_sinks([assign])
+
+    assign_nodes = model.sinks
+    assert ["Assign"] == [sink.get_type_name() for sink in assign_nodes]
+    model.remove_sink(assign)
+    assert len(model.sinks) == 0