move wrappers to 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/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.

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

@@ -6,10 +6,10 @@
 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:

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)