tracebloc 1.2.76__tar.gz → 1.2.77__tar.gz

{tracebloc-1.2.76 → tracebloc-1.2.77}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tracebloc
-Version: 1.2.76
+Version: 1.2.77
 Summary: Build AI with anyone. On data that can't move. SDK for the tracebloc collaborative AI workspace.
 Author-email: tracebloc <lukas@tracebloc.io>
 License-Expression: MIT

{tracebloc-1.2.76 → tracebloc-1.2.77}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "tracebloc"
-version = "1.2.76"
+version = "1.2.77"
 description = "Build AI with anyone. On data that can't move. SDK for the tracebloc collaborative AI workspace."
 license = "MIT"
 license-files = ["license.txt"]

{tracebloc-1.2.76 → tracebloc-1.2.77}/tracebloc/convert/families/hf/_build.py

@@ -16,6 +16,8 @@ from typing import Any, Iterable, Iterator, Mapping
 
 from tracebloc.convert.interfaces import ConversionRefusedError, RefusalCause
 from tracebloc.convert.interfaces import SourcedValue
+from tracebloc.convert.heads import collapse as head_prefixes
+from tracebloc.convert.heads import head_key
 from tracebloc.utils.constants import TaskCategory
 
 
@@ -55,8 +57,7 @@ def split_head(
     params = {name for name, _ in model.named_parameters(remove_duplicate=False)}
 
     def is_head(key: str) -> bool:
-        own = key.startswith(f"{prefix}.")
-        return inner and (not own or key.startswith(f"{prefix}.pooler."))
+        return inner and head_key(key, prefix)
 
     keys = list(keys)
     head = tuple(sorted(key for key in keys if is_head(key)))
@@ -180,23 +181,6 @@ def load(
     return model.eval(), fresh, tuple(dropped)
 
 
-def head_prefixes(fresh: Iterable[str], keys: Iterable[str]) -> tuple[str, ...]:
-    """The shortest module prefixes whose every key in ``keys`` is ``fresh``:
-    ``classifier.`` for a head whose layers all start fresh."""
-    fresh, keys = set(fresh), list(keys)
-    prefixes = set()
-    for key in fresh:
-        parts = key.split(".")
-        for depth in range(1, len(parts)):
-            prefix = ".".join(parts[:depth]) + "."
-            if all(other in fresh for other in keys if other.startswith(prefix)):
-                prefixes.add(prefix)
-                break
-        else:
-            prefixes.add(key)
-    return tuple(sorted(prefixes))
-
-
 def write_weights(model: Any, wrapper: str | None, target: Path) -> tuple[str, ...]:
     """Write the model's state as SafeTensors; return the keys written.
 

tracebloc-1.2.77/tracebloc/convert/heads.py

@@ -0,0 +1,51 @@
+"""Which keys of a converted model are its task's head, which starts fresh.
+
+Stdlib and ``tracebloc`` core at import.
+"""
+
+from __future__ import annotations
+
+from typing import Iterable
+
+from tracebloc.convert.interfaces import GeneratedPackage
+
+#: A package value naming its backbone's module: the keys the shipped weights
+#: lack outside it, or in its pooler, start fresh (:func:`head_key`).
+BACKBONE_PREFIX = "backbone_prefix"
+
+
+def head_key(key: str, backbone: str) -> bool:
+    """Whether ``key`` is the task's head of a model whose backbone is ``backbone``:
+    outside it, or in its ``pooler``, which only a head reads (Hugging Face's
+    ``base_model_prefix``)."""
+    return not key.startswith(f"{backbone}.") or key.startswith(f"{backbone}.pooler.")
+
+
+def fresh_head(package: GeneratedPackage, keys: Iterable[str]) -> set[str]:
+    """The keys of a built package, ``keys``, its weights lack in its head; empty
+    for a package that declares no :data:`BACKBONE_PREFIX` or ships no weights."""
+    declared = package.values.get(BACKBONE_PREFIX)
+    if declared is None or package.weights is None:
+        return set()
+    from safetensors import safe_open
+
+    with safe_open(str(package.weights), framework="pt") as weights:
+        shipped = set(weights.keys())
+    return {k for k in keys if k not in shipped and head_key(k, declared.value)}
+
+
+def collapse(fresh: Iterable[str], keys: Iterable[str]) -> tuple[str, ...]:
+    """The shortest module prefixes whose every key in ``keys`` is ``fresh``:
+    ``classifier.`` for a head whose layers all start fresh."""
+    fresh, keys = set(fresh), list(keys)
+    prefixes = set()
+    for key in fresh:
+        parts = key.split(".")
+        for depth in range(1, len(parts)):
+            prefix = ".".join(parts[:depth]) + "."
+            if all(other in fresh for other in keys if other.startswith(prefix)):
+                prefixes.add(prefix)
+                break
+        else:
+            prefixes.add(key)
+    return tuple(sorted(prefixes))

{tracebloc-1.2.76 → tracebloc-1.2.77}/tracebloc/convert/outputs.py

@@ -92,6 +92,18 @@ def _one_line(error: BaseException) -> str:
 
 # --- the output comparison ---------------------------------------------------
 
+#: The function a generated entry file passes its class's output through to
+#: return the task's output (``owncode.generate``); the original's goes through it too.
+OUTPUT_ADAPTER = "_task_output"
+
+
+def task_output(model: Any) -> Callable[[Any], Any]:
+    """What the generated ``MyModel.forward`` keeps of its class's output, read
+    from the entry file defining it; the output itself when it adapts none."""
+    forward = vars(type(model)).get("forward")
+    adapter = getattr(forward, "__globals__", {}).get(OUTPUT_ADAPTER)
+    return adapter if callable(adapter) else lambda output: output
+
 
 def original_builder(family: Family | None) -> Callable[..., Any] | None:
     """The family's ``build_original``, or its module's (the own-code converter's)."""
@@ -322,7 +334,8 @@ def torch_output_check(
     """The generated model gives the original's output on a fixed probe.
 
     Exact for own code; within the category's tolerance for a template-rebuilt
-    family or bf16 weights. With a fresh head, the backbone's output is
+    family or bf16 weights. The original's output is kept as the generated
+    model keeps its class's (:func:`task_output`). With a fresh head, the backbone's output is
     compared: the original's backbone against the submodule holding its
     parameters, or what reaches the head in each.
     """
@@ -338,9 +351,10 @@ def torch_output_check(
         probe = probe_for(category, header, model)
         with torch.no_grad():
             if package.compared_output is ComparedOutput.MODEL:
+                adapt = task_output(model)
                 ours, theirs = (
                     output_tensors(probe.call(model)),
-                    output_tensors(probe.call(original)),
+                    output_tensors(adapt(probe.call(original))),
                 )
                 compared = "the model's output"
             else:

{tracebloc-1.2.76 → tracebloc-1.2.77}/tracebloc/convert/owncode/__init__.py

@@ -142,7 +142,7 @@ def build_original(recognition: Recognition, package: GeneratedPackage) -> Any:
             auto = next(
                 k for k, (_, name) in state["auto_map"].items() if name == entry
             )
-        loader = vars(transformers)[auto]
+        loader = getattr(transformers, auto)  # a lazy module: not in its vars yet
         return loader.from_pretrained(
             root, trust_remote_code=True, local_files_only=True
         )

{tracebloc-1.2.76 → tracebloc-1.2.77}/tracebloc/convert/owncode/generate.py

@@ -19,6 +19,7 @@ from tracebloc.convert.interfaces import (
     UseCaseTask,
     ValueSource,
 )
+from tracebloc.convert.outputs import OUTPUT_ADAPTER
 from tracebloc.utils.constants import Framework, TaskCategory
 from tracebloc.validation.header_rules import (
     NO_OUTPUT_CLASSES,
@@ -54,7 +55,7 @@ DEFAULTS = {
 
 _ADAPTER = """
 
-def _task_output(output):
+def {adapter}(output):
     if isinstance(output, dict):
         for key in ("logits", "out"):
             if key in output:
@@ -149,7 +150,7 @@ def entry_source(
     if preamble:
         text += f"\n{preamble}\n"
     if adapter:
-        text += _ADAPTER.replace("{pick}", pick)
+        text += _ADAPTER.replace("{pick}", pick).replace("{adapter}", OUTPUT_ADAPTER)
     call = ""
     if resize is not None:
         text += _RESIZE
@@ -161,7 +162,7 @@ def entry_source(
             "    def __init__(self):\n"
             f"        super().__init__({construct}){resize_line}\n\n"
             "    def forward(self, *args, **kwargs):\n"
-            "        return _task_output(super().forward(*args, **kwargs))\n"
+            f"        return {OUTPUT_ADAPTER}(super().forward(*args, **kwargs))\n"
         )
     elif call:
         text += (

{tracebloc-1.2.76 → tracebloc-1.2.77}/tracebloc/convert/owncode/live.py

@@ -20,7 +20,7 @@ import symtable
 import sys
 import textwrap
 from pathlib import Path
-from typing import Any, Mapping
+from typing import Any, Callable, Mapping
 
 from tracebloc.convert.interfaces import (
     ConversionInput,
@@ -262,33 +262,33 @@ def _package_root(path: Path) -> Path:
     return root
 
 
-def recognise(source: ConversionInput) -> Recognition | None:
-    torch = sys.modules.get("torch")
-    obj = source.live_object
-    if torch is None or not isinstance(obj, torch.nn.Module):
-        return None
+def _reading(call: Callable[[], Any]) -> Any:
+    """``call``'s result: it runs the peer's object, so its error is a refusal.
+    Running out of memory here is this machine's, not the code's."""
     try:
-        model, hyper, notes = _model_of(obj)
-        weights = dict(model.state_dict())
-        remote_config = (
-            model.config.to_dict()
-            if type(model).__module__.startswith("transformers_modules.")
-            else None
-        )
-    except ConversionRefusedError:
+        return call()
+    except (ConversionRefusedError, MemoryError):
         raise
     except Exception as error:
         what = "reading the object passed to upload_model"
         raise refuse_code_failure(what, error) from error
+
+
+def recognise(source: ConversionInput) -> Recognition | None:
+    torch = sys.modules.get("torch")
+    obj = source.live_object
+    if torch is None or not isinstance(obj, torch.nn.Module):
+        return None
+    model, hyper, notes = _reading(lambda: _model_of(obj))
     cls = type(model)
     module = sys.modules.get(cls.__module__)
     file = getattr(module, "__file__", None)
     if file and installed(Path(file)):
         raise _refuse(f"is {cls.__module__}.{cls.__name__}, a library class")
-    if remote_config is not None and file:
-        found = remote.recognise_parts(
-            source, Path(file).parent, remote_config, weights, None
-        )
+    weights = _reading(lambda: dict(model.state_dict()))
+    if cls.__module__.startswith("transformers_modules.") and file:
+        config = _reading(lambda: model.config.to_dict())
+        found = remote.recognise_parts(source, Path(file).parent, config, weights, None)
         state = {**found.state, "live_object": model, "proof": Proof.EXACT}
         return dataclasses.replace(found, state=state)
     if cls.__module__.split(".")[0] in LIBRARY_ROOTS:

{tracebloc-1.2.76 → tracebloc-1.2.77}/tracebloc/convert/owncode/remote.py

@@ -32,6 +32,7 @@ from tracebloc.convert.owncode import assemble
 from tracebloc.convert.owncode.checkpoint import load, refuse_no_source
 from tracebloc.convert.owncode.generate import has_output_classes
 from tracebloc.convert.owncode.shake import flat_name
+from tracebloc.convert.heads import BACKBONE_PREFIX
 from tracebloc.convert.owncode.source_tree import SourceTree
 from tracebloc.utils.constants import Framework, TaskCategory
 
@@ -262,10 +263,13 @@ def build(
         "output_classes"
     )
     weights = state["weights"]
+    prefix = _base_prefix(state["tree"], module)
     if weights is not None:
-        weights = _move_weights(
-            weights, _base_prefix(state["tree"], module), bare, same and unchanged
-        )
+        weights = _move_weights(weights, prefix, bare, same and unchanged)
+    values = dict(recognition.values)
+    if prefix and not bare:
+        where = Provenance(ValueSource.REPO_FILE, f"{module}.py base_model_prefix")
+        values[BACKBONE_PREFIX] = SourcedValue(prefix, where)
     found = {}
     if isinstance(config.get("vocab_size"), int):
         found["vocab_size"] = SourcedValue(
@@ -295,4 +299,5 @@ def build(
             ComparedOutput.MODEL if same and unchanged else ComparedOutput.BACKBONE
         ),
     }
-    return assemble.build(dataclasses.replace(recognition, state=built), task, workdir)
+    built_from = dataclasses.replace(recognition, state=built, values=values)
+    return assemble.build(built_from, task, workdir)

{tracebloc-1.2.76 → tracebloc-1.2.77}/tracebloc/convert/package.py

@@ -37,6 +37,7 @@ from tracebloc.convert.interfaces import (
     RefusalCause,
     SourcedValue,
 )
+from tracebloc.convert.heads import collapse, fresh_head
 from tracebloc.convert.training_versions import refuse_code_failure
 from tracebloc.exceptions import ModelValidationError
 from tracebloc.validation.rewriter import ModelFileRewriter
@@ -384,9 +385,10 @@ def derive_seed_excluded_prefixes(
 
     ``build_for_classes(count, workdir)`` builds the package for ``count``
     output classes, which for a family is its ``build`` with that count. The
-    peer's own code raising when it runs here is its :func:`code_failure`.
+    peer's own code raising when it runs here is its :func:`code_failure`. The
+    head is also every key :func:`fresh_head` finds in the first build.
     """
-    built = []
+    built, fresh = [], set[str]()
     for count in PROBE_CLASSES:
         with tempfile.TemporaryDirectory(prefix="tracebloc_derive_") as raw:
             workdir, out = Path(raw) / "build", Path(raw) / "package"
@@ -401,7 +403,14 @@ def derive_seed_excluded_prefixes(
                 if package.family != OWN_CODE_FAMILY:
                     raise
                 raise code_failure(written, error) from error
-    return head_prefixes(*built)
+            if len(built) == 1:
+                fresh = fresh_head(package, built[0])
+    varying = head_prefixes(*built)
+    kept = [key for key in built[0] if not key.startswith(varying)]
+    found = {*varying, *collapse((k for k in fresh if k in kept), kept)}
+    return tuple(
+        sorted(p for p in found if not any(o != p and p.startswith(o) for o in found))
+    )
 
 
 def cast_weights_to_bf16(path: str | os.PathLike[str]) -> int:

{tracebloc-1.2.76 → tracebloc-1.2.77}/tracebloc.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tracebloc
-Version: 1.2.76
+Version: 1.2.77
 Summary: Build AI with anyone. On data that can't move. SDK for the tracebloc collaborative AI workspace.
 Author-email: tracebloc <lukas@tracebloc.io>
 License-Expression: MIT

{tracebloc-1.2.76 → tracebloc-1.2.77}/tracebloc.egg-info/SOURCES.txt

@@ -42,6 +42,7 @@ tracebloc/contracts/training_requirements.txt
 tracebloc/convert/__init__.py
 tracebloc/convert/acceptance.py
 tracebloc/convert/gate.py
+tracebloc/convert/heads.py
 tracebloc/convert/interfaces.py
 tracebloc/convert/outputs.py
 tracebloc/convert/package.py