[PT FE] Support dict on input (#21450)

* [PT FE] Support dict on input

* Check that keys are strings

* Revert changes in PtFrameworkNode

src/frontends/pytorch/src/frontend.cpp

@@ -200,7 +200,8 @@ void FrontEnd::normalize(const std::shared_ptr<ov::Model>& model) const {
     manager.register_pass<ov::frontend::pytorch::pass::IRFFTNComplexReplacer>();
     manager.register_pass<ov::frontend::pytorch::pass::PrimTupleUnpackReplacer>();
     manager.register_pass<ov::frontend::pytorch::pass::DecomposeListTupleResults>();
-    manager.register_pass<ov::frontend::pytorch::pass::DictResolver>();
+    manager.register_pass<ov::frontend::pytorch::pass::DictParameterResolver>();
+    manager.register_pass<ov::frontend::pytorch::pass::DictResultResolver>();
     manager.register_pass<ov::frontend::pytorch::pass::IndexLoopGetitemReplacer>();
     manager.register_pass<ov::frontend::pytorch::pass::QuantizedNodeRemover>();
     manager.register_pass<ov::frontend::pytorch::pass::SoftmaxReshapeElimination>();

src/frontends/pytorch/src/op/getitem.cpp

@@ -19,6 +19,9 @@ using namespace ov::op;
 OutputVector translate_getitem(const NodeContext& context) {
     num_inputs_check(context, 2, 2);
     auto input = context.get_input(0);
+    const auto idx_type = context.get_input_type(1);
+    FRONT_END_OP_CONVERSION_CHECK(!idx_type.is<type::Str>(),
+                                  "String index in aten::__getitem__ means dict input, this is not supported.");
     if (ov::as_type_ptr<ov::op::util::FrameworkNode>(input.get_node_shared_ptr())) {
         FRONT_END_OP_CONVERSION_CHECK(!cast_fw_node(input.get_node_shared_ptr(), "aten::split"),
                                       "special case for aten::__getitem__");

src/frontends/pytorch/src/transforms/dict_resolver.cpp

@@ -5,6 +5,7 @@
 #include "dict_resolver.hpp"
 
 #include "openvino/core/rt_info.hpp"
+#include "openvino/op/parameter.hpp"
 #include "openvino/op/result.hpp"
 #include "openvino/op/util/framework_node.hpp"
 #include "openvino/pass/pattern/op/wrap_type.hpp"
@@ -18,7 +19,56 @@ namespace pass {
 using namespace ov::pass;
 using namespace ov::op;
 
-bool DictResolver::run_on_model(const std::shared_ptr<Model>& model) {
+bool DictParameterResolver::run_on_model(const std::shared_ptr<Model>& model) {
+    bool changed = false;
+    const auto parameters = model->get_parameters();
+    ParameterVector new_params;
+
+    for (const auto& p : parameters) {
+        bool at_least_one_unused = false;
+        if (p->get_output_size() == 1) {
+            const auto targets = p->get_output_target_inputs(0);
+            for (const auto inp : targets) {
+                const auto getitem_node = cast_fw_node(inp.get_node()->shared_from_this(), "aten::__getitem__");
+                if (getitem_node) {
+                    const auto index_node = std::dynamic_pointer_cast<ov::op::util::FrameworkNode>(
+                        getitem_node->get_input_node_shared_ptr(1));
+                    if (!index_node) {
+                        at_least_one_unused = true;
+                        continue;
+                    }
+                    const auto attrs = index_node->get_attrs();
+                    if (attrs.find("string_value") == attrs.end()) {
+                        // index node must contain string value
+                        at_least_one_unused = true;
+                        continue;
+                    }
+                    const auto name = attrs.at("string_value");
+                    auto new_param = std::make_shared<v0::Parameter>(getitem_node->get_output_element_type(0),
+                                                                     getitem_node->get_output_partial_shape(0));
+                    new_param->set_friendly_name(name);
+                    getitem_node->output(0).replace(new_param);
+                    new_params.push_back(new_param);
+                    changed = true;
+                } else {
+                    at_least_one_unused = true;
+                }
+            }
+        }
+        if (changed) {
+            model->remove_parameter(p);
+            if (at_least_one_unused || p->get_output_size() != 1) {
+                new_params.push_back(p);
+            }
+        }
+    }
+    if (changed) {
+        model->add_parameters(new_params);
+    }
+    return changed;
+};
+
+bool DictResultResolver::run_on_model(const std::shared_ptr<Model>& model) {
     bool changed = false;
     const auto results = model->get_results();
     for (const auto& res : results) {

src/frontends/pytorch/src/transforms/dict_resolver.hpp

@@ -12,9 +12,16 @@ namespace frontend {
 namespace pytorch {
 namespace pass {
 
-class DictResolver : public ov::pass::ModelPass {
+// This transformation replaces pattern Parameter(Dict)->aten::__getitem__
+class DictParameterResolver : public ov::pass::ModelPass {
 public:
-    OPENVINO_RTTI("ov::frontend::pytorch::pass::DictResolver");
+    OPENVINO_RTTI("ov::frontend::pytorch::pass::DictParameterResolver");
+    bool run_on_model(const std::shared_ptr<Model>& model) override;
+};
+// This transformation replaces pattern prim::DictConstruct->Result
+class DictResultResolver : public ov::pass::ModelPass {
+public:
+    OPENVINO_RTTI("ov::frontend::pytorch::pass::DictResultResolver");
     bool run_on_model(const std::shared_ptr<Model>& model) override;
 };