Move wrappers module into networks

configs/vision/dino_vit/offline/bach.yaml

@@ -29,7 +29,7 @@ trainer:
             0: train
             1: val
           backbone:
-            class_path: eva.models.ModelFromFunction
+            class_path: eva.models.networks.wrappers.ModelFromFunction
             init_args:
               path: torch.hub.load
               arguments:

configs/vision/dino_vit/offline/crc.yaml

@@ -29,7 +29,7 @@ trainer:
             0: train
             1: val
           backbone:
-            class_path: eva.models.ModelFromFunction
+            class_path: eva.models.networks.wrappers.ModelFromFunction
             init_args:
               path: torch.hub.load
               arguments:

configs/vision/dino_vit/offline/mhist.yaml

@@ -29,7 +29,7 @@ trainer:
             0: train
             1: test
           backbone:
-            class_path: eva.models.ModelFromFunction
+            class_path: eva.models.networks.wrappers.ModelFromFunction
             init_args:
               path: torch.hub.load
               arguments:

configs/vision/dino_vit/offline/patch_camelyon.yaml

@@ -30,7 +30,7 @@ trainer:
             1: val
             2: test
           backbone:
-            class_path: eva.models.ModelFromFunction
+            class_path: eva.models.networks.wrappers.ModelFromFunction
             init_args:
               path: torch.hub.load
               arguments:

configs/vision/dino_vit/online/bach.yaml

@@ -30,7 +30,7 @@ model:
   class_path: eva.HeadModule
   init_args:
     backbone:
-      class_path: eva.models.ModelFromFunction
+      class_path: eva.models.networks.wrappers.ModelFromFunction
       init_args:
         path: torch.hub.load
         arguments:

configs/vision/dino_vit/online/crc.yaml

@@ -30,7 +30,7 @@ model:
   class_path: eva.HeadModule
   init_args:
     backbone:
-      class_path: eva.models.ModelFromFunction
+      class_path: eva.models.networks.wrappers.ModelFromFunction
       init_args:
         path: torch.hub.load
         arguments:

configs/vision/dino_vit/online/mhist.yaml

@@ -30,7 +30,7 @@ model:
   class_path: eva.HeadModule
   init_args:
     backbone:
-      class_path: eva.models.ModelFromFunction
+      class_path: eva.models.networks.wrappers.ModelFromFunction
       init_args:
         path: torch.hub.load
         arguments:

configs/vision/dino_vit/online/patch_camelyon.yaml

@@ -30,7 +30,7 @@ model:
   class_path: eva.HeadModule
   init_args:
     backbone:
-      class_path: eva.models.ModelFromFunction
+      class_path: eva.models.networks.wrappers.ModelFromFunction
       init_args:
         path: torch.hub.load
         arguments:

configs/vision/owkin/phikon/offline/bach.yaml

@@ -29,7 +29,7 @@ trainer:
             0: train
             1: val
           backbone:
-            class_path: eva.models.wrappers.HuggingFaceModel
+            class_path: eva.models.networks.wrappers.HuggingFaceModel
             init_args:
               model_name_or_path: owkin/phikon
               tensor_transforms: 

configs/vision/owkin/phikon/offline/crc.yaml

@@ -29,7 +29,7 @@ trainer:
             0: train
             1: val
           backbone:
-            class_path: eva.models.wrappers.HuggingFaceModel
+            class_path: eva.models.networks.wrappers.HuggingFaceModel
             init_args:
               model_name_or_path: owkin/phikon
               tensor_transforms: 

configs/vision/owkin/phikon/offline/mhist.yaml

@@ -29,7 +29,7 @@ trainer:
             0: train
             1: test
           backbone:
-            class_path: eva.models.wrappers.HuggingFaceModel
+            class_path: eva.models.networks.wrappers.HuggingFaceModel
             init_args:
               model_name_or_path: owkin/phikon
               tensor_transforms: 

configs/vision/owkin/phikon/offline/patch_camelyon.yaml

