mcp-automl 0.1.5__py3-none-any.whl → 0.1.7__py3-none-any.whl

mcp_automl/server.py

@@ -10,8 +10,8 @@ import argparse
 from pathlib import Path
 from mcp.server.fastmcp import FastMCP, Context
 from mcp.types import PromptMessage, TextContent
-from pycaret.classification import setup as setup_clf, compare_models as compare_models_clf, pull as pull_clf, save_model as save_model_clf, load_model as load_model_clf, predict_model as predict_model_clf, get_config as get_config_clf
-from pycaret.regression import setup as setup_reg, compare_models as compare_models_reg, pull as pull_reg, save_model as save_model_reg, load_model as load_model_reg, predict_model as predict_model_reg, get_config as get_config_reg
+from pycaret.classification import setup as setup_clf, compare_models as compare_models_clf, pull as pull_clf, save_model as save_model_clf, load_model as load_model_clf, predict_model as predict_model_clf, get_config as get_config_clf, tune_model as tune_model_clf, finalize_model as finalize_model_clf
+from pycaret.regression import setup as setup_reg, compare_models as compare_models_reg, pull as pull_reg, save_model as save_model_reg, load_model as load_model_reg, predict_model as predict_model_reg, get_config as get_config_reg, tune_model as tune_model_reg, finalize_model as finalize_model_reg
 
 # Configure logging
 logging.basicConfig(
@@ -363,12 +363,31 @@ def _train_classifier_sync(run_id: str, data_path: str, target_column: str, igno
         best_model = best_model[0]
     results = pull_clf()
     
-    # Extract feature importances
+    # Tune Model
+    logger.info("Tuning best model with Optuna...")
+    try:
+         best_model = tune_model_clf(best_model, optimize=optimize, search_library="optuna", n_trials=10)
+         results = pull_clf()
+    except Exception as e:
+         logger.warning(f"Tuning failed: {e}. Proceeding with untuned model.")
+
+    
+    # Extract feature importances (from the potentially tuned model)
     feature_importances = _get_feature_importances(best_model, get_config_clf)
     
-    # Evaluate on holdout (test_data or split)
-    predict_model_clf(best_model)
-    test_results = pull_clf()
+    # Evaluate on holdout (test_data_path)
+    test_results = None
+    if test_data_path:
+         logger.info("Evaluating on provided test data...")
+         predict_model_clf(best_model)
+         test_results = pull_clf()
+    
+    # Finalize Model
+    logger.info("Finalizing model on all data...")
+    try:
+        best_model = finalize_model_clf(best_model)
+    except Exception as e:
+        logger.warning(f"Finalization failed: {e}. Saving non-finalized model.")
     
     metadata = {
         "data_path": data_path,
@@ -543,12 +562,30 @@ def _train_regressor_sync(run_id: str, data_path: str, target_column: str, ignor
         best_model = best_model[0]
     results = pull_reg()
     
+    # Tune Model
+    logger.info("Tuning best model with Optuna...")
+    try:
+         best_model = tune_model_reg(best_model, optimize=optimize, search_library="optuna", n_trials=10)
+         results = pull_reg()
+    except Exception as e:
+         logger.warning(f"Tuning failed: {e}. Proceeding with untuned model.")
+    
     # Extract feature importances
     feature_importances = _get_feature_importances(best_model, get_config_reg)
     
     # Evaluate on holdout
-    predict_model_reg(best_model)
-    test_results = pull_reg()
+    test_results = None
+    if test_data_path:
+        logger.info("Evaluating on provided test data...")
+        predict_model_reg(best_model)
+        test_results = pull_reg()
+
+    # Finalize Model
+    logger.info("Finalizing model on all data...")
+    try:
+        best_model = finalize_model_reg(best_model)
+    except Exception as e:
+        logger.warning(f"Finalization failed: {e}. Saving non-finalized model.")
     
     metadata = {
         "data_path": data_path,

{mcp_automl-0.1.5.dist-info → mcp_automl-0.1.7.dist-info}/METADATA

@@ -1,14 +1,15 @@
 Metadata-Version: 2.4
 Name: mcp-automl
-Version: 0.1.5
+Version: 0.1.7
 Summary: MCP server for end-to-end machine learning
 Author-email: ke <idea7766@gmail.com>
 License-File: LICENSE
 Requires-Python: <3.12,>=3.10
-Requires-Dist: duckdb>=1.4.3
+Requires-Dist: duckdb[all]>=1.4.3
 Requires-Dist: joblib<1.4
 Requires-Dist: mcp>=1.21.2
 Requires-Dist: pandas<2.2.0
+Requires-Dist: pyarrow>=23.0.0
 Requires-Dist: pycaret>=3.0.0
 Requires-Dist: scikit-learn<1.4
 Requires-Dist: tabulate>=0.9.0

mcp_automl-0.1.7.dist-info/RECORD

@@ -0,0 +1,8 @@
+mcp_automl/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mcp_automl/__main__.py,sha256=ncTRY5zgcNOS7JcLnCVhPd9KsxjyREC245P2eo33BuI,74
+mcp_automl/server.py,sha256=cyxXkNnDc9qHyp3hBUvrxXwI1Ww67XU7uTLrqCw-po8,43972
+mcp_automl-0.1.7.dist-info/METADATA,sha256=z4rojoKfG5twMzP616RdOrOE5zHKN0AHoAlyaS2_zBY,3674
+mcp_automl-0.1.7.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+mcp_automl-0.1.7.dist-info/entry_points.txt,sha256=7QuAE_HatGpFE7Ul7hqNHmpaMf0Ug86aFkaCXofjhLg,54
+mcp_automl-0.1.7.dist-info/licenses/LICENSE,sha256=MLSjoBTdoZnIBt__JNue8D2KnKqNg7Mt_LYWn9vNNws,1065
+mcp_automl-0.1.7.dist-info/RECORD,,

mcp_automl-0.1.5.dist-info/RECORD

@@ -1,8 +0,0 @@
-mcp_automl/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-mcp_automl/__main__.py,sha256=ncTRY5zgcNOS7JcLnCVhPd9KsxjyREC245P2eo33BuI,74
-mcp_automl/server.py,sha256=rk8mQFSm-Y-p5-6DqvPkdiUN6WQrC7jXGXTb4Byedgw,42435
-mcp_automl-0.1.5.dist-info/METADATA,sha256=iB054Xp5j3YMPsr2CmpuOuq2kyQaq20stDRsrwcJBfY,3638
-mcp_automl-0.1.5.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-mcp_automl-0.1.5.dist-info/entry_points.txt,sha256=7QuAE_HatGpFE7Ul7hqNHmpaMf0Ug86aFkaCXofjhLg,54
-mcp_automl-0.1.5.dist-info/licenses/LICENSE,sha256=MLSjoBTdoZnIBt__JNue8D2KnKqNg7Mt_LYWn9vNNws,1065
-mcp_automl-0.1.5.dist-info/RECORD,,