diff --git a/services/ml/app/db.py b/services/ml/app/db.py
index 13811b7..8fc8475 100644
--- a/services/ml/app/db.py
+++ b/services/ml/app/db.py
@@ -10,8 +10,7 @@ from contextlib import contextmanager
 from typing import Generator
 
 from sqlalchemy import create_engine, Column, Integer, String, DateTime, JSON
-from sqlalchemy.ext.declarative import declarative_base
-from sqlalchemy.orm import sessionmaker, Session
+from sqlalchemy.orm import DeclarativeBase, sessionmaker, Session
 from sqlalchemy.sql import func
 
 
@@ -35,7 +34,8 @@ engine = create_engine(
 SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
 
 # Base class for declarative models
-Base = declarative_base()
+class Base(DeclarativeBase):
+    pass
 
 
 # Training runs model