Merge pull request #95 from ayasyrev/pydantic_v2

Pydantic v2

Author: ayasyrev

requirements.txt

@@ -1,2 +1,3 @@
+pydantic
 # pytorch
 # numpy

src/model_constructor/model_constructor.py

@@ -3,7 +3,7 @@
 from typing import Any, Callable, Optional, TypeVar, Union
 
 import torch
-from pydantic import BaseModel, validator
+from pydantic import BaseModel, field_validator
 from torch import nn
 
 from .helpers import nn_seq
@@ -35,7 +35,7 @@ def init_cnn(module: nn.Module) -> None:
 
 class BasicBlock(nn.Module):
     """Basic Resnet block.
-    Configurable - can use pool to reduce at identity path, change act etc. """
+    Configurable - can use pool to reduce at identity path, change act etc."""
 
     def __init__(
         self,
@@ -119,7 +119,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:  # type: ignore
 
 class BottleneckBlock(nn.Module):
     """Bottleneck Resnet block.
-    Configurable - can use pool to reduce at identity path, change act etc. """
+    Configurable - can use pool to reduce at identity path, change act etc."""
 
     def __init__(
         self,
@@ -285,7 +285,7 @@ def make_head(cfg: TModelCfg) -> nn.Sequential:  # type: ignore
     return nn_seq(head)
 
 
-class ModelCfg(BaseModel):
+class ModelCfg(BaseModel, arbitrary_types_allowed=True, extra="forbid"):
     """Model constructor Config. As default - xresnet18"""
 
     name: Optional[str] = None
@@ -320,10 +320,6 @@ class ModelCfg(BaseModel):
     make_body: Callable[[TModelCfg], Union[nn.Module, nn.Sequential]] = make_body  # type: ignore
     make_head: Callable[[TModelCfg], Union[nn.Module, nn.Sequential]] = make_head  # type: ignore
 
-    class Config:  # pylint: disable=too-few-public-methods
-        arbitrary_types_allowed = True
-        extra = "forbid"
-
     def _get_str_value(self, field: str) -> str:
         value = getattr(self, field)
         if isinstance(value, type):
@@ -340,15 +336,15 @@ def __repr__(self) -> str:
     def __repr_args__(self) -> list[tuple[str, str]]:
         return [
             (field, str_value)
-            for field in self.__fields__
+            for field in self.model_fields
             if (str_value := self._get_str_value(field))
         ]
 
     def __repr_changed_args__(self) -> list[str]:
         """Return list repr for changed fields"""
         return [
             f"{field}: {self._get_str_value(field)}"
-            for field in self.__fields_set__
+            for field in self.model_fields_set
             if field != "name"
         ]
 
@@ -370,7 +366,7 @@ def print_changed(self) -> None:
 class ModelConstructor(ModelCfg):
     """Model constructor. As default - resnet18"""
 
-    @validator("se")
+    @field_validator("se")
     def set_se(  # pylint: disable=no-self-argument
         cls, value: Union[bool, type[nn.Module]]
     ) -> Union[bool, type[nn.Module]]:
@@ -379,7 +375,7 @@ def set_se(  # pylint: disable=no-self-argument
                 return SEModule
         return value
 
-    @validator("sa")
+    @field_validator("sa")
     def set_sa(  # pylint: disable=no-self-argument
         cls, value: Union[bool, type[nn.Module]]
     ) -> Union[bool, type[nn.Module]]:
@@ -388,7 +384,7 @@ def set_sa(  # pylint: disable=no-self-argument
                 return SimpleSelfAttention  # default: ks=1, sym=sym
         return value
 
-    @validator("se_module", "se_reduction")  # pragma: no cover
+    @field_validator("se_module", "se_reduction")  # pragma: no cover
     def deprecation_warning(  # pylint: disable=no-self-argument
         cls, value: Union[bool, int, None]
     ) -> Union[bool, int, None]:
@@ -409,12 +405,14 @@ def body(self):
 
     @classmethod
     def from_cfg(cls, cfg: ModelCfg):
-        return cls(**cfg.dict())
+        return cls(**cfg.model_dump())
 
     @classmethod
-    def create_model(cls, cfg: Union[ModelCfg, None] = None, **kwargs: dict[str, Any]) -> nn.Sequential:
+    def create_model(
+        cls, cfg: Union[ModelCfg, None] = None, **kwargs: dict[str, Any]
+    ) -> nn.Sequential:
         if cfg:
-            return cls(**cfg.dict())()
+            return cls(**cfg.model_dump())()
         return cls(**kwargs)()
 
     def __call__(self) -> nn.Sequential: