TP2 Partie 1 : package lab + DVC (split full_history v1)

lab/split/cli.py

@@ -0,0 +1,35 @@
+import logging
+
+import pandas as pd
+
+from .. import constants
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def main():
+    # Entrees : dataset source deja prepare (hors git)
+    feature_file_path = constants.SOURCE_DIR / constants.FEATURE_FILENAME
+    target_file_path = constants.SOURCE_DIR / constants.TARGET_FILENAME
+
+    logger.info(f"Read dataset from {feature_file_path} and {target_file_path}")
+    df_features = pd.read_parquet(feature_file_path)
+    df_target = pd.read_parquet(target_file_path)
+
+    df = df_features.join(df_target)
+    timestamps = df.index.get_level_values("timestamp")
+
+    # Sorties : splits versionnes par DVC dans le depot
+    constants.DATASET_DIR.mkdir(parents=True, exist_ok=True)
+    logger.info(f"Split strategy: {constants.CHOSEN_SPLIT_STRATEGY.value}")
+    for dataset_name, (start_date, end_date) in constants.DATASET_SPLIT_DATES[constants.CHOSEN_SPLIT_STRATEGY].items():
+        file_path = constants.DATASET_DIR / f"{dataset_name}.parquet"
+        mask = (timestamps.date >= start_date) & (timestamps.date <= end_date)
+        df_split = df[mask]
+        logger.info(f"Split dataset into {dataset_name} ({start_date} -> {end_date}) with shape: {df_split.shape}")
+        df_split.to_parquet(file_path)
+
+
+if __name__ == "__main__":
+    main()