lecrapaud 0.5.0__py3-none-any.whl → 0.6.2__py3-none-any.whl

lecrapaud/api.py

@@ -36,81 +36,66 @@ import logging
 from lecrapaud.utils import logger
 from lecrapaud.db.session import init_db
 from lecrapaud.feature_selection import FeatureSelectionEngine, PreprocessModel
-from lecrapaud.model_selection import ModelSelectionEngine, ModelEngine
+from lecrapaud.model_selection import ModelSelectionEngine, ModelEngine, evaluate
 from lecrapaud.feature_engineering import FeatureEngineeringEngine, PreprocessFeature
-from lecrapaud.experiment import create_dataset
-from lecrapaud.db import Dataset
+from lecrapaud.experiment import create_experiment
+from lecrapaud.db import Experiment
+from lecrapaud.search_space import normalize_models_idx
 
 
 class LeCrapaud:
     def __init__(self, uri: str = None):
         init_db(uri=uri)
 
-    def create_experiment(self, **kwargs):
-        return Experiment(**kwargs)
+    def create_experiment(self, data: pd.DataFrame, **kwargs):
+        return App(data=data, **kwargs)
 
-    def get_experiment(self, id: int):
-        return Experiment(id)
+    def get_experiment(self, id: int, **kwargs):
+        return App(id=id, **kwargs)
 
 
-class Experiment:
-    def __init__(self, id=None, **kwargs):
+class App:
+    def __init__(self, id=None, data=None, **kwargs):
         if id:
-            self.dataset = Dataset.get(id)
+            self.experiment = Experiment.get(id)
+            kwargs.update(self.experiment.context)
         else:
-            self.dataset = create_dataset(**kwargs)
+            self.experiment = create_experiment(data=data, **kwargs)
 
         for key, value in kwargs.items():
+            if key == "models_idx":
+                value = normalize_models_idx(value)
             setattr(self, key, value)
 
-        self.context = {
-            # generic
-            "dataset": self.dataset,
-            # for FeatureEngineering
-            "columns_drop": self.columns_drop,
-            "columns_boolean": self.columns_boolean,
-            "columns_date": self.columns_date,
-            "columns_te_groupby": self.columns_te_groupby,
-            "columns_te_target": self.columns_te_target,
-            # for PreprocessFeature
-            "time_series": self.time_series,
-            "date_column": self.date_column,
-            "group_column": self.group_column,
-            "val_size": self.val_size,
-            "test_size": self.test_size,
-            "columns_pca": self.columns_pca,
-            "columns_onehot": self.columns_onehot,
-            "columns_binary": self.columns_binary,
-            "columns_frequency": self.columns_frequency,
-            "columns_ordinal": self.columns_ordinal,
-            "target_numbers": self.target_numbers,
-            "target_clf": self.target_clf,
-            # for PreprocessModel
-            "models_idx": self.models_idx,
-            "max_timesteps": self.max_timesteps,
-            # for ModelSelection
-            "perform_hyperopt": self.perform_hyperopt,
-            "number_of_trials": self.number_of_trials,
-            "perform_crossval": self.perform_crossval,
-            "plot": self.plot,
-            "preserve_model": self.preserve_model,
-            # not yet
-            "target_mclf": self.target_mclf,
-        }
-
     def train(self, data):
+        logger.info("Running training...")
+
         data_eng = self.feature_engineering(data)
+        logger.info("Feature engineering done.")
+
         train, val, test = self.preprocess_feature(data_eng)
-        all_features = self.feature_selection(train)
+        logger.info("Feature preprocessing done.")
+
+        self.feature_selection(train)
+        logger.info("Feature selection done.")
+
         std_data, reshaped_data = self.preprocess_model(train, val, test)
+        logger.info("Model preprocessing done.")
+
         self.model_selection(std_data, reshaped_data)
+        logger.info("Model selection done.")
 
     def predict(self, new_data, verbose: int = 0):
+        # for scores if TARGET is in columns
+        scores_reg = []
+        scores_clf = []
+
         if verbose == 0:
             logger.setLevel(logging.WARNING)
 
         logger.warning("Running prediction...")
 
