orivane 0.1.0__py3-none-any.whl

orivane/__init__.py

@@ -0,0 +1,8 @@
+from .patch import (
+    patch,
+    inject,
+    PatchType,
+    InvalidModuleException,
+    InvalidPatchMethodException,
+    PatchBase,
+)

orivane/patch/__init__.py

@@ -0,0 +1,5 @@
+from .patch import patch
+from .inject import inject
+from .patch_base import PatchBase
+from .patch_type import PatchType
+from .patch_errors import InvalidModuleException, InvalidPatchMethodException

orivane/patch/inject.py

@@ -0,0 +1,25 @@
+from functools import wraps
+
+from .patch_errors import InvalidPatchMethodException
+from .patch_type import PatchType
+
+def inject(method: str, at: PatchType):
+    def decorator(func):    
+        @wraps(func)
+        def wrapper(self, *args, **kwargs):
+            if not getattr(self, "__orivane_patched__", False):
+                raise InvalidPatchMethodException(f"{func.__name__} can only be used inside a class decorated with @patch.")
+
+            if getattr(self, "__orivane_model__", None) is None:
+                raise InvalidPatchMethodException(f"{type(self).__name__} doesn't contain an __orivane_model__.")
+
+            return func(self, *args, **kwargs)
+
+        wrapper.__orivane_injection__ = {
+            "method": method,
+            "at": at,
+        }
+
+        return wrapper
+
+    return decorator

orivane/patch/patch.py

@@ -0,0 +1,201 @@
+from functools import wraps
+
+from torch import nn
+
+from .patch_base import PatchBase
+from .patch_errors import InvalidModuleException, InvalidPatchMethodException
+from .patch_type import PatchType
+
+def _resolve_method(model: nn.Module, path: str):
+    module_path, separator, method_name = path.rpartition(".")
+    target = model.get_submodule(module_path) if separator else model
+    method = getattr(target, method_name, None)
+
+    if not callable(method):
+        raise InvalidPatchMethodException(
+            f"Method {path!r} does not exist on {type(model).__name__}."
+        )
+
+    return target, method_name, method
+
+
+def _create_wrapper(original, injections):
+    before = [
+        callback for at, callback in injections
+        if at is PatchType.BEFORE
+    ]
+    after = [
+        callback for at, callback in injections
+        if at is PatchType.AFTER
+    ]
+    overwrite = [
+        callback for at, callback in injections
+        if at is PatchType.OVERWRITE
+    ]
+
+    if len(overwrite) > 1:
+        raise InvalidPatchMethodException(
+            "Only one OVERWRITE injection is allowed per method."
+        )
+
+    implementation = overwrite[0] if overwrite else original
+
+    @wraps(original)
+    def wrapper(*args, **kwargs):
+        for callback in before:
+            result = callback(*args, **kwargs)
+
+            if result is not None:
+                if (
+                    not isinstance(result, tuple)
+                    or len(result) != 2
+                    or not isinstance(result[0], tuple)
+                    or not isinstance(result[1], dict)
+                ):
+                    raise TypeError(
+                        "BEFORE injections must return None or "
+                        "(args: tuple, kwargs: dict)."
+                    )
+
+                args, kwargs = result
+
+        output = implementation(*args, **kwargs)
+
+        for callback in after:
+            output = callback(output)
+
+        return output
+
+    return wrapper
+
+
+def _next_patch_name(model: nn.Module) -> str:
+    index = 0
+    while f"_orivane_patch_{index}" in model._modules:
+        index += 1
+    return f"_orivane_patch_{index}"
+
+
+def patch(*, model: nn.Module, train_only_added: bool = True):
+    if not isinstance(model, nn.Module):
+        raise InvalidModuleException("'model' must inherit from torch.nn.Module.")
+
+    def decorator(cls):
+        if not issubclass(cls, PatchBase):
+            raise InvalidPatchMethodException(
+                f"{cls.__name__} must inherit from PatchBase."
+            )
+
+        cls.__orivane_patch_decorator_applied__ = True
+        instance = cls()
+        instance._set_model(model)
+
+        reference = next(model.parameters(), None)
+        if reference is not None:
+            instance.to(device=reference.device, dtype=reference.dtype)
+
+        injections = {}
+        for name, member in vars(cls).items():
+            metadata = getattr(member, "__orivane_injection__", None)
+            if metadata is None:
+                continue
+
+            method_path = metadata["method"]
+            at = metadata["at"]
+            if not isinstance(method_path, str) or not method_path:
+                raise InvalidPatchMethodException(
+                    "Injected method paths must be non-empty strings."
+                )
+            if not isinstance(at, PatchType):
+                raise InvalidPatchMethodException(
+                    f"Invalid patch type for method {method_path!r}."
+                )
+
+            injections.setdefault(method_path, []).append(
+                (at, getattr(instance, name))
+            )
+
+        restorations = []
+        for method_path, callbacks in injections.items():
+            target, method_name, original = _resolve_method(model, method_path)
+            had_instance_attr = method_name in target.__dict__
+            old_instance_attr = target.__dict__.get(method_name)
+            restorations.append(
+                (
+                    target,
+                    method_name,
+                    had_instance_attr,
+                    old_instance_attr,
+                    _create_wrapper(original, callbacks),
+                )
+            )
+
+        patch_name = _next_patch_name(model)
+        instance.__orivane_patched__ = True
+        original_trainability = []
+        applied_restorations = []
+        registered = False
+
+        try:
+            model.add_module(patch_name, instance)
+            registered = True
+
+            if train_only_added:
+                patch_parameter_ids = {
+                    id(parameter) for parameter in instance.parameters()
+                }
+                for parameter in model.parameters():
+                    if id(parameter) not in patch_parameter_ids:
+                        original_trainability.append(
+                            (parameter, parameter.requires_grad)
+                        )
+                        parameter.requires_grad_(False)
+
+            for target, method_name, existed, old_value, wrapped in restorations:
+                setattr(target, method_name, wrapped)
+                applied_restorations.append(
+                    (target, method_name, existed, old_value)
+                )
+        except Exception:
+            for target, method_name, existed, old_value in reversed(
+                applied_restorations
+            ):
+                if existed:
+                    setattr(target, method_name, old_value)
+                else:
+                    delattr(target, method_name)
+            for parameter, requires_grad in original_trainability:
+                parameter.requires_grad_(requires_grad)
+            if registered:
+                delattr(model, patch_name)
+            instance.__orivane_patched__ = False
+            raise
+
+        active = True
+
+        def unpatch():
+            nonlocal active
+            if not active:
+                return
+
+            for target, method_name, existed, old_value in reversed(
+                applied_restorations
+            ):
+                if existed:
+                    setattr(target, method_name, old_value)
+                else:
+                    delattr(target, method_name)
+
+            for parameter, requires_grad in original_trainability:
+                parameter.requires_grad_(requires_grad)
+
+            if model._modules.get(patch_name) is instance:
+                delattr(model, patch_name)
+
+            instance.__orivane_patched__ = False
+            active = False
+
+        instance.unpatch = unpatch
+        return instance
+
+    return decorator

orivane/patch/patch_base.py