@@ -30,7 +30,7 @@ trainer:
             1: val
             2: test
           backbone:
-            class_path: eva.models.wrappers.HuggingFaceModel
+            class_path: eva.models.networks.wrappers.HuggingFaceModel
             init_args:
               model_name_or_path: owkin/phikon
               tensor_transforms: 

configs/vision/tests/offline/patch_camelyon.yaml

@@ -15,13 +15,13 @@ trainer:
             1: val
             2: test
           backbone:
-            class_path: eva.models.ModelFromFunction
+            class_path: eva.models.networks.wrappers.ModelFromFunction
             init_args:
               path: torch.hub.load
               arguments:
                 repo_or_dir: facebookresearch/dino:main
                 model: dino_vits16
-                pretrained: true
+                pretrained: false
               checkpoint_path: &CHECKPOINT_PATH ${oc.env:CHECKPOINT_PATH, null}
       - class_path: pytorch_lightning.callbacks.LearningRateMonitor
         init_args:

src/eva/models/__init__.py

@@ -1,6 +1,5 @@
 """Models API."""
 
 from eva.models.modules import HeadModule, InferenceModule
-from eva.models.networks import ModelFromFunction
 
-__all__ = ["HeadModule", "ModelFromFunction", "InferenceModule"]
+__all__ = ["HeadModule", "InferenceModule"]

src/eva/models/networks/__init__.py

@@ -1,6 +1,6 @@
 """Networks API."""
 
+from eva.models.networks import wrappers
 from eva.models.networks.mlp import MLP
-from eva.models.wrappers.from_function import ModelFromFunction
 
-__all__ = ["ModelFromFunction", "MLP"]
+__all__ = ["wrappers", "MLP"]

src/eva/models/networks/wrappers/__init__.py

@@ -0,0 +1,7 @@
+"""Model Wrappers API."""
+
+from eva.models.networks.wrappers.from_function import ModelFromFunction
+from eva.models.networks.wrappers.huggingface import HuggingFaceModel
+from eva.models.networks.wrappers.onnx import ONNXModel
+
+__all__ = ["ModelFromFunction", "HuggingFaceModel", "ONNXModel"]

src/eva/models/wrappers/base.py → src/eva/models/networks/wrappers/base.py

@@ -15,7 +15,8 @@ def __init__(self, tensor_transforms: Callable | None = None) -> None:
         """Initializes the model.
 
         Args:
-            tensor_transforms: The transforms to apply to the output tensor produced by the model.
+            tensor_transforms: The transforms to apply to the output tensor
+                produced by the model.
         """
         super().__init__()
 

src/eva/models/wrappers/from_function.py → src/eva/models/networks/wrappers/from_function.py

@@ -7,11 +7,11 @@
 from torch import nn
 from typing_extensions import override
 
-from eva.models import wrappers
 from eva.models.networks import _utils
+from eva.models.networks.wrappers import base
 
 
-class ModelFromFunction(wrappers.BaseModel):
+class ModelFromFunction(base.BaseModel):
     """Wrapper class for models which are initialized from functions.
 
     This is helpful for initializing models in a `.yaml` configuration file.
@@ -29,10 +29,12 @@ def __init__(
         Args:
             path: The path to the callable object (class or function).
             arguments: The extra callable function / class arguments.
-            checkpoint_path: The path to the checkpoint to load the model weights from. This is
-                currently only supported for torch model checkpoints. For other formats, the
-                checkpoint loading should be handled within the provided callable object in <path>.
-            tensor_transforms: The transforms to apply to the output tensor produced by the model.
+            checkpoint_path: The path to the checkpoint to load the model
+                weights from. This is currently only supported for torch
+                model checkpoints. For other formats, the checkpoint loading
+                should be handled within the provided callable object in <path>.
+            tensor_transforms: The transforms to apply to the output tensor
+                produced by the model.
         """
         super().__init__()
 

src/eva/models/wrappers/huggingface.py → src/eva/models/networks/wrappers/huggingface.py