+        # feature engineering + preprocessing
         data = self.feature_engineering(
             data=new_data,
             for_training=False,
@@ -123,16 +108,16 @@ class Experiment:
         for target_number in self.target_numbers:
 
             # loading model
-            training_target_dir = f"{self.dataset.path}/TARGET_{target_number}"
-            all_features = self.dataset.get_all_features(
+            training_target_dir = f"{self.experiment.path}/TARGET_{target_number}"
+            all_features = self.experiment.get_all_features(
                 date_column=self.date_column, group_column=self.group_column
             )
-            if self.dataset.name == "data_28_X_X":
+            if self.experiment.name == "data_28_X_X":
                 features = joblib.load(
-                    f"{self.dataset.path}/preprocessing/features_{target_number}.pkl"
+                    f"{self.experiment.path}/preprocessing/features_{target_number}.pkl"
                 )  # we keep this for backward compatibility
             else:
-                features = self.dataset.get_features(target_number)
+                features = self.experiment.get_features(target_number)
             model = ModelEngine(path=training_target_dir)
 
             # getting data
@@ -151,7 +136,7 @@ class Experiment:
             if model.recurrent:
                 y_pred.index = (
                     new_data.index
-                )  # TODO: not sure this will work for old dataset not aligned with data_for_training for test use case (done, this is why we decode the test set)
+                )  # TODO: not sure this will work for old experiment not aligned with data_for_training for test use case (done, this is why we decode the test set)
 
             # unscaling prediction
             if (
@@ -165,6 +150,26 @@ class Experiment:
                     ).flatten(),
                     index=new_data.index,
                 )
+                y_pred.name = "PRED"
+
+            # evaluate if TARGET is in columns
+            if f"TARGET_{target_number}" in new_data.columns:
+                y_true = new_data[f"TARGET_{target_number}"]
+                prediction = pd.concat([y_true, y_pred], axis=1)
+                prediction.rename(
+                    columns={f"TARGET_{target_number}": "TARGET"}, inplace=True
+                )
+                print(prediction)
+                score = evaluate(
+                    prediction,
+                    target_type=model.target_type,
+                )
+                score["TARGET"] = f"TARGET_{target_number}"
+
+                if model.target_type == "classification":
+                    scores_clf.append(score)
+                else:
+                    scores_reg.append(score)
 
             # renaming pred column and concatenating with initial data
             if isinstance(y_pred, pd.DataFrame):
@@ -179,7 +184,11 @@ class Experiment:
                 y_pred.name = f"TARGET_{target_number}_PRED"
                 new_data = pd.concat([new_data, y_pred], axis=1)
 
