Skip to content

Commit e5ae3a5

Browse files
committed
part fix loader
1 parent 9f00d05 commit e5ae3a5

File tree

1 file changed

+3
-0
lines changed

1 file changed

+3
-0
lines changed

autointent/_dump_tools.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import joblib
77
import numpy as np
88
import numpy.typing as npt
9+
from pydantic import BaseModel
910
from sklearn.base import BaseEstimator
1011

1112
from autointent import Embedder, Ranker, VectorIndex
@@ -65,6 +66,8 @@ def dump(obj: Any, path: Path) -> None: # noqa: ANN401
6566
joblib.dump(val, path / Dumper.estimators / key)
6667
elif isinstance(val, Ranker):
6768
val.save(str(path / Dumper.cross_encoders / key))
69+
elif isinstance(val, BaseModel):
70+
simple_attrs[f"base_model_{key}"] = val.model_dump()
6871
else:
6972
msg = f"Attribute {key} of type {type(val)} cannot be dumped to file system."
7073
logger.error(msg)

0 commit comments

Comments
 (0)