code-loader 1.0.202.dev1__tar.gz → 1.0.203.dev0__tar.gz

{code_loader-1.0.202.dev1 → code_loader-1.0.203.dev0}/PKG-INFO

@@ -1,9 +1,9 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.1
 Name: code-loader
-Version: 1.0.202.dev1
+Version: 1.0.203.dev0
 Summary: 
+Home-page: https://github.com/tensorleap/code-loader
 License: MIT
-License-File: LICENSE
 Author: dorhar
 Author-email: doron.harnoy@tensorleap.ai
 Requires-Python: >=3.8,<3.13
@@ -20,7 +20,6 @@ Requires-Dist: numpy (>=2.3.2,<3.0.0) ; python_version >= "3.11" and python_vers
 Requires-Dist: psutil (>=5.9.5,<6.0.0)
 Requires-Dist: pyyaml (>=6.0.2,<7.0.0)
 Requires-Dist: requests (>=2.32.3,<3.0.0)
-Project-URL: Homepage, https://github.com/tensorleap/code-loader
 Project-URL: Repository, https://github.com/tensorleap/code-loader
 Description-Content-Type: text/markdown
 

{code_loader-1.0.202.dev1 → code_loader-1.0.203.dev0}/code_loader/contract/datasetclasses.py

@@ -332,11 +332,6 @@ class AutoregressiveStepHandler:
     # time. Fills the role InputHandler.shape plays for input encoders.
     input_shapes: Optional[Dict[str, List[int]]] = None
     latent_space_aggregation: str = 'last_step'
-    # {model_input_name: model_output_name} for a KV-cache (or any threaded state): inputs whose
-    # value is the named output from the previous step. The platform threads these itself — they
-    # are never passed to the hook and never reported — so the hook must NOT return them. Empty
-    # when the model has no cache, in which case the chain re-feeds its full prompt each step.
-    cache_tensors: Optional[Dict[str, str]] = None
 
 
 # Per-chain, unbatched callables: called once per finished chain with the final step's tensors —

{code_loader-1.0.202.dev1 → code_loader-1.0.203.dev0}/code_loader/inner_leap_binder/leapbinder.py

@@ -590,41 +590,8 @@ class LeapBinder:
         self.setup_container.custom_latent_spaces[name] = CustomLatentSpaceHandler(
             function, name, use_ls_for_analysis)
 
-    @staticmethod
-    def _validated_cache_tensors(cache_tensors: Optional[Dict[str, str]]
-                                 ) -> Optional[Dict[str, str]]:
-        """Structural checks only — shape/dtype/existence are validated by the platform, which
-        is the side that has the model graph. This catches the typos that are cheap to catch
-        here, at integration-test time, instead of at push."""
-        if not cache_tensors:
-            return None
-        if not isinstance(cache_tensors, dict):
-            raise Exception(f'tensorleap_autoregressive_step: cache_tensors must be a dict of '
-                            f'{{model_input_name: model_output_name}}, got '
-                            f'{type(cache_tensors).__name__}.')
-        seen_outputs: Dict[str, str] = {}
-        for input_name, output_name in cache_tensors.items():
-            for role, name in (('key', input_name), ('value', output_name)):
-                if not isinstance(name, str) or not name:
-                    raise Exception(
-                        f'tensorleap_autoregressive_step: cache_tensors {role}s must be '
-                        f'non-empty model tensor names, got {name!r}.')
-            if input_name == output_name:
-                raise Exception(
-                    f'tensorleap_autoregressive_step: cache_tensors maps "{input_name}" to '
-                    f'itself. A model input and output cannot share a name — the mapping is '
-                    f'{{input_name: output_name}}.')
-            if output_name in seen_outputs:
-                raise Exception(
-                    f'tensorleap_autoregressive_step: model output "{output_name}" is declared '
-                    f'as the cache source for both "{seen_outputs[output_name]}" and '
-                    f'"{input_name}". Each cache output feeds exactly one input.')
-            seen_outputs[output_name] = input_name
-        return dict(cache_tensors)
-
     def set_autoregressive_step(self, function: AutoregressiveStepCallableInterface,
-                                latent_space_aggregation: str = 'last_step',
-                                cache_tensors: Optional[Dict[str, str]] = None) -> None:
+                                latent_space_aggregation: str = 'last_step') -> None:
         """
         Set the autoregressive step hook — the feedback function that drives a chain:
         it supplies the model's initial inputs on its first call (prev_inputs=None, prev_outputs=None)
@@ -632,8 +599,6 @@ class LeapBinder:
         An autoregressive integration has no input encoders; the hook is the sole input source.
         latent_space_aggregation declares how the chain's latent-space vectors are derived from
         its steps (see AUTOREGRESSIVE_LATENT_SPACE_AGGREGATIONS).
-        cache_tensors declares a KV-cache as {model_input_name: model_output_name}; the platform
-        threads those tensors itself and the hook neither receives nor returns them.
         """
         if self.setup_container.autoregressive_step is not None:
             raise Exception('tensorleap_autoregressive_step is already defined. '
@@ -642,10 +607,8 @@ class LeapBinder:
             raise Exception(f'tensorleap_autoregressive_step: unknown latent_space_aggregation '
                             f'{latent_space_aggregation!r}. Supported values: '
                             f'{", ".join(AUTOREGRESSIVE_LATENT_SPACE_AGGREGATIONS)}.')
-        cache_tensors = self._validated_cache_tensors(cache_tensors)
         self.setup_container.autoregressive_step = AutoregressiveStepHandler(
-            function, latent_space_aggregation=latent_space_aggregation,
-            cache_tensors=cache_tensors)
+            function, latent_space_aggregation=latent_space_aggregation)
 
         # Builtin chain metadata, declared at parse time so it survives the reporter's
         # metadata type mapping; the placeholder values are overwritten by the engine when a
