faiss to ML pipeline

hypex/executor/__init__.py

@@ -1,6 +1,5 @@
 from .calculators import MinSampleSize
 from .executor import Calculator, Executor, IfExecutor
-from .feature_ml_executor import FeatureMLExecutor
 from .ml_executor import MLExecutor
 from .state import MLExecutorParams
 
@@ -10,6 +9,5 @@
 	"IfExecutor",
 	"MLExecutor",
 	"MLExecutorParams",
-	"FeatureMLExecutor",
 	"MinSampleSize",
 ]

hypex/experiments/artifact.py

@@ -194,15 +194,32 @@ def _save_metadata(self) -> None:
     def _save_ml_executor_states(self) -> None:
         """
         Save fitted ML executor states.
-        
+
         Structure:
         ml_executors/
-            StandardScaler_abc123.json
+            StandardScaler_abc123.json  (JSON-serializable states)
+            FaissMLExecutor_def456.pkl  (binary states via pickle)
         """
+        import pickle
+
         for executor_id, state in self.ml_executor_states.items():
-            state_file = os.path.join(self.ml_executors_dir, f"{executor_id}.json")
-            with open(state_file, 'w') as f:
-                json.dump(state.to_dict(), f, indent=2)
+            if hasattr(state, 'to_dict'):
+                state_dict = state.to_dict()
+                # Try JSON serialization, fall back to pickle for binary objects
+                try:
+                    state_file = os.path.join(self.ml_executors_dir, f"{executor_id}.json")
+                    with open(state_file, 'w') as f:
+                        json.dump(state_dict, f, indent=2)
+                except (TypeError, ValueError):
+                    # Contains non-serializable objects, use pickle
+                    state_file = os.path.join(self.ml_executors_dir, f"{executor_id}.pkl")
+                    with open(state_file, 'wb') as f:
+                        pickle.dump(state, f)
+            else:
+                # No to_dict method, use pickle directly
+                state_file = os.path.join(self.ml_executors_dir, f"{executor_id}.pkl")
+                with open(state_file, 'wb') as f:
+                    pickle.dump(state, f)
 
     def _save_models(self) -> None:
         """
@@ -242,10 +259,11 @@ def _save_models(self) -> None:
     def load_ml_executor_states(self) -> Dict[str, Any]:
         """
         Load fitted ML executor states from disk.
-        
+
         Returns:
-            Dict of executor_id -> MLExecutorParams
+            Dict of executor_id -> MLExecutorParams or pickled state
         """
+        import pickle
         from ..executor.state import MLExecutorParams
 
         ml_executor_states = {}
@@ -254,15 +272,21 @@ def load_ml_executor_states(self) -> Dict[str, Any]:
             return ml_executor_states
 
         for filename in os.listdir(self.ml_executors_dir):
+            state_file = os.path.join(self.ml_executors_dir, filename)
+
             if filename.endswith('.json'):
                 executor_id = filename[:-5]
-                state_file = os.path.join(self.ml_executors_dir, filename)
-
                 with open(state_file, 'r') as f:
                     state_dict = json.load(f)
                 state = MLExecutorParams.from_dict(state_dict)
                 ml_executor_states[executor_id] = state
 
+            elif filename.endswith('.pkl'):
+                executor_id = filename[:-4]
+                with open(state_file, 'rb') as f:
+                    state = pickle.load(f)
+                ml_executor_states[executor_id] = state
+
         return ml_executor_states
 
     def load_models(self) -> Dict[str, Dict[str, Dict[str, Any]]]:

hypex/extensions/faiss.py

@@ -51,30 +51,38 @@ def _calc_pandas(
         self,
         data: Dataset,
         test_data: Dataset | None = None,
-        mode: Literal["auto", "fit", "predict"] | None = None,
+        mode: Literal["auto", "fit", "predict", "fit_predict"] | None = None,
         **kwargs,
     ):
         mode = mode or "auto"
         X = data.data.values
-        test = test_data.data.values
-        if mode in ["auto", "fit"]:
+
+        # Fit mode: train index on data
+        if mode in ["auto", "fit", "fit_predict"]:
             self.index = faiss.IndexFlatL2(X.shape[1])
-            if (
-                (
-                    (len(X) > 1_000_000 and self.faiss_mode == "auto")
-                    or self.faiss_mode == "fast"
-                )
-                and len(X) > 1_000
-                and len(test) > 1_000
-            ):
-                self.index = faiss.IndexIVFFlat(self.index, X.shape[1], 1000)
-                self.index.train(X)
+            if test_data is not None:
+                test = test_data.data.values
+                if (
+                    (
+                        (len(X) > 1_000_000 and self.faiss_mode == "auto")
+                        or self.faiss_mode == "fast"
+                    )
+                    and len(X) > 1_000
+                    and len(test) > 1_000
+                ):
+                    self.index = faiss.IndexIVFFlat(self.index, X.shape[1], 1000)
+                    self.index.train(X)
             self.index.add(X)
-        if mode in ["auto", "predict"]:
+
+        # Predict mode: use index to find neighbors
+        if mode in ["auto", "predict", "fit_predict"]:
             if test_data is None:
                 raise ValueError("test_data is needed for evaluation")
-            X = test_data.data.values if mode == "auto" else data.data.values
-            return self._predict(data, test_data, X)
+            # Always query using test_data (data to find neighbors for)
+            X_predict = test_data.data.values
+            return self._predict(data, test_data, X_predict)
+
+        # Fit only mode: return self
         return self
 
     def fit(self, X: Dataset, Y: Dataset | None = None, **kwargs):

hypex/matching.py

@@ -10,7 +10,8 @@
 from .executor import Executor
 from .experiments import GroupExperiment
 from .experiments.base import Experiment, OnRoleExperiment
-from .ml.faiss import FaissNearestNeighbors
+from .experiments.ml import MLExperiment
+from .ml.faiss import FaissMLExecutor
 from .operators.operators import Bias, MatchingMetrics
 from .reporters.matching import MatchingDatasetReporter
 from .transformers import TypeCaster
@@ -147,12 +148,17 @@ def _make_experiment(
                 dtype={int: float},
                 roles=[FeatureRole(), TargetRole()],
             ),
-            FaissNearestNeighbors(
-                grouping_role=TreatmentRole(),
-                two_sides=two_sides,
-                test_pairs=test_pairs,
-                faiss_mode=faiss_mode,
-                n_neighbors=n_neighbors,
+            MLExperiment(
+                ml_executors=[
+                    FaissMLExecutor(
+                        grouping_role=TreatmentRole(),
+                        two_sides=two_sides,
+                        test_pairs=test_pairs,
+                        faiss_mode=faiss_mode,
+                        n_neighbors=n_neighbors,
+                    ),
+                ],
+                mode="fit_predict",
             ),
         ]
         if bias_estimation:

hypex/ml/__init__.py

@@ -1,12 +1,12 @@
 from .cupac import CUPACExecutor
-from .faiss import FaissNearestNeighbors
+from .faiss import FaissMLExecutor
 from .model_selection import ModelSelectionExecutor
 from .models import MLModel
 from .stats import ModelStats
 
 __all__ = [
 	"CUPACExecutor",
-	"FaissNearestNeighbors",
+	"FaissMLExecutor",
 	"ModelSelectionExecutor",
 	"MLModel",
 	"ModelStats",