code-loader 1.0.201__tar.gz → 1.0.202.dev0__tar.gz

{code_loader-1.0.201 → code_loader-1.0.202.dev0}/PKG-INFO

@@ -1,9 +1,9 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: code-loader
-Version: 1.0.201
+Version: 1.0.202.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,6 +20,7 @@ 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.201 → code_loader-1.0.202.dev0}/code_loader/contract/datasetclasses.py

@@ -332,6 +332,11 @@ 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.201 → code_loader-1.0.202.dev0}/code_loader/contract/visualizer_classes.py

@@ -353,7 +353,7 @@ class LeapAudio:
     waveform (a LeapGraph). Build one with `default_audio_visualizer`.
 
     Attributes:
-    audio (npt.NDArray): Mono waveform samples, shaped [N]. Multi-channel input
+    audio (npt.NDArray[np.float32]): Mono waveform samples, shaped [N]. Multi-channel input
         should be mixed down before constructing LeapAudio.
     sample_rate (int): Samples per second — needed to encode the clip and to map
         the playhead position to time.
@@ -362,7 +362,7 @@ class LeapAudio:
         driven by the audio element's currentTime, so no time range is needed here.
     type (LeapDataType): The data type, default is LeapDataType.Audio.
     """
-    audio: npt.NDArray
+    audio: npt.NDArray[np.float32]
     sample_rate: int
     visual: Union[LeapImage, LeapGraph]
     type: LeapDataType = LeapDataType.Audio

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

@@ -534,6 +534,11 @@ class LeapBinder:
             leap_binder.set_metadata(metadata_handler_index, name='metadata_index')
             leap_binder.set_metadata(metadata_handler_image_mean, name='metadata_image_mean')
         """
+        if not isinstance(name, str):
+            raise TypeError(f"set_metadata validation failed.\n`name` must be a string, got type "
+                            f"{type(name).__name__}. When using @tensorleap_metadata, the metadata "
+                            f"types go in the second argument: "
+                            f"@tensorleap_metadata('my_metadata', {{'key': DatasetMetadataType.int}}).")
         self.setup_container.metadata.append(MetadataHandler(name, function, metadata_type))
 
     def set_custom_latent_space(self, function: SectionCallableInterface,
@@ -585,8 +590,41 @@ 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') -> None:
+                                latent_space_aggregation: str = 'last_step',
+                                cache_tensors: Optional[Dict[str, str]] = None) -> 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)
@@ -594,6 +632,8 @@ 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. '
@@ -602,8 +642,10 @@ 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)
+            function, latent_space_aggregation=latent_space_aggregation,
+            cache_tensors=cache_tensors)
 
         # 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

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

@@ -560,12 +560,22 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
                     return result
 
                 def validate_declared_prediction_types(self, ret):
-                    if not (len(self.prediction_types) == len(ret) if isinstance(ret, list) else 1) and len(
-                            self.prediction_types) != 0:
+                    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 _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 prediction types({len(prediction_types)}) != number of model outputs({len(ret) if isinstance(ret, list) else 1})")
+                            f"tensorleap_load_model validation failed: number of declared "
+                            f"prediction types({len(prediction_types)}) != number of model "
+                            f"outputs({n_outputs}){cache_note}")
 
                 def _convert_onnx_inputs_to_correct_type(
                         self, float_arrays_inputs: Dict[str, np.ndarray]
@@ -1883,12 +1893,29 @@ class _MappingStatePlaceholder:
         self._unsupported()
 
 
+def _shares_base(leaf, original):
+    """True when `leaf` is `original` or a view onto it (e.g. original[None])."""
+    if leaf is original:
+        return True
+    base = getattr(leaf, 'base', None)
+    while base is not None:
+        if base is original:
+            return True
+        base = getattr(base, 'base', None)
+    return False
+
+
 class _ModelLoopContext:
-    def __init__(self, sample_id, preprocess_response, prediction_names, is_mapping):
+    def __init__(self, sample_id, preprocess_response, prediction_names, is_mapping,
+                 cache_tensors=None):
         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
@@ -1897,6 +1924,9 @@ 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}')
@@ -1981,12 +2011,39 @@ class _ModelLoopContext:
                        'next_inputs the loop must exit without further hook or model calls.')
         if self.is_mapping:
             return
-        fed_ids = {id(leaf) for leaf in _ndarray_leaves(fed)}
+        # 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}
         hook_ids = {id(value) for value in self.fed_inputs.values()}
