enable Torch mixin with init args

[Borda]

Before submitting

• Was this discussed/agreed via a Github issue? (no need for typos and docs improvements)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure to update the docs?
• Did you write any new necessary tests?

What does this PR do?

Until now, the model has to have only optional kwargs and could not diverge from

For example, this class could not be instantiated without knowing creation arguments

class DummyTorchModel(PyTorchRegistryMixin, torch.nn.Module):
    def __init__(self, input_size: int, output_size: int = 10):
        # PyTorchRegistryMixin.__init__ will capture these arguments
        super().__init__()
        self.fc = torch.nn.Linear(input_size, output_size)

    def forward(self, x):
        x = x.view(x.size(0), -1)

the only other way how to pull this model is Pickle which is not very space-efficient

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in GitHub issues there's a high chance it will not be merged.

Did you have fun?

Make sure you had fun coding 🙃

[Copilot]

Pull Request Overview

This PR enables Torch mixin support with initialization arguments by capturing the arguments passed to the model's init method and integrating them into push/pull registry operations. Key changes include:

• Capturing and storing init arguments in the PyTorchRegistryMixin via the new override.
• Updating file uploads and downloads to use upload_model_files/download_model_files.
• Adjusting dummy model definitions and mocks in tests to reflect the new mixin inheritance order.

Reviewed Changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
src/litmodels/integrations/mixins.py	Added new override to capture init kwargs and updated registry functions.
src/litmodels/integrations/duplicate.py	Modified workdir creation by adding a redundant os.makedirs call.
tests/integrations/test_cloud.py	Updated dummy model inheritance and mock targets for registry functions.
src/litmodels/io/gateway.py	Adjusted type hints and directory creation in upload_model.
src/litmodels/io/cloud.py	Updated type hints for upload_model_files.

[Copilot]

Pull Request Overview

This PR enhances the Torch mixin functionality by capturing model initialization arguments and updating the registry push/pull methods. The changes include replacing deprecated upload/download functions with file-based variants, updating the PyTorch mixin to capture init arguments, and revising test cases for proper inheritance order and argument types.

Reviewed Changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
src/litmodels/integrations/mixins.py	Updated mixin to capture init arguments and use upload/download files
src/litmodels/io/gateway.py	Modified type annotations and model type checks
tests/integrations/test_cloud.py	Revised dummy model inheritance order and constructor parameters
tests/integrations/test_mixins.py	Adjusted mocks to reflect new file-based upload/download functions
src/litmodels/io/cloud.py	Refined type usage for file paths in upload_model_files

[ethanwharris]

Looks nice, but could use a demo in the PR description, bit tricky to see the expected usage from the code 🙂

[Copilot]

Pull Request Overview

This PR enables the Torch mixin to capture and later reinitialize model constructor arguments, allowing models with non-optional init parameters to be registered without relying solely on pickle.

• Refactors tests to use new upload/download file functions and validates behavior with two dummy model implementations.
• Updates file type hints and integration methods in io and mixins modules to support both strings and Paths.
• Implements init argument capture in PyTorchRegistryMixin and adjusts model creation in pull_from_registry.

Reviewed Changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
tests/integrations/test_mixins.py	Updated tests for PyTorch and pickle registries with new file upload/download APIs and renamed dummy model classes.
tests/integrations/test_cloud.py	Adjusts dummy model inheritance order and constructor signature to align with mixin changes.
src/litmodels/io/gateway.py	Refines type hints and logic in upload_model to better distinguish between model types.
src/litmodels/io/cloud.py	Updates type hints for file paths to support Path objects.
src/litmodels/integrations/mixins.py	Implements init args capture in new, adds JSON dumping of constructor kwargs, and updates registry push/pull logic.

Comments suppressed due to low confidence (3)

src/litmodels/integrations/mixins.py:148

• The condition 'if self.__init_kwargs:' will evaluate to False if __init_kwargs is an empty dict, which may be a valid case. Consider using 'hasattr(self, "__init_kwargs")' to check for its presence instead.

if self.__init_kwargs:

tests/integrations/test_mixins.py:68

• [nitpick] Ensure that the expected type for the 'path' parameter in the test matches the production code. If upload_model_files accepts a Path object, either update the test to pass a Path (or its string conversion) consistently or adjust the production code to normalize the type.

mock_upload_model.assert_called_once_with(name=torch_class.__name__, path=str(tmp_path))

src/litmodels/integrations/mixins.py:217

• [nitpick] Pulling the model instance with an empty init_kwargs dict may raise errors if the init method requires non-optional arguments. Consider adding a fallback or ensuring that init_kwargs always contains the necessary parameters.

instance = cls(**init_kwargs)