From 26cf3f4c01b708f186853c2b388cc3285028d8d8 Mon Sep 17 00:00:00 2001 From: Pawel Date: Thu, 29 May 2025 01:00:20 +0300 Subject: [PATCH] Add deep learning wrappers with MCP integration --- README.md | 6 +- faster_llm/ML/__init__.py | 17 +++++ faster_llm/ML/keras_model.py | 53 ++++++++++++++++ faster_llm/ML/pytorch_lightning_model.py | 36 +++++++++++ faster_llm/ML/pytorch_model.py | 81 ++++++++++++++++++++++++ tests/test_dl_wrappers.py | 49 ++++++++++++++ 6 files changed, 241 insertions(+), 1 deletion(-) create mode 100644 faster_llm/ML/keras_model.py create mode 100644 faster_llm/ML/pytorch_lightning_model.py create mode 100644 faster_llm/ML/pytorch_model.py create mode 100644 tests/test_dl_wrappers.py diff --git a/README.md b/README.md index 887d22f..d4bc870 100644 --- a/README.md +++ b/README.md @@ -13,8 +13,9 @@ - Helpers for passing results to LLM frameworks - Feature selection, model evaluation and hyperparameter tuning tools - Wrappers around common scikit‑learn and XGBoost models +- Lightweight wrappers for Keras, PyTorch and PyTorch Lightning - Simple generator for producing fake data for experiments -- MCP client for forwarding messages to LLM services +- MCP client for forwarding messages to LLM services and AI agents ## Install from PyPI @@ -58,6 +59,9 @@ model on the included Iris dataset and prints evaluation metrics via the python examples/classification_example.py ``` +The same helper can forward metrics from Keras, PyTorch or PyTorch Lightning +models to AI agents by using the library's MCP integration. + ## Bug Reports & Feature Requests Please open an issue on GitHub if you encounter a bug or have a feature request. diff --git a/faster_llm/ML/__init__.py b/faster_llm/ML/__init__.py index df92005..1af4926 100644 --- a/faster_llm/ML/__init__.py +++ b/faster_llm/ML/__init__.py @@ -1,3 +1,20 @@ # -*- coding: utf-8 -*- """Machine learning model wrappers.""" +from .classification import Model as ClassificationModel +from .regression import Model as RegressionModel +from .time_series import TimeSeries +from .clasterization import ClusterModel +from .keras_model import KerasModel +from .pytorch_model import PyTorchModel +from .pytorch_lightning_model import PyTorchLightningModel + +__all__ = [ + "ClassificationModel", + "RegressionModel", + "TimeSeries", + "ClusterModel", + "KerasModel", + "PyTorchModel", + "PyTorchLightningModel", +] diff --git a/faster_llm/ML/keras_model.py b/faster_llm/ML/keras_model.py new file mode 100644 index 0000000..83100c4 --- /dev/null +++ b/faster_llm/ML/keras_model.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +"""Minimal wrapper for Keras style models.""" + +from dataclasses import dataclass, field +from typing import Any + +import numpy as np +import pandas as pd +from sklearn.model_selection import train_test_split +from sklearn.metrics import accuracy_score + +from faster_llm.LLM import send_to_llm + + +@dataclass +class KerasModel: + """Train and evaluate a Keras-like model and optionally report metrics.""" + + model: Any + X: pd.DataFrame + y: pd.Series + epochs: int = 1 + test_size: float = 0.2 + send_to_llm_flag: bool = False + server_url: str | None = None + y_pred: np.ndarray | None = field(init=False, default=None) + + def __post_init__(self) -> None: + self.X_train, self.X_test, self.y_train, self.y_test = train_test_split( + self.X, self.y, test_size=self.test_size + ) + self.model.fit(self.X_train, self.y_train, epochs=self.epochs, verbose=0) + self.y_pred = np.asarray(self.model.predict(self.X_test)) + metrics = self._compute_metrics() + if self.send_to_llm_flag: + self.send_metrics_to_llm(metrics) + + def _compute_metrics(self) -> dict: + preds = self.y_pred + if preds is None: + return {} + preds = preds.squeeze() + if preds.ndim > 1: + preds = preds.argmax(axis=1) + else: + preds = (preds > 0.5).astype(int) + return {"accuracy": accuracy_score(self.y_test, preds)} + + def send_metrics_to_llm(self, metrics: dict | None = None) -> None: + if metrics is None: + metrics = self._compute_metrics() + send_to_llm(f"Keras metrics: {metrics}", server_url=self.server_url) diff --git a/faster_llm/ML/pytorch_lightning_model.py b/faster_llm/ML/pytorch_lightning_model.py new file mode 100644 index 0000000..be75dea --- /dev/null +++ b/faster_llm/ML/pytorch_lightning_model.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +"""Minimal wrapper for PyTorch Lightning style models.""" + +from dataclasses import dataclass +from typing import Any + +import pandas as pd +from sklearn.model_selection import train_test_split + +from faster_llm.LLM import send_to_llm +from .pytorch_model import PyTorchModel + + +@dataclass +class PyTorchLightningModel(PyTorchModel): + """Wrapper around a PyTorch Lightning model.""" + + def __post_init__(self) -> None: + # Reuse PyTorchModel behaviour assuming model has `fit` and `predict`. + self.X_train, self.X_test, self.y_train, self.y_test = train_test_split( + self.X, self.y, test_size=self.test_size + ) + if hasattr(self.model, "fit"): + self.model.fit(self.X_train, self.y_train, epochs=self.epochs) + self.y_pred = self.model.predict(self.X_test) + metrics = self._compute_metrics() + if self.send_to_llm_flag: + self.send_metrics_to_llm(metrics) + + def send_metrics_to_llm(self, metrics: dict | None = None) -> None: + if metrics is None: + metrics = self._compute_metrics() + send_to_llm( + f"PyTorch Lightning metrics: {metrics}", server_url=self.server_url + ) diff --git a/faster_llm/ML/pytorch_model.py b/faster_llm/ML/pytorch_model.py new file mode 100644 index 0000000..8868f24 --- /dev/null +++ b/faster_llm/ML/pytorch_model.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +"""Minimal wrapper for PyTorch style models.""" + +from dataclasses import dataclass, field +from typing import Any + +import numpy as np +import pandas as pd +from sklearn.model_selection import train_test_split +from sklearn.metrics import accuracy_score + +from faster_llm.LLM import send_to_llm + + +@dataclass +class PyTorchModel: + """Train and evaluate a PyTorch-like model and optionally report metrics.""" + + model: Any + X: pd.DataFrame + y: pd.Series + epochs: int = 1 + test_size: float = 0.2 + send_to_llm_flag: bool = False + server_url: str | None = None + y_pred: np.ndarray | None = field(init=False, default=None) + + def __post_init__(self) -> None: + self.X_train, self.X_test, self.y_train, self.y_test = train_test_split( + self.X, self.y, test_size=self.test_size + ) + if hasattr(self.model, "fit"): + self.model.fit(self.X_train, self.y_train, epochs=self.epochs) + else: + try: + import torch + self.model.train() + optimizer = getattr(self.model, "optimizer", torch.optim.SGD(self.model.parameters(), lr=0.01)) + criterion = getattr(self.model, "criterion", torch.nn.CrossEntropyLoss()) + X_tensor = torch.tensor(self.X_train.values, dtype=torch.float32) + y_tensor = torch.tensor(self.y_train.values) + for _ in range(self.epochs): + optimizer.zero_grad() + output = self.model(X_tensor) + loss = criterion(output, y_tensor) + loss.backward() + optimizer.step() + except Exception: + raise RuntimeError("PyTorch model training not available") + if hasattr(self.model, "predict"): + preds = self.model.predict(self.X_test) + else: + try: + import torch + self.model.eval() + with torch.no_grad(): + preds = self.model(torch.tensor(self.X_test.values, dtype=torch.float32)) + preds = preds.detach().numpy() + except Exception: + raise RuntimeError("PyTorch model prediction not available") + self.y_pred = np.asarray(preds) + metrics = self._compute_metrics() + if self.send_to_llm_flag: + self.send_metrics_to_llm(metrics) + + def _compute_metrics(self) -> dict: + preds = self.y_pred + if preds is None: + return {} + preds = preds.squeeze() + if preds.ndim > 1: + preds = preds.argmax(axis=1) + else: + preds = (preds > 0.5).astype(int) + return {"accuracy": accuracy_score(self.y_test, preds)} + + def send_metrics_to_llm(self, metrics: dict | None = None) -> None: + if metrics is None: + metrics = self._compute_metrics() + send_to_llm(f"PyTorch metrics: {metrics}", server_url=self.server_url) diff --git a/tests/test_dl_wrappers.py b/tests/test_dl_wrappers.py new file mode 100644 index 0000000..46a36f2 --- /dev/null +++ b/tests/test_dl_wrappers.py @@ -0,0 +1,49 @@ +import pytest +pd = pytest.importorskip("pandas") + +from faster_llm.ML.keras_model import KerasModel +from faster_llm.ML.pytorch_model import PyTorchModel +from faster_llm.ML.pytorch_lightning_model import PyTorchLightningModel + + +class DummyModel: + def fit(self, X, y, epochs=1, verbose=0): + self.fitted = True + + def predict(self, X): + return [0 for _ in range(len(X))] + + +def _run_wrapper(wrapper_cls): + recorded = {} + + def fake_send(message, server_url=None): + recorded["msg"] = message + recorded["url"] = server_url + + df = pd.DataFrame({"a": [1, 2, 3, 4], "b": [5, 6, 7, 8]}) + y = pd.Series([0, 0, 0, 0]) + wrapper_cls = wrapper_cls # rename + + from faster_llm import LLM + original = LLM.send_to_llm + LLM.send_to_llm = fake_send + try: + wrapper_cls(DummyModel(), df, y, send_to_llm_flag=True, server_url="http://host") + finally: + LLM.send_to_llm = original + + assert recorded["url"] == "http://host" + assert "accuracy" in recorded["msg"] + + +def test_keras_wrapper(): + _run_wrapper(KerasModel) + + +def test_pytorch_wrapper(): + _run_wrapper(PyTorchModel) + + +def test_pytorch_lightning_wrapper(): + _run_wrapper(PyTorchLightningModel)