-        return new_data
+        if len(scores_reg) > 0:
+            scores_reg = pd.DataFrame(scores_reg).set_index("TARGET")
+        if len(scores_clf) > 0:
+            scores_clf = pd.DataFrame(scores_clf).set_index("TARGET")
+        return new_data, scores_reg, scores_clf
 
     def feature_engineering(self, data, for_training=True):
         app = FeatureEngineeringEngine(
@@ -197,7 +206,7 @@ class Experiment:
     def preprocess_feature(self, data, for_training=True):
         app = PreprocessFeature(
             data=data,
-            dataset=self.dataset,
+            experiment=self.experiment,
             time_series=self.time_series,
             date_column=self.date_column,
             group_column=self.group_column,
@@ -223,12 +232,12 @@ class Experiment:
             app = FeatureSelectionEngine(
                 train=train,
                 target_number=target_number,
-                dataset=self.dataset,
+                experiment=self.experiment,
                 target_clf=self.target_clf,
             )
             app.run()
-        self.dataset = Dataset.get(self.dataset.id)
-        all_features = self.dataset.get_all_features(
+        self.experiment = Experiment.get(self.experiment.id)
+        all_features = self.experiment.get_all_features(
             date_column=self.date_column, group_column=self.group_column
         )
         return all_features
@@ -238,7 +247,7 @@ class Experiment:
             train=train,
             val=val,
             test=test,
-            dataset=self.dataset,
+            experiment=self.experiment,
             target_numbers=self.target_numbers,
             target_clf=self.target_clf,
             models_idx=self.models_idx,
@@ -260,15 +269,16 @@ class Experiment:
                 data=data,
                 reshaped_data=reshaped_data,
                 target_number=target_number,
-                dataset=self.dataset,
+                experiment=self.experiment,
                 target_clf=self.target_clf,
                 models_idx=self.models_idx,
                 time_series=self.time_series,
                 date_column=self.date_column,
                 group_column=self.group_column,
+                target_clf_thresholds=self.target_clf_thresholds,
             )
             app.run(
-                self.session_name,
+                self.experiment_name,
                 perform_hyperopt=self.perform_hyperopt,
                 number_of_trials=self.number_of_trials,
                 perform_crossval=self.perform_crossval,

lecrapaud/config.py

@@ -12,9 +12,9 @@ DB_USER = (
     os.getenv("TEST_DB_USER") if PYTHON_ENV == "Test" else os.getenv("DB_USER", None)
 )
 DB_PASSWORD = (
-    os.getenv("TEST_DB_PASSWORD")
+    os.getenv("TEST_DB_PASSWORD", "")
     if PYTHON_ENV == "Test"
-    else os.getenv("DB_PASSWORD", None)
+    else os.getenv("DB_PASSWORD", "")
 )
 DB_HOST = (
     os.getenv("TEST_DB_HOST") if PYTHON_ENV == "Test" else os.getenv("DB_HOST", None)
@@ -25,5 +25,9 @@ DB_PORT = (
 DB_NAME = (
     os.getenv("TEST_DB_NAME") if PYTHON_ENV == "Test" else os.getenv("DB_NAME", None)
 )
-DB_URI = os.getenv("DB_URI", None)
+DB_URI = (
+    os.getenv("TEST_DB_URI", None)
+    if PYTHON_ENV == "Test"
+    else os.getenv("DB_URI", None)
+)
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

lecrapaud/db/alembic/versions/{2025_06_20_1924-1edada319fd7_initial_setup.py → 2025_06_23_1748-f089dfb7e3ba_.py}

@@ -1,8 +1,8 @@
-"""initial_setup
+"""
 
-Revision ID: 1edada319fd7
+Revision ID: f089dfb7e3ba
 Revises: 
-Create Date: 2025-06-20 19:24:25.033055
+Create Date: 2025-06-23 17:48:32.842030
 
 """
 from typing import Sequence, Union
@@ -12,7 +12,7 @@ import sqlalchemy as sa
 
 
 # revision identifiers, used by Alembic.
-revision: str = '1edada319fd7'
+revision: str = 'f089dfb7e3ba'
 down_revision: Union[str, None] = None
 branch_labels: Union[str, Sequence[str], None] = None
 depends_on: Union[str, Sequence[str], None] = None
@@ -20,7 +20,7 @@ depends_on: Union[str, Sequence[str], None] = None
 
 def upgrade() -> None:
     # ### commands auto generated by Alembic - please adjust! ###
-    op.create_table('lecrapaud_datasets',
+    op.create_table('lecrapaud_experiments',
     sa.Column('id', sa.BigInteger(), autoincrement=True, nullable=False),
     sa.Column('created_at', sa.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
     sa.Column('updated_at', sa.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
@@ -45,9 +45,9 @@ def upgrade() -> None:
     sa.Column('test_start_date', sa.DateTime(), nullable=True),
     sa.Column('test_end_date', sa.DateTime(), nullable=True),
     sa.PrimaryKeyConstraint('id'),
-    sa.UniqueConstraint('name', name='uq_datasets_composite')
+    sa.UniqueConstraint('name', name='uq_experiments_composite')
     )
-    op.create_index(op.f('ix_lecrapaud_datasets_id'), 'lecrapaud_datasets', ['id'], unique=False)
+    op.create_index(op.f('ix_lecrapaud_experiments_id'), 'lecrapaud_experiments', ['id'], unique=False)
     op.create_table('lecrapaud_features',
     sa.Column('id', sa.BigInteger(), autoincrement=True, nullable=False),
     sa.Column('created_at', sa.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
@@ -79,12 +79,12 @@ def upgrade() -> None:
     sa.UniqueConstraint('name', 'type', name='uq_target_composite')
     )
     op.create_index(op.f('ix_lecrapaud_targets_id'), 'lecrapaud_targets', ['id'], unique=False)
-    op.create_table('lecrapaud_dataset_target_association',
-    sa.Column('dataset_id', sa.BigInteger(), nullable=False),
+    op.create_table('lecrapaud_experiment_target_association',
+    sa.Column('experiment_id', sa.BigInteger(), nullable=False),
     sa.Column('target_id', sa.BigInteger(), nullable=False),
-    sa.ForeignKeyConstraint(['dataset_id'], ['lecrapaud_datasets.id'], ondelete='CASCADE'),
+    sa.ForeignKeyConstraint(['experiment_id'], ['lecrapaud_experiments.id'], ondelete='CASCADE'),
     sa.ForeignKeyConstraint(['target_id'], ['lecrapaud_targets.id'], ondelete='CASCADE'),
-    sa.PrimaryKeyConstraint('dataset_id', 'target_id')
+    sa.PrimaryKeyConstraint('experiment_id', 'target_id')
     )
     op.create_table('lecrapaud_feature_selections',
     sa.Column('id', sa.BigInteger(), autoincrement=True, nullable=False),
@@ -92,12 +92,12 @@ def upgrade() -> None:
     sa.Column('updated_at', sa.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
     sa.Column('training_time', sa.Integer(), nullable=True),
     sa.Column('best_features_path', sa.String(length=255), nullable=True),
-    sa.Column('dataset_id', sa.BigInteger(), nullable=False),
+    sa.Column('experiment_id', sa.BigInteger(), nullable=False),
     sa.Column('target_id', sa.BigInteger(), nullable=False),
-    sa.ForeignKeyConstraint(['dataset_id'], ['lecrapaud_datasets.id'], ondelete='CASCADE'),
+    sa.ForeignKeyConstraint(['experiment_id'], ['lecrapaud_experiments.id'], ondelete='CASCADE'),
     sa.ForeignKeyConstraint(['target_id'], ['lecrapaud_targets.id'], ondelete='CASCADE'),
     sa.PrimaryKeyConstraint('id'),
-    sa.UniqueConstraint('dataset_id', 'target_id', name='uq_feature_selection_composite')
+    sa.UniqueConstraint('experiment_id', 'target_id', name='uq_feature_selection_composite')
     )
     op.create_index(op.f('ix_lecrapaud_feature_selections_id'), 'lecrapaud_feature_selections', ['id'], unique=False)
     op.create_table('lecrapaud_model_selections',
@@ -108,12 +108,12 @@ def upgrade() -> None:
     sa.Column('best_model_path', sa.String(length=255), nullable=True),
     sa.Column('best_model_id', sa.BigInteger(), nullable=True),
     sa.Column('target_id', sa.BigInteger(), nullable=False),
-    sa.Column('dataset_id', sa.BigInteger(), nullable=False),
+    sa.Column('experiment_id', sa.BigInteger(), nullable=False),
     sa.ForeignKeyConstraint(['best_model_id'], ['lecrapaud_models.id'], ondelete='CASCADE'),
-    sa.ForeignKeyConstraint(['dataset_id'], ['lecrapaud_datasets.id'], ondelete='CASCADE'),
+    sa.ForeignKeyConstraint(['experiment_id'], ['lecrapaud_experiments.id'], ondelete='CASCADE'),
     sa.ForeignKeyConstraint(['target_id'], ['lecrapaud_targets.id'], ondelete='CASCADE'),
     sa.PrimaryKeyConstraint('id'),
-    sa.UniqueConstraint('target_id', 'dataset_id', name='uq_model_selection_composite')
+    sa.UniqueConstraint('target_id', 'experiment_id', name='uq_model_selection_composite')
     )
     op.create_index(op.f('ix_lecrapaud_model_selections_id'), 'lecrapaud_model_selections', ['id'], unique=False)
     op.create_table('lecrapaud_feature_selection_association',
@@ -202,13 +202,13 @@ def downgrade() -> None:
     op.drop_table('lecrapaud_model_selections')
     op.drop_index(op.f('ix_lecrapaud_feature_selections_id'), table_name='lecrapaud_feature_selections')
     op.drop_table('lecrapaud_feature_selections')
-    op.drop_table('lecrapaud_dataset_target_association')
+    op.drop_table('lecrapaud_experiment_target_association')
     op.drop_index(op.f('ix_lecrapaud_targets_id'), table_name='lecrapaud_targets')
     op.drop_table('lecrapaud_targets')
     op.drop_index(op.f('ix_lecrapaud_models_id'), table_name='lecrapaud_models')
     op.drop_table('lecrapaud_models')
     op.drop_index(op.f('ix_lecrapaud_features_id'), table_name='lecrapaud_features')
     op.drop_table('lecrapaud_features')
-    op.drop_index(op.f('ix_lecrapaud_datasets_id'), table_name='lecrapaud_datasets')
-    op.drop_table('lecrapaud_datasets')
+    op.drop_index(op.f('ix_lecrapaud_experiments_id'), table_name='lecrapaud_experiments')
+    op.drop_table('lecrapaud_experiments')
     # ### end Alembic commands ###

lecrapaud/db/alembic/versions/2025_06_24_1216-c62251b129ed_.py

@@ -0,0 +1,30 @@
+"""
+
+Revision ID: c62251b129ed
+Revises: f089dfb7e3ba
+Create Date: 2025-06-24 12:16:21.949079
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = 'c62251b129ed'
+down_revision: Union[str, None] = 'f089dfb7e3ba'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column('lecrapaud_experiments', sa.Column('context', sa.JSON(), nullable=True))
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column('lecrapaud_experiments', 'context')
+    # ### end Alembic commands ###

lecrapaud/db/alembic/versions/2025_06_24_1711-86457e2f333f_.py

@@ -0,0 +1,34 @@
+"""
+
+Revision ID: 86457e2f333f
+Revises: c62251b129ed
+Create Date: 2025-06-24 17:11:25.187876
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import mysql
+
+# revision identifiers, used by Alembic.
+revision: str = '86457e2f333f'
+down_revision: Union[str, None] = 'c62251b129ed'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column('lecrapaud_scores', sa.Column('thresholds', sa.JSON(), nullable=True))
+    op.add_column('lecrapaud_scores', sa.Column('f1_at_threshold', sa.Float(), nullable=True))
+    op.drop_column('lecrapaud_scores', 'threshold')
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column('lecrapaud_scores', sa.Column('threshold', mysql.FLOAT(), nullable=True))
+    op.drop_column('lecrapaud_scores', 'f1_at_threshold')
+    op.drop_column('lecrapaud_scores', 'thresholds')
+    # ### end Alembic commands ###

lecrapaud/db/models/__init__.py

@@ -1,6 +1,5 @@
 from lecrapaud.db.models.base import Base
-
-from lecrapaud.db.models.dataset import Dataset
+from lecrapaud.db.models.experiment import Experiment
 from lecrapaud.db.models.feature_selection_rank import FeatureSelectionRank
 from lecrapaud.db.models.feature_selection import FeatureSelection
 from lecrapaud.db.models.feature import Feature
@@ -9,3 +8,16 @@ from lecrapaud.db.models.model_training import ModelTraining
 from lecrapaud.db.models.model import Model
 from lecrapaud.db.models.score import Score
 from lecrapaud.db.models.target import Target
+
+__all__ = [
+    'Base',
+    'Experiment',
+    'FeatureSelectionRank',
+    'FeatureSelection',
+    'Feature',
+    'ModelSelection',
+    'ModelTraining',
+    'Model',
+    'Score',
+    'Target',
+]

lecrapaud/db/models/base.py

@@ -9,6 +9,7 @@ from sqlalchemy.inspection import inspect
 from sqlalchemy.orm.attributes import InstrumentedAttribute
 from lecrapaud.db.session import get_db
 from sqlalchemy.ext.declarative import declared_attr
+from sqlalchemy.dialects.mysql import insert as mysql_insert
 
 
 def with_db(func):
@@ -98,9 +99,53 @@ class Base(DeclarativeBase):
                 }
                 for row in results
             ]
-
         return results
 
+    @classmethod
+    @with_db
+    def upsert_bulk(cls, db=None, match_fields: list[str] = None, **kwargs):
+        """
+        Performs a bulk upsert into the database using ON DUPLICATE KEY UPDATE.
+
+        Args:
+            db (Session): SQLAlchemy DB session
+            match_fields (list[str]): Fields to match on for deduplication
+            **kwargs: Column-wise keyword arguments (field_name=[...])
+        """
+        # Ensure all provided fields have values of equal length
+        value_lengths = [len(v) for v in kwargs.values()]
+        if not value_lengths or len(set(value_lengths)) != 1:
+            raise ValueError(
+                "All field values must be non-empty lists of the same length."
+            )
+
+        # Convert column-wise kwargs to row-wise list of dicts
+        items = [dict(zip(kwargs.keys(), row)) for row in zip(*kwargs.values())]
+        if not items:
+            return
+
+        stmt = mysql_insert(cls.__table__).values(items)
+
+        # Default to primary keys if match_fields not provided
+        if not match_fields:
+            match_fields = [col.name for col in cls.__table__.primary_key.columns]
+
+        # Ensure all columns to be updated are in the insert
+        update_dict = {
+            c.name: stmt.inserted[c.name]
+            for c in cls.__table__.columns
+            if c.name not in match_fields and c.name in items[0]
+        }
+
+        if not update_dict:
+            # Avoid triggering ON DUPLICATE KEY UPDATE with empty dict
+            db.execute(stmt.prefix_with("IGNORE"))
+        else:
+            upsert_stmt = stmt.on_duplicate_key_update(**update_dict)
+            db.execute(upsert_stmt)
+
+        db.commit()
+
     @classmethod
     @with_db
     def filter(cls, db=None, **kwargs):
@@ -165,7 +210,8 @@ class Base(DeclarativeBase):
 
         if instance:
             for key, value in kwargs.items():
-                setattr(instance, key, value)
+                if key != "id":
+                    setattr(instance, key, value)
         else:
             instance = cls(**kwargs)
             db.add(instance)

lecrapaud/db/models/{dataset.py → experiment.py}

@@ -1,33 +1,31 @@
+from itertools import chain
+
 from sqlalchemy import (
     Column,
     Integer,
     String,
     DateTime,
-    Date,
     Float,
     JSON,
     Table,
     ForeignKey,
     BigInteger,
-    Index,
     TIMESTAMP,
     UniqueConstraint,
+    func,
 )
-from sqlalchemy import desc, asc, cast, text, func
-from sqlalchemy.orm import relationship, Mapped, mapped_column, DeclarativeBase
-from itertools import chain
+from sqlalchemy.orm import relationship
 
-from lecrapaud.db.session import get_db
 from lecrapaud.db.models.base import Base
 
 # jointures
-lecrapaud_dataset_target_association = Table(
-    "lecrapaud_dataset_target_association",
+lecrapaud_experiment_target_association = Table(
+    "lecrapaud_experiment_target_association",
     Base.metadata,
     Column(
-        "dataset_id",
+        "experiment_id",
         BigInteger,
-        ForeignKey("lecrapaud_datasets.id", ondelete="CASCADE"),
+        ForeignKey("lecrapaud_experiments.id", ondelete="CASCADE"),
         primary_key=True,
     ),
     Column(
@@ -39,7 +37,7 @@ lecrapaud_dataset_target_association = Table(
 )
 
 
-class Dataset(Base):
+class Experiment(Base):
 
     id = Column(BigInteger, primary_key=True, index=True, autoincrement=True)
     created_at = Column(
@@ -71,30 +69,31 @@ class Dataset(Base):
     val_end_date = Column(DateTime)
     test_start_date = Column(DateTime)
     test_end_date = Column(DateTime)
+    context = Column(JSON)
 
     feature_selections = relationship(
         "FeatureSelection",
-        back_populates="dataset",
+        back_populates="experiment",
         cascade="all, delete-orphan",
         lazy="selectin",
     )
     model_selections = relationship(
         "ModelSelection",
-        back_populates="dataset",
+        back_populates="experiment",
         cascade="all, delete-orphan",
         lazy="selectin",
     )
     targets = relationship(
         "Target",
-        secondary=lecrapaud_dataset_target_association,
-        back_populates="datasets",
+        secondary=lecrapaud_experiment_target_association,
+        back_populates="experiments",
         lazy="selectin",
     )
 
     __table_args__ = (
         UniqueConstraint(
             "name",
-            name="uq_datasets_composite",
+            name="uq_experiments_composite",
         ),
     )
 
@@ -106,23 +105,22 @@ class Dataset(Base):
         feature_selection = [
             fs for fs in feature_selections if fs.target_id == target_id
         ][0]
-        feature = [f.name for f in feature_selection.features]
-        return feature
+        features = [f.name for f in feature_selection.features]
+        return features
 
     def get_all_features(self, date_column: str = None, group_column: str = None):
         target_idx = [target.id for target in self.targets]
+        _all_features = chain.from_iterable(
+            [f.name for f in fs.features]
+            for fs in self.feature_selections
+            if fs.target_id in target_idx
+        )
         all_features = []
         if date_column:
             all_features.append(date_column)
         if group_column:
             all_features.append(group_column)
-        all_features += list(
-            chain.from_iterable(
-                [f.name for f in fs.features]
-                for fs in self.feature_selections
-                if fs.target_id in target_idx
-            )
-        )
+        all_features += list(_all_features)
         all_features = list(dict.fromkeys(all_features))
 
         return all_features