code-loader 1.0.128__tar.gz → 1.0.129__tar.gz

{code_loader-1.0.128 → code_loader-1.0.129}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: code-loader
-Version: 1.0.128
+Version: 1.0.129
 Summary: 
 Home-page: https://github.com/tensorleap/code-loader
 License: MIT

{code_loader-1.0.128 → code_loader-1.0.129}/code_loader/inner_leap_binder/leapbinder_decorators.py

@@ -93,6 +93,9 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
                 def run(self, output_names, input_dict):
                     return self.model.run(output_names, input_dict)
 
+                def get_inputs(self):
+                    return self.model.get_inputs()
+
             return ModelPlaceholder()
 
         def mapping_inner():
@@ -124,11 +127,31 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
                     assert output_names is None
                     assert isinstance(input_dict, dict), \
                         f'Expected input_dict to be a dict, got {type(input_dict)} instead.'
-                    for i, elem in enumerate(input_dict.values()):
-                        elem.node_mapping.type = NodeMappingType[f'Input{str(i)}']
+                    for i, (input_key, elem) in enumerate(input_dict.items()):
+                        if isinstance(input_key, NodeMappingType):
+                            elem.node_mapping.type = input_key
+                        else:
+                            elem.node_mapping.type = NodeMappingType[f'Input{str(i)}']
 
                     return ModelOutputPlaceholder()
 
+                def get_inputs(self):
+                    class FollowIndex:
+                        def __init__(self, index):
+                            self.name =  NodeMappingType[f'Input{str(i)}']
+
+                    class FollowInputIndex:
+                        def __init__(self):
+                            pass
+
+                        def __getitem__(self, index):
+                            assert isinstance(index, int), \
+                                f'Expected key to be an int, got {type(index)} instead.'
+
+                            return FollowIndex(index)
+
+                    return FollowInputIndex()
+
             return ModelPlaceholder()
 
         def final_inner():

{code_loader-1.0.128 → code_loader-1.0.129}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "code-loader"
-version = "1.0.128"
+version = "1.0.129"
 description = ""
 authors = ["dorhar <doron.harnoy@tensorleap.ai>"]
 license = "MIT"