@@ -6,19 +6,21 @@
 import transformers
 from typing_extensions import override
 
-from eva.models import wrappers
+from eva.models.networks.wrappers import base
 
 
-class HuggingFaceModel(wrappers.BaseModel):
+class HuggingFaceModel(base.BaseModel):
     """Wrapper class for loading HuggingFace `transformers` models."""
 
     def __init__(self, model_name_or_path: str, tensor_transforms: Callable | None = None) -> None:
         """Initializes the model.
 
         Args:
-            model_name_or_path: The model name or path to load the model from. This can be a local
-                path or a model name from the HuggingFace model hub.
-            tensor_transforms: The transforms to apply to the output tensor produced by the model.
+            model_name_or_path: The model name or path to load the model from.
+                This can be a local path or a model name from the `HuggingFace`
+                model hub.
+            tensor_transforms: The transforms to apply to the output tensor
+                produced by the model.
         """
         super().__init__(tensor_transforms=tensor_transforms)
 

src/eva/models/wrappers/onnx.py → src/eva/models/networks/wrappers/onnx.py

@@ -6,10 +6,10 @@
 import torch
 from typing_extensions import override
 
-from eva.models import wrappers
+from eva.models.networks.wrappers import base
 
 
-class ONNXModel(wrappers.BaseModel):
+class ONNXModel(base.BaseModel):
     """Wrapper class for loading ONNX models."""
 
     def __init__(

tests/eva/models/wrappers/test_from_function.py → tests/eva/models/networks/wrappers/test_from_function.py

@@ -6,7 +6,7 @@
 import torch
 from torch import nn
 
-from eva.models import networks
+from eva.models.networks import wrappers
 
 
 @pytest.mark.parametrize(
@@ -17,7 +17,7 @@
     ],
 )
 def test_model_from_function(
-    model_from_function: networks.ModelFromFunction,
+    model_from_function: wrappers.ModelFromFunction,
 ) -> None:
     """Tests the model_from_function network."""
     input_tensor = torch.Tensor(4, 10)
@@ -38,13 +38,13 @@ def test_error_model_from_function(
 ) -> None:
     """Tests the model_from_function network."""
     with pytest.raises(TypeError):
-        networks.ModelFromFunction(path=path, arguments=arguments)
+        wrappers.ModelFromFunction(path=path, arguments=arguments)
 
 
 @pytest.fixture(scope="function")
 def model_from_function(
     path: Callable[..., nn.Module],
     arguments: Dict[str, Any] | None,
-) -> networks.ModelFromFunction:
+) -> wrappers.ModelFromFunction:
     """ModelFromFunction fixture."""
-    return networks.ModelFromFunction(path=path, arguments=arguments)
+    return wrappers.ModelFromFunction(path=path, arguments=arguments)

tests/eva/models/wrappers/test_huggingface.py → tests/eva/models/networks/wrappers/test_huggingface.py

@@ -6,7 +6,7 @@
 import torch
 from transformers import modeling_outputs
 
-from eva.models import wrappers
+from eva.models.networks import wrappers
 from eva.vision.data.transforms import ExtractCLSFeatures
 
 

tests/eva/models/wrappers/test_onnx.py → tests/eva/models/networks/wrappers/test_onnx.py

@@ -8,7 +8,7 @@
 import torch
 from pytorch_lightning.demos import boring_classes
 
-from eva.models.wrappers import ONNXModel
+from eva.models.networks import wrappers
 
 
 @pytest.mark.parametrize(
@@ -22,7 +22,7 @@ def test_onnx_model(
     model_path: str, input_shape: Tuple[int, ...], expected_output_shape: Tuple[int, ...]
 ) -> None:
     """Tests the forward pass using the ONNXModel wrapper."""
-    model = ONNXModel(path=model_path)
+    model = wrappers.ONNXModel(path=model_path)
     model.eval()
 
     input_tensor = torch.rand(1, 32)