-        if fed_ids != hook_ids:
+        # A leaf may be the hook's array itself, or a VIEW of it: `arr[None]` adds the batch
+        # axis an ONNX model needs and shares its memory, so `.base` is the original. Allowing
+        # views keeps the guarantee that matters — no tensor may be COMPUTED between the hook
+        # and the model, since the platform would not see that work — while not forcing every
+        # integration to choose between the identity check and a model that expects a batch
+        # dimension. A genuinely new array has no base among the hook's tensors.
+        def _from_hook(leaf):
+            if id(leaf) in hook_ids:
+                return True
+            base = getattr(leaf, 'base', None)
+            while base is not None:
+                if id(base) in hook_ids:
+                    return True
+                base = getattr(base, 'base', None)
+            return False
+
+        fed_leaves = list(_ndarray_leaves(fed))
+        matched = {id(value) for value in self.fed_inputs.values()
+                   if any(_from_hook(leaf) and (id(leaf) == id(value)
+                                                or _shares_base(leaf, value))
+                          for leaf in fed_leaves)}
+        if not all(_from_hook(leaf) for leaf in fed_leaves) or matched != 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 '
-                       'platform. Move it into the hook.')
+                       f'platform. Move it into the hook.{extra}')
         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 '
@@ -1998,12 +2055,24 @@ class _ModelLoopContext:
             self.last_outputs = raw_outputs
             return raw_outputs
         outputs_list = raw_outputs if isinstance(raw_outputs, list) else [raw_outputs]
-        if len(outputs_list) != len(self.prediction_names):
+        # 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 '')
             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 — declare one prediction type per model output.')
-        named_outputs = {name: np.asarray(output)
-                         for name, output in zip(self.prediction_names, outputs_list)}
+                       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):])
         self.last_outputs = named_outputs
         self.last_outputs_fingerprint = _nest_fingerprint(named_outputs)
         return named_outputs
@@ -2025,8 +2094,19 @@ 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'):
+
+def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step',
+                                   cache_tensors: Optional[Dict[str, str]] = None):
     """The feedback hook that drives an autoregressive chain.
 
     Signature of the decorated function:
@@ -2050,6 +2130,22 @@ 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: '
@@ -2197,7 +2293,8 @@ 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)
+                                            latent_space_aggregation=latent_space_aggregation,
+                                            cache_tensors=cache_tensors)
 
         def inner(*args, **kwargs):
             if not _call_from_tl_platform:
@@ -2397,8 +2494,9 @@ 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)
+            context = _ModelLoopContext(
+                sample_id, preprocess_response, prediction_names, is_mapping,
+                cache_tensors=leap_binder.setup_container.autoregressive_step.cache_tensors)
             _active_model_loop = context
             try:
                 result = user_function(_ModelLoopModelProxy(model, context), sample_id,

{code_loader-1.0.201 → code_loader-1.0.202.dev0}/code_loader/leaploader.py

@@ -1382,6 +1382,20 @@ 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.201 → code_loader-1.0.202.dev0}/code_loader/leaploaderbase.py

@@ -176,6 +176,11 @@ 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.201 → code_loader-1.0.202.dev0}/code_loader/visualizers/default_visualizers.py

@@ -33,7 +33,7 @@ def default_video_visualizer(data: npt.NDArray[np.float32]) -> LeapVideo:
 _MAX_WAVEFORM_POINTS = 2000
 
 
-def _to_mono(audio: npt.NDArray) -> npt.NDArray:
+def _to_mono(audio: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
     a = np.squeeze(np.asarray(audio))
     if a.ndim == 2:  # mix down: channels are the shorter axis
         a = a.mean(axis=int(np.argmin(a.shape)))
@@ -42,7 +42,7 @@ def _to_mono(audio: npt.NDArray) -> npt.NDArray:
     return a.astype(np.float32)
 
 
-def _waveform_graph(series: npt.NDArray, duration: float) -> LeapGraph:
+def _waveform_graph(series: npt.NDArray[np.float32], duration: float) -> LeapGraph:
     s = np.squeeze(np.asarray(series)).astype(np.float32)
     # ponytail: recharts chokes on 100k+ points; stride down for the default view. x_range keeps full time span.
     step = max(1, s.shape[0] // _MAX_WAVEFORM_POINTS)
@@ -50,9 +50,9 @@ def _waveform_graph(series: npt.NDArray, duration: float) -> LeapGraph:
     return LeapGraph(s, x_label='Time [s]', y_label='Amplitude', x_range=(0.0, duration))
 
 
-def default_audio_visualizer(audio: npt.NDArray,
+def default_audio_visualizer(audio: npt.NDArray[np.float32],
                              sample_rate: int,
-                             visual: Optional[npt.NDArray] = None,
+                             visual: Optional[npt.NDArray[np.float32]] = None,
                              time_axis_index: int = -1) -> LeapAudio:
     """Package a playable audio sample plus the visual a playhead runs through.
 

{code_loader-1.0.201 → code_loader-1.0.202.dev0}/pyproject.toml

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