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

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

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

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

@@ -1,5 +1,6 @@
 # mypy: ignore-errors
 import os
+from functools import lru_cache
 from typing import Optional, Union, Callable, List, Dict
 
 import numpy as np
@@ -21,6 +22,7 @@ import inspect
 import functools
 
 _called_from_inside_tl_decorator = 0
+_called_from_inside_tl_integration_test_decorator = False
 
 
 def _add_mapping_connection(user_unique_name, connection_destinations, arg_names, name, node_mapping_type):
@@ -46,17 +48,22 @@ def integration_test():
         leap_binder.integration_test_func = integration_test_function
 
         def inner(*args, **kwargs):
-            ret = integration_test_function(*args, **kwargs)
-
+            global _called_from_inside_tl_integration_test_decorator
             try:
-                os.environ[mapping_runtime_mode_env_var_mame] = 'True'
-                integration_test_function(None, None)
-            except Exception as e:
-                print(f'Error during integration test: Make sure to disable any non tensorleap decorators '
-                      f'functions before pushing a new TL version')
+                _called_from_inside_tl_integration_test_decorator = True
+                ret = integration_test_function(*args, **kwargs)
+
+                try:
+                    os.environ[mapping_runtime_mode_env_var_mame] = 'True'
+                    integration_test_function(None, None)
+                except Exception as e:
+                    print(f'Error during integration test: Make sure to disable any non tensorleap decorators '
+                          f'functions before pushing a new TL version')
+                finally:
+                    if mapping_runtime_mode_env_var_mame in os.environ:
+                        del os.environ[mapping_runtime_mode_env_var_mame]
             finally:
-                if mapping_runtime_mode_env_var_mame in os.environ:
-                    del os.environ[mapping_runtime_mode_env_var_mame]
+                _called_from_inside_tl_integration_test_decorator = False
 
         return inner
 
@@ -71,6 +78,7 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
         class TempMapping:
             pass
 
+        @lru_cache()
         def inner():
             class ModelPlaceholder:
                 def __init__(self):
@@ -85,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():
@@ -116,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():
@@ -762,7 +793,7 @@ def tensorleap_input_encoder(name: str, channel_dim=-1, model_input_index=None):
 
             _validate_result(result)
 
-            if _called_from_inside_tl_decorator == 0:
+            if _called_from_inside_tl_decorator == 0 and _called_from_inside_tl_integration_test_decorator:
                 result = np.expand_dims(result, axis=0)
 
             return result
@@ -846,7 +877,7 @@ def tensorleap_gt_encoder(name: str):
 
             _validate_result(result)
 
-            if _called_from_inside_tl_decorator == 0:
+            if _called_from_inside_tl_decorator == 0 and _called_from_inside_tl_integration_test_decorator:
                 result = np.expand_dims(result, axis=0)
 
             return result

{code_loader-1.0.127 → code_loader-1.0.129}/code_loader/plot_functions/plot_functions.py

@@ -29,7 +29,6 @@ def run_only_on_non_mapping_mode():
     def decorator(func):
         def wrapper(*args, **kwargs):
             if os.environ.get(mapping_runtime_mode_env_var_mame):
-                print(f"Skipping {func.__name__} in mapping mode.")
                 return
             return func(*args, **kwargs)
         return wrapper

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

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