@@ -0,0 +1,70 @@
+from torch import nn
+
+from .patch_errors import InvalidModuleException, InvalidPatchMethodException
+
+class _PatchMeta(type(nn.Module)):
+    def __call__(cls, *args, **kwargs):
+        has_injections = any(
+            getattr(member, "__orivane_injection__", None) is not None
+            for member in vars(cls).values()
+        )
+        decorator_applied = cls.__dict__.get(
+            "__orivane_patch_decorator_applied__", False
+        )
+
+        if has_injections and not decorator_applied:
+            raise InvalidPatchMethodException(
+                f"{cls.__name__} contains injected methods and must be "
+                "decorated with @patch."
+            )
+
+        return super().__call__(*args, **kwargs)
+
+
+class PatchBase(nn.Module, metaclass=_PatchMeta):
+    def __init__(self):
+        super().__init__()
+        object.__setattr__(self, "__orivane_model__", None)
+        self.__orivane_added_modules__ = {}
+        self.__orivane_patched__ = False
+
+    def _set_model(self, model: nn.Module):
+        object.__setattr__(self, "__orivane_model__", model)
+
+    def add_module(
+        self,
+        name: str,
+        module: nn.Module,
+        *,
+        trainable: bool = True,
+        device=None,
+    ):
+        model = self.__orivane_model__
+
+        if not isinstance(module, nn.Module):
+            raise InvalidModuleException(
+                "'module' has to inherit from torch.nn.Module."
+            )
+
+        if name in self._modules or hasattr(self, name):
+            raise ValueError(f"Patch attribute {name!r} already exists.")
+
+        if device is None and model is not None:
+            reference = next(model.parameters(), None)
+            if reference is not None:
+                device = reference.device
+
+        if device is not None:
+            module.to(device)
+
+        module.requires_grad_(trainable)
+        super().add_module(name, module)
+        self.__orivane_added_modules__[name] = module
+        return module
+
+    def trainable_parameters(self):
+        return (
+            parameter
+            for parameter in self.parameters()
+            if parameter.requires_grad
+        )

orivane/patch/patch_errors.py

@@ -0,0 +1,7 @@
+
+
+class InvalidPatchMethodException(Exception):
+    pass
+
+class InvalidModuleException(Exception):
+    pass

orivane/patch/patch_type.py

@@ -0,0 +1,6 @@
+from enum import Enum
+
+class PatchType(Enum):
+    BEFORE = 0 # Changes the input for the original function, it's output is the new input for the orignal function
+    AFTER = 1 # Changes the output, it's input is the output of the original function
+    OVERWRITE = 2 # Overwrites the whole function

orivane-0.1.0.dist-info/METADATA