@@ -1054,17 +1017,18 @@ class LeapBinder:
 
         Each entry must be one of:
           - "foreground" (the friendly name for the max-pooled default LS)
+          - "balanced" (the friendly name for the universal fallback default LS),
+            allowed only when a custom latent space is registered with
+            use_ls_for_analysis=True — that flagged LS is what replaces it
           - An input name registered via set_input(name=...)
           - A prediction name registered via add_prediction(name=...)
 
         Forbidden tokens (raise with a clear message):
-          - "balanced" — the friendly name for the universal fallback default
-            LS; cannot be disabled.
           - "user_custom" — populated by set_custom_latent_space; if the user
             registered a custom LS they should use it, not ignore it.
 
-        Order-independent: this runs after all set_input / add_prediction calls
-        because it consults setup_container.inputs and prediction_types.
+        Order-independent: this runs after all set_input / add_prediction /
+        set_custom_latent_space calls because it consults setup_container.
         """
         if not self.leap_analysis_configuration:
             return
@@ -1072,24 +1036,37 @@ class LeapBinder:
         if not names:
             return
 
-        forbidden = {"balanced", "user_custom"}
+        forbidden = {"user_custom"}
         input_names = {h.name for h in self.setup_container.inputs}
         prediction_names = {h.name for h in self.setup_container.prediction_types}
-        allowed_tokens = {"foreground"} | input_names | prediction_names
+        allowed_tokens = {"foreground", "balanced"} | input_names | prediction_names
+        analysis_ls_names = [
+            ls_name
+            for ls_name, handler in self.setup_container.custom_latent_spaces.items()
+            if handler.use_ls_for_analysis
+        ]
 
         for name in names:
             if name in forbidden:
                 raise Exception(
-                    f"Latent space '{name}' cannot be ignored. "
-                    f"'balanced' is the universal fallback default LS and 'user_custom' is "
-                    f"reserved for set_custom_latent_space."
+                    f"Latent space '{name}' cannot be ignored. 'user_custom' is reserved "
+                    f"for set_custom_latent_space."
+                )
+            if name == "balanced" and not analysis_ls_names:
+                raise Exception(
+                    "Latent space 'balanced' (the model-derived default latent space) can "
+                    "only be ignored when a replacement is registered. Add "
+                    "@tensorleap_custom_latent_space(..., use_ls_for_analysis=True) to the "
+                    "latent space that should be analyzed instead, or drop 'balanced' from "
+                    "ignore_latent_spaces — ignoring it with nothing in its place would "
+                    "leave the analyses with no latent space."
                 )
             if name not in allowed_tokens:
                 raise Exception(
                     f"Latent space name '{name}' in ignore_latent_spaces did not match "
                     f"any registered input or prediction. Available names: "
                     f"inputs={sorted(input_names)}, predictions={sorted(prediction_names)}, "
-                    f"tokens=['foreground']."
+                    f"tokens=['foreground', 'balanced']."
                 )
 
     def set_batch_size_to_validate(self, batch_size: int) -> None:

{code_loader-1.0.202.dev1 → code_loader-1.0.203.dev0}/code_loader/inner_leap_binder/leapbinder_decorators.py

@@ -560,22 +560,12 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
                     return result
 
                 def validate_declared_prediction_types(self, ret):
-                    n_outputs = len(ret) if isinstance(ret, list) else 1
-                    # A declared KV-cache is plumbing: those outputs are fed straight back into
-                    # the model and never reported, so they carry no prediction type. Count only
-                    # the rest, or every cached model would demand 60 meaningless declarations.
-                    step = leap_binder.setup_container.autoregressive_step
-                    n_cache = len((step.cache_tensors if step else None) or {})
-                    if len(self.prediction_types) != 0 and \
-                            n_outputs - n_cache != len(self.prediction_types):
+                    if not (len(self.prediction_types) == len(ret) if isinstance(ret, list) else 1) and len(
+                            self.prediction_types) != 0:
                         if not _call_from_tl_platform:
                             update_env_params_func("tensorleap_load_model", "x")
-                        cache_note = (f' ({n_cache} of them are the declared KV-cache, which '
-                                      f'needs none)' if n_cache else '')
                         raise Exception(
-                            f"tensorleap_load_model validation failed: number of declared "
-                            f"prediction types({len(prediction_types)}) != number of model "
-                            f"outputs({n_outputs}){cache_note}")
+                            f"tensorleap_load_model validation failed: number of declared prediction types({len(prediction_types)}) != number of model outputs({len(ret) if isinstance(ret, list) else 1})")
 
                 def _convert_onnx_inputs_to_correct_type(
                         self, float_arrays_inputs: Dict[str, np.ndarray]
@@ -733,22 +723,12 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
                     assert isinstance(input_dict, dict), \
                         f'Expected input_dict to be a dict, got {type(input_dict)} instead.'
                     seen: set = set()
-                    # A declared KV-cache is threaded by the platform, not wired from an input
-                    # encoder, so those slots have no mapping node to assign — the loop feeds
-                    # them real arrays even in mapping mode. Skip them, and count only the wired
-                    # inputs so the remaining slot indices stay contiguous.
-                    step_handler = leap_binder.setup_container.autoregressive_step
-                    cache_tensors = (step_handler.cache_tensors if step_handler else None) or {}
-                    slot = 0
-                    for input_key, elem in input_dict.items():
-                        if input_key in cache_tensors:
-                            continue
+                    for i, (input_key, elem) in enumerate(input_dict.items()):
                         self._reject_reused_input_source(elem, seen)
                         if isinstance(input_key, NodeMappingType):
                             elem.node_mapping.type = input_key
                         else:
-                            elem.node_mapping.type = _safe_get_item(slot)
-                        slot += 1
+                            elem.node_mapping.type = _safe_get_item(i)
 
                     return ModelOutputPlaceholder()
 
@@ -1904,16 +1884,11 @@ class _MappingStatePlaceholder:
 
 
 class _ModelLoopContext:
-    def __init__(self, sample_id, preprocess_response, prediction_names, is_mapping,
-                 cache_tensors=None):
+    def __init__(self, sample_id, preprocess_response, prediction_names, is_mapping):
         self.sample_id = sample_id
         self.preprocess_response = preprocess_response
         self.prediction_names = prediction_names
         self.is_mapping = is_mapping
-        # Declared KV-cache, {model_input_name: model_output_name}. These tensors are supplied
-        # by the loop (and by the platform at runtime), NOT by the hook, so both the fed-inputs
-        # check and the output-count check have to account for them.
-        self.cache_tensors = dict(cache_tensors or {})
         self.phase = 'awaiting_first_hook'
         self.steps = 0
         self.fed_inputs = None
@@ -1922,9 +1897,6 @@ class _ModelLoopContext:
         self.last_outputs_fingerprint = None
         self.last_state = None
         self.last_state_fingerprint = None
-        # Declared-cache outputs of the last model call, in graph order (predictions first, so
-        # this is the tail). Read by the loop through the model proxy.
-        self.last_cache_outputs = []
 
     def _fail(self, message):
         raise LeapValidationError(f'tensorleap_model_loop validation failed: {message}')
@@ -2009,19 +1981,12 @@ class _ModelLoopContext:
                        'next_inputs the loop must exit without further hook or model calls.')
         if self.is_mapping:
             return
-        # A declared KV-cache is threaded by the platform, so the hook neither receives nor
-        # returns it and the loop is expected to supply it. Exclude those keys from the
-        # identity check rather than demanding the hook produce tensors it cannot see.
-        if isinstance(fed, dict) and self.cache_tensors:
-            fed = {key: value for key, value in fed.items() if key not in self.cache_tensors}
         fed_ids = {id(leaf) for leaf in _ndarray_leaves(fed)}
         hook_ids = {id(value) for value in self.fed_inputs.values()}
         if fed_ids != hook_ids:
-            extra = (' (the declared KV-cache inputs are exempt — those are threaded for you)'
-                     if self.cache_tensors else '')
             self._fail('the model must be fed exactly the tensors the last hook call returned — '
                        'any computation between the hook and the model is invisible to the '
-                       f'platform. Move it into the hook.{extra}')
+                       'platform. Move it into the hook.')
         if _nest_fingerprint(self.fed_inputs) != self.fed_inputs_fingerprint:
             self._fail('the model inputs were mutated in place after the hook returned them — '
                        'the platform feeds the model the tensors exactly as the hook returned '
@@ -2033,24 +1998,12 @@ class _ModelLoopContext:
             self.last_outputs = raw_outputs
             return raw_outputs
         outputs_list = raw_outputs if isinstance(raw_outputs, list) else [raw_outputs]
-        # Cache outputs are plumbing, not predictions: they are fed straight back into the model
-        # and never reported, so they need no prediction type. Only the remainder is keyed.
-        n_cache_outputs = len(self.cache_tensors)
-        if len(outputs_list) - n_cache_outputs != len(self.prediction_names):
-            cache_note = (f' ({n_cache_outputs} of them are the declared KV-cache, which needs '
-                          f'no prediction types)' if n_cache_outputs else '')
+        if len(outputs_list) != len(self.prediction_names):
             self._fail(f'the model returned {len(outputs_list)} outputs but '
                        f'{len(self.prediction_names)} prediction types are declared on '
-                       f'tensorleap_load_model{cache_note} — declare one prediction type per '
-                       f'non-cache model output.')
-        # Predictions come first in graph order, as the export declares them; the cache follows.
-        named_outputs = {name: np.asarray(output) for name, output
-                         in zip(self.prediction_names, outputs_list)}
-        # The hook must receive predictions only, exactly as on the platform, and must receive
-        # this very dict (on_hook_call checks identity). The loop still needs the cache, because
-        # locally it stands in for the platform and threads it by hand — so that goes out on the
-        # proxy rather than being mixed into the hook's view.
-        self.last_cache_outputs = list(outputs_list[len(self.prediction_names):])
+                       f'tensorleap_load_model — declare one prediction type per model output.')
+        named_outputs = {name: np.asarray(output)
+                         for name, output in zip(self.prediction_names, outputs_list)}
         self.last_outputs = named_outputs
         self.last_outputs_fingerprint = _nest_fingerprint(named_outputs)
         return named_outputs
@@ -2072,19 +2025,8 @@ class _ModelLoopModelProxy:
     def get_inputs(self):
         return self._model.get_inputs()
 
-    @property
-    def cache_outputs(self):
-        """Declared-cache outputs of the last model call, in graph order.
-
-        Only the local loop needs these: it stands in for the platform, which threads the cache
-        for you at runtime. They are kept off the hook's `prev_outputs` deliberately — the hook
-        must not see the cache in either place, or the local run would diverge from the real one.
-        """
-        return self._context.last_cache_outputs
 
-
-def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step',
-                                   cache_tensors: Optional[Dict[str, str]] = None):
+def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step'):
     """The feedback hook that drives an autoregressive chain.
 
     Signature of the decorated function:
@@ -2108,22 +2050,6 @@ def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step',
     space comes from the final step's forward pass, except input-kind latent spaces which come
     from the first step — the original sample, before generated content dominates the model
     inputs. 'mean': every latent space is the elementwise mean over all steps of the chain.
-
-    cache_tensors declares a KV-cache (or any threaded state) as
-    {model_input_name: model_output_name} — model inputs whose value is the named output from
-    the previous step. The platform threads those tensors itself and holds them on the GPU, so
-    the hook neither receives them in prev_inputs/prev_outputs nor returns them in next_inputs;
-    return only the inputs you actually compute (e.g. input_ids, position_ids, attention_mask).
-    The cache must be fixed-capacity: the declared input and output must have the same, fully
-    concrete shape, since chains at different steps are batched together and the cache is
-    preallocated. Exports with a growing cache (past length P in, P+1 out) are rejected at push.
-
-        @tensorleap_autoregressive_step(cache_tensors={
-            f'past_key_values.{i}.{kv}': f'present.{i}.{kv}'
-            for i in range(30) for kv in ('key', 'value')
-        })
-
-    Omit it for a model with no cache — the chain then re-feeds its whole prompt every step.
     """
     assert isinstance(latent_space_aggregation, str), \
         ('tensorleap_autoregressive_step must be called with parentheses: '
@@ -2271,8 +2197,7 @@ def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step',
             return result
 
         leap_binder.set_autoregressive_step(inner_without_validate,
-                                            latent_space_aggregation=latent_space_aggregation,
-                                            cache_tensors=cache_tensors)
+                                            latent_space_aggregation=latent_space_aggregation)
 
         def inner(*args, **kwargs):
             if not _call_from_tl_platform:
@@ -2472,9 +2397,8 @@ def tensorleap_model_loop():
             if not prediction_names:
                 _fail('tensorleap_model_loop requires prediction types declared on '
                       'tensorleap_load_model — the loop keys the model outputs by their names.')
-            context = _ModelLoopContext(
-                sample_id, preprocess_response, prediction_names, is_mapping,
-                cache_tensors=leap_binder.setup_container.autoregressive_step.cache_tensors)
+            context = _ModelLoopContext(sample_id, preprocess_response, prediction_names,
+                                        is_mapping)
             _active_model_loop = context
             try:
                 result = user_function(_ModelLoopModelProxy(model, context), sample_id,

{code_loader-1.0.202.dev1 → code_loader-1.0.203.dev0}/code_loader/leaploader.py

@@ -1382,20 +1382,6 @@ class LeapLoader(LeapLoaderBase):
             return 'last_step'
         return handler.latent_space_aggregation
 
-    @lru_cache()
-    def get_autoregressive_cache_tensors(self) -> Dict[str, str]:
-        """Declared KV-cache pairing, {model_input_name: model_output_name}.
-
-        Empty when nothing is declared — the platform then runs the chain uncached. The platform
-        validates these names against the real model graph at push time; only structural checks
-        happen here.
-        """
-        self.exec_script()
-        handler = global_leap_binder.setup_container.autoregressive_step
-        if handler is None or not handler.cache_tensors:
-            return {}
-        return dict(handler.cache_tensors)
-
     def run_autoregressive_step(self, sample_id: Union[int, str],
                                 prev_inputs: Optional[Dict[str, npt.NDArray[np.float32]]],
                                 prev_outputs: Optional[Dict[str, npt.NDArray[np.float32]]],

{code_loader-1.0.202.dev1 → code_loader-1.0.203.dev0}/code_loader/leaploaderbase.py

@@ -176,11 +176,6 @@ class LeapLoaderBase:
         raise NotImplementedError(f'{type(self).__name__} does not implement '
                                   'get_autoregressive_latent_space_aggregation.')
 
-    @abstractmethod
-    def get_autoregressive_cache_tensors(self) -> Dict[str, str]:
-        raise NotImplementedError(f'{type(self).__name__} does not implement '
-                                  'get_autoregressive_cache_tensors.')
-
     @abstractmethod
     def get_autoregressive_decorator_names(self) -> Dict[str, List[str]]:
         raise NotImplementedError(f'{type(self).__name__} does not implement '

{code_loader-1.0.202.dev1 → code_loader-1.0.203.dev0}/pyproject.toml

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