Add deep learning wrappers

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.

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",
+]

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)

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
+        )

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)

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)