@@ -0,0 +1,105 @@
+Metadata-Version: 2.3
+Name: orivane
+Version: 0.1.0
+Summary: Orivane is a python package that makes modifing architecture of an already pretrained model easier.
+Author: orven
+License: MIT
+Requires-Dist: torch>=2.14.1
+Requires-Python: >=3.14
+Description-Content-Type: text/markdown
+
+# Orivane
+
+Orivane is a small Python library in alpha for making experimental changes to PyTorch models. It wraps methods with input, output, or replacement callbacks and registers patch modules with the model so they participate in device movement, training mode, checkpoints, and optimizers.
+
+## Installation
+
+```bash
+uv add orivane
+```
+
+## Add a trainable adapter
+
+Patch classes inherit from `PatchBase` and call `super().__init__()` so PyTorch can register modules assigned to the patch. By default, applying a patch freezes the model's existing parameters while leaving the patch's trainable parameters trainable.
+
+```python
+import torch
+import torch.nn as nn
+
+from orivane import PatchBase, PatchType, inject, patch
+
+model = nn.Linear(10, 5)
+
+@patch(model=model)
+class AdapterPatch(PatchBase):
+    def __init__(self):
+        super().__init__()
+        self.adapter = nn.Sequential(
+            nn.Linear(5, 16),
+            nn.GELU(),
+            nn.Linear(16, 5),
+        )
+
+    @inject(method="forward", at=PatchType.AFTER)
+    def add_adapter(self, output):
+        return output + self.adapter(output)
+
+optimizer = torch.optim.AdamW(
+    (parameter for parameter in model.parameters() if parameter.requires_grad),
+    lr=1e-3,
+)
+
+x = torch.randn(32, 10)
+target = torch.randn(32, 5)
+loss = nn.MSELoss()(model(x), target)
+optimizer.zero_grad()
+loss.backward()
+optimizer.step()
+
+# Restore the original method and parameter trainability when finished.
+AdapterPatch.unpatch()
+```
+
+The decorated class name refers to the applied patch instance. Its added modules are registered under the model, so they are included in `model.parameters()` and `model.state_dict()`, and follow `model.to(...)` and `model.train()` calls. Patch modules are moved to the device and dtype of the model's first parameter when the patch is applied.
+
+Use `train_only_added=False` to leave the model's existing trainability settings unchanged:
+
+```python
+@patch(model=model, train_only_added=False)
+class FineTunePatch(PatchBase):
+    def __init__(self):
+        super().__init__()
+```
+
+Unpatching is safe to call more than once. It restores the wrapped method and the original `requires_grad` state of model parameters, and removes the patch module from the model.
+
+## Target a nested module
+
+`inject(method=...)` accepts a method name on the model or a dotted path to a nested module method. For example, `"layers.0.forward"` targets the first module in a `ModuleList` named `layers`:
+
+```python
+class BlockPatch(PatchBase):
+    def __init__(self):
+        super().__init__()
+        self.scale = nn.Parameter(torch.tensor(0.1))
+
+    @inject(method="layers.0.forward", at=PatchType.AFTER)
+    def scale_output(self, output):
+        return output * (1 + self.scale)
+```
+
+Apply it with `@patch(model=model)` as in the adapter example. `BEFORE` callbacks receive the method's positional and keyword arguments and may return `(new_args, new_kwargs)`; `AFTER` callbacks receive and return the method result; `OVERWRITE` callbacks replace the method implementation.
+
+## Patch types
+
+| Patch type | Behavior |
+|---|---|
+| `PatchType.BEFORE` | Optionally changes the arguments before the original method runs. |
+| `PatchType.AFTER` | Changes the original method's result. |
+| `PatchType.OVERWRITE` | Replaces the original method; only one overwrite is allowed per method. |
+
+Patches operate on PyTorch modules and methods. They do not provide model-specific integrations, automatic layer discovery, distributed-training support, or quantization support.
+
+## License
+
+See the [LICENSE](LICENSE) file for license information.

orivane-0.1.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+orivane/__init__.py,sha256=NLnDfMzpR5u3-ipdtJD3QvoXhUmDCa6bp_IBlEp0cDk,136
+orivane/patch/__init__.py,sha256=EvbMGquBUEA5RR1dRsH7EZlu3_pelMYJtbmer46HxdE,197
+orivane/patch/inject.py,sha256=sZGQ6rquwqFyDOaKXBGVRvQ586mkqFIwNWr4IrWCD_U,819
+orivane/patch/patch.py,sha256=tpiW8r5oVBgqNEuVDA-HVMi2rMxCEN7vWx5pJZEXZ-A,6563
+orivane/patch/patch_base.py,sha256=62gwIUw_mCfpv467EecxlYTI4YVG1T-xbu5jQiUxaSo,2146
+orivane/patch/patch_errors.py,sha256=SLg2OOhPR4_yzZr8jc6GlZiXuHPR9K43_sK1H5Q24qM,107
+orivane/patch/patch_type.py,sha256=RJsI4GP-6GFM2wEwKw6rYerzxOifJG5bWEPvhQw9_fQ,297
+orivane-0.1.0.dist-info/WHEEL,sha256=oy45bIW3ehK02fnLDUKD-FzeaR3VpJxY5Jdblkktnpw,80
+orivane-0.1.0.dist-info/METADATA,sha256=5C_aowUJozQR1ZXvlRjSbsrNG0R-IaVKCmec5jH1UOc,3879
+orivane-0.1.0.dist-info/RECORD,,

orivane-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: uv 0.12.0
+Root-Is-Purelib: true
+Tag: py3-none-any