FedModelKit 0.5.0__tar.gz → 0.6.0__tar.gz

{fedmodelkit-0.5.0 → fedmodelkit-0.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: FedModelKit
-Version: 0.5.0
+Version: 0.6.0
 Summary: This package contains the core components and protocols for creating, managing, and registering federated learning models using MLflow. It provides utilities for defining local learners, aggregation strategies, and integrating them with MLflow for tracking and deployment.
 Author-email: ceresale <alessandro.ceresi@upm.es>
 License-File: LICENSE
@@ -101,7 +101,7 @@ Follow the steps below to install and set up the uploading environment:
 2. Go to the directory from where you want to upload the model and run this command to initialize it (pay attention to files that might get overwritten):
 
     ```bash
-    msi init
+    fmk init
     ```
     A `src` folder will be created where the dependencies files must be stored, an `example.py` script will be created in the main directory with an example of how to upload the model and a `README.md` file with a description of how to use the package functionalities. They are provided with clear documentation on how to define your local learner and aggregation strategy, and how to log the model to the Platform Model Catalogue
 

{fedmodelkit-0.5.0 → fedmodelkit-0.6.0}/README.md

@@ -81,7 +81,7 @@ Follow the steps below to install and set up the uploading environment:
 2. Go to the directory from where you want to upload the model and run this command to initialize it (pay attention to files that might get overwritten):
 
     ```bash
-    msi init
+    fmk init
     ```
     A `src` folder will be created where the dependencies files must be stored, an `example.py` script will be created in the main directory with an example of how to upload the model and a `README.md` file with a description of how to use the package functionalities. They are provided with clear documentation on how to define your local learner and aggregation strategy, and how to log the model to the Platform Model Catalogue
 

{fedmodelkit-0.5.0 → fedmodelkit-0.6.0}/examples/simulation-scikit-model/simulation_example.ipynb

@@ -58,7 +58,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "!msi init app"
+    "!fmk init app"
    ]
   },
   {
@@ -106,7 +106,7 @@
    "outputs": [],
    "source": [
     "%%writefile model_example.py\n",
-    "import FedModelKit as msi\n",
+    "import FedModelKit as fmk\n",
     "\n",
     "# Function defining the local learner\n",
     "def create_local_learner():\n",
@@ -313,11 +313,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import FedModelKit as msi\n",
+    "import FedModelKit as fmk\n",
     "from model_example import create_local_learner as create_ll_check # Use alias here to avoid name conflict when uploading the model\n",
     "from model_example import create_aggregator as create_agg_check # Use alias here to avoid name conflict when uploading the model\n",
     "\n",
-    "msi.FederatedModel(create_local_learner=create_ll_check, \n",
+    "fmk.FederatedModel(create_local_learner=create_ll_check, \n",
     "                    model_name='simple_lr',\n",
     "                    create_aggregator=create_agg_check,\n",
     "                    aggregator_name='custom_aggregator')"
@@ -574,13 +574,13 @@
    "outputs": [],
    "source": [
     "%run model_example.py # Run the model_example.py file to directly get the functions defining the local learner and aggregator \n",
-    "import FedModelKit as msi\n",
+    "import FedModelKit as fmk\n",
     "\n",
     "# Create an instance of the FederatedModel class from the FedModelKit module.\n",
     "# This class handles the federated learning process, including model creation, training, and aggregation.\n",
     "# The create_local_learner and create_aggregator functions is passed as an argument to define the local learner (model)\n",
     "# for each client. Here the functions used are the ones defined in the model_example.py file.\n",
-    "federated_model = msi.FederatedModel(create_local_learner=create_local_learner, # type: ignore # Make sure you already type-checked the function\n",
+    "federated_model = fmk.FederatedModel(create_local_learner=create_local_learner, # type: ignore # Make sure you already type-checked the function\n",
     "                                     model_name='simple_lr',\n",
     "                                     create_aggregator=create_aggregator, # type: ignore # Make sure you already type-checked the function\n",
     "                                     aggregator_name='custom_aggregator')"
@@ -631,7 +631,7 @@
     "# HERE you should produce the information about the whole dataset needed by you model and store it in\n",
     "# the src directory, UNLESS you performed the simulation and the clients already stored it.\n",
     "\n",
-    "msi.submit_fl_model(model=federated_model,\n",
+    "fmk.submit_fl_model(model=federated_model,\n",
     "                    platform_url='http://localhost:5000',  # URL of the local MLflow server\n",
     "                    username='username',  # Username for authentication (If no username is required , put a mock username)\n",
     "                    password='password',  # Password for authentication (if no password is required, put a mock password)\n",

{fedmodelkit-0.5.0 → fedmodelkit-0.6.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "FedModelKit"
-version = "0.5.0"
+version = "0.6.0"
 description = "This package contains the core components and protocols for creating, managing, and registering federated learning models using MLflow. It provides utilities for defining local learners, aggregation strategies, and integrating them with MLflow for tracking and deployment."
 readme = "README.md"
 authors = [

{fedmodelkit-0.5.0 → fedmodelkit-0.6.0}/src/FedModelKit/cli.py

@@ -62,7 +62,7 @@ def create_structure(exp_name: str = "new_experiment") -> None:
 
 '''def main():
     if len(sys.argv) < 2:
-        print("Usage: msi <command> [options]")
+        print("Usage: fmk <command> [options]")
         sys.exit(1)
 
     command = sys.argv[1]

fedmodelkit-0.6.0/src/FedModelKit/templates/client_app_template.py

@@ -0,0 +1,49 @@
+from flwr.client import ClientApp, NumPyClient
+from flwr.common import Context
+from loguru import logger
+
+from EXPERIMENT_NAME.task import (
+    Net,
+    evaluate,
+    get_weights,
+    load_flwr_data,
+    set_weights,
+    train,
+)
+
+
+class FlowerClient(NumPyClient):
+    def __init__(self, net, trainloader, testloader):
+        self.net = net
+        self.trainloader = trainloader
+        self.testloader = testloader
+
+    def fit(self, parameters, config):
+        set_weights(self.net, parameters)
+        train(self.net, self.trainloader)
+        return get_weights(self.net), len(self.trainloader), {}
+
+    def evaluate(self, parameters, config):
+        set_weights(self.net, parameters)
+        loss, accuracy = evaluate(self.net, self.testloader)
+        return loss, len(self.testloader), {"accuracy": accuracy}
+
+
+def client_fn(context: Context):
+    from EXPERIMENT_NAME.task import load_syftbox_dataset
+    from syft_flwr.utils import run_syft_flwr
+
+    if not run_syft_flwr():
+        logger.info("Running flwr locally")
+        train_loader, test_loader = load_flwr_data(
+            partition_id=context.node_config["partition-id"],
+            num_partitions=context.node_config["num-partitions"],
+        )
+    else:
+        logger.info("Running with syft_flwr")
+        train_loader, test_loader = load_syftbox_dataset()
+    net = Net()
+    return FlowerClient(net, train_loader, test_loader).to_client()
+
+
+app = ClientApp(client_fn=client_fn)

fedmodelkit-0.6.0/src/FedModelKit/templates/server_app_template.py

@@ -0,0 +1,38 @@
+"""fltabular: Flower Example on Adult Census Income Tabular Dataset."""
+
+from flwr.common import Context, ndarrays_to_parameters
+from flwr.server import ServerApp, ServerAppComponents, ServerConfig
+
+from EXPERIMENT_NAME.task import Net, get_weights
+
+
+def weighted_average(metrics):
+    accuracies = [num_examples * m["accuracy"] for num_examples, m in metrics]
+    examples = [num_examples for num_examples, _ in metrics]
+
+    return {"accuracy": sum(accuracies) / sum(examples)}
+
+
+def server_fn(context: Context) -> ServerAppComponents:
+    net = Net()
+    params = ndarrays_to_parameters(get_weights(net))
+
+    from pathlib import Path
+
+    from syft_flwr.strategy import FedAvgWithModelSaving
+
+    strategy = FedAvgWithModelSaving(
+        save_path=Path(__file__).parent.parent.parent / "weights",
+        fraction_fit=1.0,
+        fraction_evaluate=1.0,
+        min_available_clients=2,
+        initial_parameters=params,
+        evaluate_metrics_aggregation_fn=weighted_average,
+    )
+    num_rounds = context.run_config["num-server-rounds"]
+    config = ServerConfig(num_rounds=num_rounds)
+
+    return ServerAppComponents(config=config, strategy=strategy)
+
+
+app = ServerApp(server_fn=server_fn)

fedmodelkit-0.6.0/src/FedModelKit/templates/task_template.py

@@ -0,0 +1,243 @@
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from flwr_datasets import FederatedDataset
+from flwr_datasets.partitioner import IidPartitioner
+from imblearn.over_sampling import SMOTE
+from loguru import logger
+from pandas import DataFrame
+from sklearn.model_selection import train_test_split
+from sklearn.preprocessing import StandardScaler
+from torch.utils.data import DataLoader, TensorDataset
+
+
+def get_device():
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    elif torch.backends.mps.is_available():
+        return torch.device("mps")
+    elif torch.xpu.is_available():
+        return torch.device("xpu")
+    else:
+        return torch.device("cpu")
+
+
+DEVICE = get_device()
+logger.info(f"Using device: {DEVICE}")
+
+
+class Net(nn.Module):
+    def __init__(self, input_dim=6):
+        super(Net, self).__init__()
+        # First layer with more units and batch normalization
+        self.layer1 = nn.Sequential(
+            nn.Linear(input_dim, 32),  # Increased from 20 to 32
+            nn.BatchNorm1d(32),  # Added batch normalization
+            nn.LeakyReLU(0.1),  # LeakyReLU instead of ReLU
+            nn.Dropout(0.2),  # Increased dropout
+        )
+
+        # Second layer with more units
+        self.layer2 = nn.Sequential(
+            nn.Linear(32, 24),  # Increased from 14 to 24
+            nn.BatchNorm1d(24),  # Added batch normalization
+            nn.LeakyReLU(0.1),
+            nn.Dropout(0.25),
+        )
+
+        # Third layer
+        self.layer3 = nn.Sequential(
+            nn.Linear(24, 16), nn.BatchNorm1d(16), nn.LeakyReLU(0.1)
+        )
+
+        # Output layer
+        self.output_layer = nn.Sequential(nn.Linear(16, 1), nn.Sigmoid())
+
+    def forward(self, x):
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.output_layer(x)
+        return x
+
+
+def dataset_processing(
+    train_df: DataFrame, test_df: DataFrame
+) -> tuple[DataLoader, DataLoader]:
+    def preprocess_df(df: DataFrame) -> DataFrame:
+        columns_to_drop = ["SkinThickness", "Insulin"]
+        df_new: DataFrame = df.drop(columns_to_drop, axis=1)
+
+        # Calculate mean and median (excluding zeros)
+        mean_glucose = df_new[df_new["Glucose"] != 0]["Glucose"].mean()
+        median_bmi = df_new[df_new["BMI"] != 0]["BMI"].median()
+        median_bp = df_new[df_new["BloodPressure"] != 0]["BloodPressure"].median()
+
+        # Replace zeros values with mean/median
+        df_new.replace(
+            {
+                "Glucose": {0: mean_glucose},
+                "BMI": {0: median_bmi},
+                "BloodPressure": {0: median_bp},
+            },
+            inplace=True,
+        )
+
+        return df_new
+
+    # Preprocess both datasets
+    train_processed = preprocess_df(train_df)
+    test_processed = preprocess_df(test_df)
+
+    # Split features and labels for both sets
+    X_train = train_processed.values[:, :6]
+    y_train = train_processed.values[:, 6:]
+    X_test = test_processed.values[:, :6]
+    y_test = test_processed.values[:, 6:]
+
+    from collections import Counter
+
+    def get_minority_class_count(y):
+        return min(Counter(y.flatten()).values())
+
+    minority_count = get_minority_class_count(y_train)
+    k_neighbors = min(5, minority_count - 1) if minority_count > 1 else 1
+
+    # Resample the training data to fix the class imbalance
+    smote = SMOTE(random_state=42, k_neighbors=k_neighbors)
+    X_train_resampled, y_train_resampled = smote.fit_resample(X_train, y_train)
+
+    # Scale the data to have zero mean and unit variance
+    scaler = StandardScaler()
+    X_train_resampled = scaler.fit_transform(X_train_resampled)
+    X_test = scaler.transform(X_test)
+
+    # Convert numpy arrays to PyTorch tensors
+    X_train_tensor = torch.FloatTensor(X_train_resampled)
+    y_train_tensor = torch.FloatTensor(y_train_resampled).reshape(
+        -1, 1
+    )  # Add this reshape
+    X_test_tensor = torch.FloatTensor(X_test)
+    y_test_tensor = torch.FloatTensor(y_test).reshape(-1, 1)
+
+    # Create datasets and dataloaders
+    train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
+    test_dataset = TensorDataset(X_test_tensor, y_test_tensor)
+
+    train_loader = DataLoader(dataset=train_dataset, batch_size=10, shuffle=True)
+    test_loader = DataLoader(
+        dataset=test_dataset, batch_size=len(test_dataset), shuffle=False
+    )
+
+    return train_loader, test_loader
+
+
+def load_syftbox_dataset() -> tuple[DataLoader, DataLoader]:
+    import pandas as pd
+
+    from syft_flwr.utils import get_syftbox_dataset_path
+
+    data_dir = get_syftbox_dataset_path()
+    logger.info(f"Loading dataset from {data_dir}")
+
+    train_df = pd.read_csv(data_dir / "train.csv")
+    test_df = pd.read_csv(data_dir / "test.csv")
+
+    return dataset_processing(train_df, test_df)
+
+
+fds = None  # Cache FederatedDataset
+
+
+def load_flwr_data(
+    partition_id: int, num_partitions: int
+) -> tuple[DataLoader, DataLoader]:
+    """
+    Load the `fl-diabetes-prediction` dataset to memory
+    """
+    global fds
+    if fds is None:
+        partitioner = IidPartitioner(num_partitions=num_partitions)
+        fds = FederatedDataset(
+            dataset="khoaguin/pima-indians-diabetes-database",
+            partitioners={"train": partitioner},
+        )
+
+    partition: DataFrame = fds.load_partition(partition_id, "train").with_format(
+        "pandas"
+    )[:]
+    train_df, test_df = train_test_split(partition, test_size=0.2, random_state=95)
+
+    return dataset_processing(train_df, test_df)
+
+
+def train(model, train_loader, local_epochs=1):
+    criterion = nn.BCELoss()
+    optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=0.0005)
+    history = {"train_loss": [], "train_acc": []}
+    model.to(DEVICE)
+
+    for epoch in range(local_epochs):
+        model.train()
+        running_loss = 0.0
+        correct = 0
+        total = 0
+
+        for inputs, labels in train_loader:
+            inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)
+
+            optimizer.zero_grad()
+            outputs = model(inputs)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+
+            running_loss += loss.item() * inputs.size(0)
+            predicted = (outputs > 0.5).float()
+            total += labels.size(0)
+            correct += (predicted == labels).sum().item()
+
+        epoch_loss = running_loss / len(train_loader.dataset)
+        epoch_acc = correct / total
+        history["train_loss"].append(epoch_loss)
+        history["train_acc"].append(epoch_acc)
+
+    return history
+
+
+def evaluate(model, data_loader):
+    model.to(DEVICE)
+    model.eval()
+    criterion = nn.BCELoss()
+
+    running_loss = 0.0
+    correct = 0
+    total = 0
+
+    with torch.no_grad():
+        for inputs, labels in data_loader:
+            inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)
+            outputs = model(inputs)
+            loss = criterion(outputs, labels)
+            running_loss += loss.item() * inputs.size(0)
+            predicted = (outputs > 0.5).float()
+            total += labels.size(0)
+            correct += (predicted == labels).sum().item()
+
+    epoch_loss = running_loss / len(data_loader.dataset)
+    epoch_acc = correct / total
+
+    return epoch_loss, epoch_acc
+
+
+def set_weights(model, parameters):
+    params_dict = zip(model.state_dict().keys(), parameters)
+    state_dict = OrderedDict({k: torch.tensor(v) for k, v in params_dict})
+    model.load_state_dict(state_dict, strict=True)
+
+
+def get_weights(model):
+    ndarrays = [val.cpu().numpy() for _, val in model.state_dict().items()]
+    return ndarrays

fedmodelkit-0.5.0/src/FedModelKit/templates/client_app_template.py

@@ -1,118 +0,0 @@
-
-import numpy as np
-import pandas as pd
-import pickle
-from pathlib import Path
-from flwr.client import ClientApp
-from flwr.common import Message, Context
-from flwr.common.record import RecordSet, MetricsRecord, ConfigsRecord
-from sklearn.preprocessing import OneHotEncoder
-import FedModelKit as msi
-
-from model_example import create_local_learner #type: ignore[import]
-from load_data import load_data #type: ignore[import]
-
-
-# Initialize the Flower ClientApp
-app = ClientApp()
-
-@app.query()
-def query(msg: Message, ctx: Context) -> Message:
-    """
-    Query function to be executed by the Flower client. This function handles the
-    initial configuration sent by the server.
-    """
-
-    # Retrieve the configuration sent by the server
-    fancy_config = msg.content.configs_records['fancy_config']
-
-    # Load the client split data using the load_data function
-    data = load_data(fancy_config['num_clients'], fancy_config['client_id'])
-
-    # Instantiate the federated model
-    federated_model = msi.FederatedModel(create_local_learner=create_local_learner, model_name='simple_lr')
-
-    # Store the local learner and the data split in the context
-    # To store in context other objects, you can use ctx.state.<object_name> = <object>
-    ctx.state.local_learner = federated_model.create_local_learner()
-    ctx.state.data = data
-    ctx.state.undestand = "che succede"
-
-    return msg.create_reply(RecordSet())
-
-@app.train()
-def train(msg: Message, ctx: Context):
-    """
-    Train function to be executed by the Flower client.
-    This function handles the training of the local model using the data provided.
-    """
-
-    # Retrieve the local learner and the client split from the context
-    local_learner = ctx.state.local_learner
-    data = ctx.state.data
-
-    # Retrieve configuration sent by the server - example
-    #fancy_config = msg.content.configs_records['fancy_config']
-    #local_epochs = fancy_config['local_epochs']
-
-    # Retrieve the model parameters sent by the server
-    fancy_parameters = msg.content.parameters_records['fancy_model']
-    local_learner.set_parameters(fancy_parameters)    
-
-    # Prepare the data using the local learner
-    local_learner.prepare_data(data)
-
-    # Perform local training and obtain training metrics
-    train_metrics = local_learner.train_round()
-
-    # Retrieve the trained model parameters
-    new_parameters_records = local_learner.get_parameters()
-    assert ctx.state.undestand.startswith("che"), "The context state is not being stored correctly"
-
-    # Construct a reply message carrying updated model parameters and generated metrics
-    reply_content = RecordSet()
-    reply_content.parameters_records['fancy_model_returned'] = new_parameters_records
-    reply_content.metrics_records['train_metrics'] = train_metrics
-
-    # Store the metrics and the local learner in the context for future reference
-    ctx.state.metrics_records['prev'] = train_metrics
-    ctx.state.local_learner =  local_learner
-
-    # Return the reply message to the server
-    return msg.create_reply(reply_content)
-
-@app.evaluate()
-def eval(msg: Message, ctx: Context):
-    """
-    Evaluate function to be executed by the Flower client.
-    This function handles the evaluation of the local model using the data provided.
-    """
-
-    # Retrieve the local learner and the client split from the context
-    local_learner = ctx.state.local_learner
-    data = ctx.state.data
-
-    # Retrieve configuration sent by the server - example
-    #fancy_config = msg.content.configs_records['fancy_config']
-    #local_epochs = fancy_config['local_epochs']
-
-    # Retrieve the model parameters sent by the server
-    fancy_parameters = msg.content.parameters_records['fancy_model']
-    local_learner.set_parameters(fancy_parameters)
-
-    # Prepare the data using the local learner
-    local_learner.prepare_data(data)
-
-    # Evaluate the model and obtain evaluation metrics
-    eval_metrics = local_learner.evaluate()
-
-    # Construct a reply message with evaluation metrics
-    reply_content = RecordSet()
-    reply_content.metrics_records['eval_metrics'] = eval_metrics
-
-    # Store the metrics and the local learner in the context for future reference
-    ctx.state.metrics_records['prev'] = eval_metrics
-    ctx.state.local_learner =  local_learner
-
-    # Return the reply message to the server
-    return msg.create_reply(reply_content)

fedmodelkit-0.5.0/src/FedModelKit/templates/server_app_template.py

@@ -1,204 +0,0 @@
-
-from typing import List
-import time
-
-import flwr as fl
-from flwr.common import (
-    Context,
-    NDArrays,
-    Message,
-    MessageType,
-    Metrics,
-    RecordSet,
-    ConfigsRecord,
-    DEFAULT_TTL,
-)
-from flwr.server import Driver
-
-import FedModelKit as msi
-
-from model_example import create_local_learner #type: ignore[import]
-
-
-# Run via `flower-server-app server:app`
-app = fl.server.ServerApp()
-
-
-
-
-@app.main()
-def main(driver: Driver, context: Context) -> None:
-    """
-    Main function to run the federated learning server.
-
-    Structure:
-    - Send a query message to clients for creating the local learner and loading the data
-    - Start global epochs loop for training and evaluation
-        - Send training messages to clients
-        - Aggregate parameters received from clients
-        - Send evaluation messages to clients
-        - Aggregate evaluation metrics
-    """
-    print("Starting test run")
-
-    # Get node IDs of connected clients
-    node_ids = driver.get_node_ids()
-
-    # Initialize the federated model
-    federated_model = msi.FederatedModel(create_local_learner=create_local_learner,
-                                        model_name='simple_lr')
-    global_model = federated_model.create_local_learner()
-    aggregation_strategy = federated_model.create_aggregator()
-
-    # Send a query message to clients for creating the local learner and loading the data
-    messages = []
-    for idx, node_id in enumerate(node_ids):
-        # Create messages to send to clients
-        recordset = RecordSet()
-
-        # Add a config with information to send the client for the query
-        recordset.configs_records["fancy_config"] = ConfigsRecord({"num_clients": len(node_ids), "client_id": idx})
-
-        # Create a query message for each client
-        message = driver.create_message(
-            content=recordset,
-            message_type=MessageType.QUERY,
-            dst_node_id=node_id,
-            group_id=str(1),
-            ttl=DEFAULT_TTL,
-        )
-        messages.append(message)
-
-    # Send training messages to clients
-    message_ids = driver.push_messages(messages)
-    print(f"Pushed {len(list(message_ids))} messages: {message_ids}")
-
-    # Wait for results from clients
-    message_ids = [message_id for message_id in message_ids if message_id != ""]
-    all_replies: List[Message] = []
-    while True:
-        replies = driver.pull_messages(message_ids=message_ids)
-        print(f"Got {len(list(replies))} results")
-        all_replies += replies
-        if len(all_replies) == len(message_ids):
-            break
-        time.sleep(12)
-
-    # Filter out messages with errors
-    all_replies = [
-        msg
-        for msg in all_replies
-        if msg.has_content()
-    ]
-    print(f"Received {len(all_replies)} answers")
-    
-
-    # Run federated training and evaluation for a fixed number of rounds
-    for server_round in range(3):
-        print(f"Commencing server train and evaluation round {server_round + 1}")
-
-        messages = []
-        for idx, node_id in enumerate(node_ids):
-            # Create messages to send to clients
-            recordset = RecordSet()
-
-            # Add model parameters to record
-            recordset.parameters_records["fancy_model"] = global_model.get_parameters()
-            # Add a config with information to send the client for training
-            recordset.configs_records["fancy_config"] = ConfigsRecord({"local_epochs": 3})
-
-            # Create a training message for each client
-            message = driver.create_message(
-                content=recordset,
-                message_type=MessageType.TRAIN,
-                dst_node_id=node_id,
-                group_id=str(server_round),
-                ttl=DEFAULT_TTL,
-            )
-            messages.append(message)
-
-        # Send training messages to clients
-        message_ids = driver.push_messages(messages)
-        print(f"Pushed {len(list(message_ids))} messages: {message_ids}")
-
-        # Wait for results from clients
-        message_ids = [message_id for message_id in message_ids if message_id != ""]
-        all_replies: List[Message] = []
-        while True:
-            replies = driver.pull_messages(message_ids=message_ids)
-            print(f"Got {len(list(replies))} results")
-            all_replies += replies
-            if len(all_replies) == len(message_ids):
-                break
-            time.sleep(12)
-
-        # Filter out messages with errors
-        all_replies = [
-            msg
-            for msg in all_replies
-            if msg.has_content()
-        ]
-        print(f"Received {len(all_replies)} results")
-
-        # Print metrics received from clients
-        for reply in all_replies:
-            print(reply.content.metrics_records)
-
-        # Aggregate parameters received from clients
-        parameter_records_list = [reply.content.parameters_records["fancy_model_returned"] for reply in all_replies]
-        new_parameter_record = aggregation_strategy.aggregate_parameters(parameter_records_list)
-        global_model.set_parameters(new_parameter_record)
-
-        # Evaluate the updated global model
-        messages = []
-        for idx, node_id in enumerate(node_ids):
-            # Create evaluation messages for clients
-            recordset = RecordSet()
-
-            # Add updated model parameters to record
-            recordset.parameters_records["fancy_model"] = new_parameter_record
-            # Add a config with information to send the client for evaluation
-            recordset.configs_records["fancy_config"] = ConfigsRecord({"local_epochs": 3})
-
-            # Create an evaluation message for each client
-            message = driver.create_message(
-                content=recordset,
-                message_type=MessageType.EVALUATE,
-                dst_node_id=node_id,
-                group_id=str(server_round),
-                ttl=DEFAULT_TTL,
-            )
-            messages.append(message)
-
-        # Send evaluation messages to clients
-        message_ids = driver.push_messages(messages)
-        print(f"Pushed {len(list(message_ids))} messages: {message_ids}")
-
-        # Wait for evaluation results from clients
-        message_ids = [message_id for message_id in message_ids if message_id != ""]
-        all_replies: List[Message] = []
-        while True:
-            replies = driver.pull_messages(message_ids=message_ids)
-            print(f"Got {len(list(replies))} results")
-            all_replies += replies
-            if len(all_replies) == len(message_ids):
-                break
-            time.sleep(3)
-
-        # Filter out messages with errors
-        all_replies = [
-            msg
-            for msg in all_replies
-            if msg.has_content()
-        ]
-        print(f"Received {len(all_replies)} results")
-
-        # Print evaluation metrics received from clients
-        metrics_records_list = [reply.content.metrics_records['eval_metrics'] for reply in all_replies]
-        for i, reply in enumerate(all_replies):
-            print(f"Client {i+1} metrics:   ", reply.content.metrics_records['eval_metrics'])
-
-        # Aggregate evaluation metrics
-        print("Aggregated metrics result:   ", aggregation_strategy.aggregate_metrics(metrics_records_list))
-    
-    print("🎉🎉🎉 Successfully completed federated learning run! 🎉🎉🎉")

fedmodelkit-0.5.0/src/FedModelKit/templates/task_template.py

@@ -1,140 +0,0 @@
-"""
-Model upload example script. 
-
-To run this code you need to install pytorch library, you can find it at this link: https://pytorch.org/get-started/locally/
-"""
-from FedModelKit import FederatedModel, submit_fl_model
-
-# Create your custom function defining a PyTorch-based model and returning its instance
-def create_local_learner():
-    import pandas as pd
-    from torch import nn, optim
-    from torch.utils.data import DataLoader, Dataset
-    import torch
-    import flwr
-
-    from src.utils import Utils # Dependency script stored inside the src directory
-    
-    class CustomLocalLearner(nn.Module):
-            def __init__(self, input_size: int) -> None:
-                super(CustomLocalLearner, self).__init__()
-                self.linear = nn.Linear(input_size, 3)
-                self.softmax = nn.Softmax(dim=1)
-                self.loss_fn = nn.CrossEntropyLoss()
-                self.optimizer = optim.Adam(self.parameters(), lr=0.001)
-
-            def _forward(self, x):
-                x = self.linear(x)
-                return self.softmax(x)
-
-            def get_parameters(self) -> flwr.common.ParametersRecord:
-                return Utils.pytorch_to_parameter_record(self.state_dict())
-
-            def set_parameters(self, parameters: flwr.common.ParametersRecord) -> None:
-                self.load_state_dict(Utils.parameters_to_pytorch_state_dict(parameters))
-
-            def prepare_data(self, data: pd.DataFrame) -> None:
-                class IrisDataset(Dataset):
-                    def __init__(self, dataframe: pd.DataFrame) -> None:
-                        self.dataframe = dataframe
-                        self.dataframe.loc[:, "class"] = dataframe.loc[:, "class"].replace(
-                            {"Iris-setosa": 0, "Iris-versicolor": 1, "Iris-virginica": 2}
-                        )
-
-                    def __getitem__(self, idx: int):
-                        x = self.dataframe.iloc[idx, :-1].to_numpy("float32")
-                        y = torch.tensor(self.dataframe.iloc[idx, -1], dtype=torch.long)
-                        return x, y
-
-                    def __len__(self) -> int:
-                        return len(self.dataframe)
-
-                dataset = IrisDataset(data)
-                dataloader = DataLoader(dataset, batch_size=32)
-                self.dataloader = dataloader
-
-            def train_round(self) -> flwr.common.MetricsRecord:
-                loss = torch.tensor(0.0)
-                for batch in self.dataloader:
-                    x, y = batch
-                    self.optimizer.zero_grad()
-                    y_hat = self._forward(x)
-                    loss = self.loss_fn(y_hat, y)
-                    loss.backward()
-                    self.optimizer.step()
-
-                return flwr.common.MetricsRecord({"loss": loss.item()})
-
-            def evaluate(self) -> flwr.common.MetricsRecord:
-                correct = 0
-                total = 0
-                with torch.no_grad():
-                    for batch in self.dataloader:
-                        x, y = batch
-                        y_hat = self._forward(x)
-                        _, predicted = torch.max(y_hat, 1)
-                        total += y.size(0)
-                        correct += (predicted == y).sum().item()
-                return flwr.common.MetricsRecord({"accuracy": correct / total})
-    
-    return CustomLocalLearner(4)
-
-
-# Create your custom function defining an aggregation strategy and returning its instance
-def create_aggregator():
-    from collections import OrderedDict
-    import numpy as np
-    import flwr
-    from typing import Optional
-
-    from src.utils import Utils # Dependency script stored inside the src directory #type: ignore[import]
-
-    class CustomAggregator:
-
-        def aggregate_parameters(self, results: list[flwr.common.ParametersRecord], config: Optional[flwr.common.ConfigsRecord]=None
-            ) -> flwr.common.ParametersRecord:
-                parameters = [Utils.parameters_to_dict(param) for param in results]
-                keys = parameters[0].keys()
-                result = OrderedDict()
-                for key in keys:
-                    # Init array
-                    this_array: np.ndarray = np.zeros_like(parameters[0][key])
-                    for p in parameters:
-                        this_array += p[key]
-                    result[key] = this_array / len(results)
-                return Utils.dict_to_parameter_record(result)
-
-        def aggregate_metrics(self, results: list[flwr.common.MetricsRecord], config: Optional[flwr.common.ConfigsRecord]=None) -> flwr.common.MetricsRecord:
-                keys = results[0].keys()
-                result = OrderedDict()
-                for key in keys:
-                    # Init array
-                    cumsum = 0.0
-                    for m in results:
-                        if not isinstance(m[key], (int, float)):
-                            raise ValueError(
-                                f"flwr.common.MetricsRecord value type not supported: {type(m[key])}"
-                            )
-                        cumsum += m[key]  # type: ignore
-                    result[key] = cumsum / len(results)
-                return flwr.common.MetricsRecord(result)
-    
-    return CustomAggregator()
-
-
-if __name__ == "__main__":
-    # Create FederatedModel with custom aggregator and local learner
-    federated_model = FederatedModel(create_local_learner=create_local_learner,
-                               model_name="your_model_name",
-                               create_aggregator=create_aggregator,
-                               aggregator_name="custom_aggregator")
-
-    # Submit the FederatedModel to the Federated Platform Model Catalogue with credentials and experiment name
-    submit_fl_model(federated_model,
-                    platform_url="your_platform_url",
-                    username="your_username", 
-                    password="your_password",
-                    experiment_name="your_experiment_name",
-                    disease="AML",
-                    trained=False
-                    )