Update app.py

app.py

@@ -1,4 +1,5 @@
 # %%writefile app.py
+# %%writefile app.py
 
 import streamlit as st
 import matplotlib.pyplot as plt
@@ -6,12 +7,15 @@ import torch
 from transformers import AutoTokenizer, DataCollatorWithPadding, AutoModelForSequenceClassification, AdamW
 from datasets import load_dataset
 from evaluate import load as load_metric
-from torch.utils.data import DataLoader, random_split
+from torch.utils.data import DataLoader
 import random
+import warnings
+from collections import OrderedDict
+import flwr as fl
 
 DEVICE = torch.device("cpu")
 
-def load_data(dataset_name, train_size=20, test_size=20):
+def load_data(dataset_name, train_size=20, test_size=20, num_clients=2):
     raw_datasets = load_dataset(dataset_name)
     raw_datasets = raw_datasets.shuffle(seed=42)
     del raw_datasets["unsupervised"]
@@ -25,14 +29,21 @@ def load_data(dataset_name, train_size=20, test_size=20):
     tokenized_datasets = tokenized_datasets.remove_columns("text")
     tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
 
-    train_dataset = tokenized_datasets["train"].select(random.sample(range(len(tokenized_datasets["train"])), train_size))
-    test_dataset = tokenized_datasets["test"].select(random.sample(range(len(tokenized_datasets["test"])), test_size))
+    train_datasets = []
+    test_datasets = []
+
+    for _ in range(num_clients):
+        train_dataset = tokenized_datasets["train"].select(random.sample(range(len(tokenized_datasets["train"])), train_size))
+        test_dataset = tokenized_datasets["test"].select(random.sample(range(len(tokenized_datasets["test"])), test_size))
+        train_datasets.append(train_dataset)
+        test_datasets.append(test_dataset)
 
     data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
-    trainloader = DataLoader(train_dataset, shuffle=True, batch_size=32, collate_fn=data_collator)
-    testloader = DataLoader(test_dataset, batch_size=32, collate_fn=data_collator)
 
-    return trainloader, testloader
+    trainloaders = [DataLoader(ds, shuffle=True, batch_size=32, collate_fn=data_collator) for ds in train_datasets]
+    testloaders = [DataLoader(ds, batch_size=32, collate_fn=data_collator) for ds in test_datasets]
+
+    return trainloaders, testloaders
 
 def train(net, trainloader, epochs):
     optimizer = AdamW(net.parameters(), lr=5e-5)
@@ -62,6 +73,30 @@ def test(net, testloader):
     accuracy = metric.compute()["accuracy"]
     return loss, accuracy
 
+class CustomClient(fl.client.NumPyClient):
+    def __init__(self, net, trainloader, testloader):
+        self.net = net
+        self.trainloader = trainloader
+        self.testloader = testloader
+
+    def get_parameters(self, config):
+        return [val.cpu().numpy() for _, val in self.net.state_dict().items()]
+
+    def set_parameters(self, parameters):
+        params_dict = zip(self.net.state_dict().keys(), parameters)
+        state_dict = OrderedDict({k: torch.Tensor(v) for k, v in params_dict})
+        self.net.load_state_dict(state_dict, strict=True)
+
+    def fit(self, parameters, config):
+        self.set_parameters(parameters)
+        train(self.net, self.trainloader, epochs=1)
+        return self.get_parameters(config={}), len(self.trainloader.dataset), {}
+
+    def evaluate(self, parameters, config):
+        self.set_parameters(parameters)
+        loss, accuracy = test(self.net, self.testloader)
+        return float(loss), len(self.testloader.dataset), {"accuracy": float(accuracy)}
+
 def main():
     st.write("## Federated Learning with Dynamic Models and Datasets for Mobile Devices")
     dataset_name = st.selectbox("Dataset", ["imdb", "amazon_polarity", "ag_news"])
@@ -72,52 +107,39 @@ def main():
     NUM_CLIENTS = st.slider("Number of Clients", min_value=1, max_value=10, value=2)
     NUM_ROUNDS = st.slider("Number of Rounds", min_value=1, max_value=10, value=3)
 
-    trainloader, testloader = load_data(dataset_name)
+    trainloaders, testloaders = load_data(dataset_name, num_clients=NUM_CLIENTS)
 
     if st.button("Start Training"):
         round_losses = []
         round_accuracies = []
 
-        for round_num in range(1, NUM_ROUNDS + 1):
-            st.write(f"## Round {round_num}")
-
-            st.write("### Training Metrics for Each Client")
-            client_losses = []
-            client_accuracies = []
+        clients = [CustomClient(net, trainloaders[i], testloaders[i]) for i in range(NUM_CLIENTS)]
 
-            for client in range(1, NUM_CLIENTS + 1):
-                train_subset, _ = random_split(trainloader.dataset, [len(trainloader.dataset) // NUM_CLIENTS] * NUM_CLIENTS)
-                trainloader_client = DataLoader(train_subset, shuffle=True, batch_size=32, collate_fn=trainloader.collate_fn)
-                train(net, trainloader_client, epochs=1)
-                client_loss, client_accuracy = test(net, testloader)
-                st.write(f"Client {client}: Loss: {client_loss:.4f}, Accuracy: {client_accuracy:.4f}")
-                client_losses.append(client_loss)
-                client_accuracies.append(client_accuracy)
+        def client_fn(cid):
+            return clients[int(cid)]
 
-            avg_client_loss = sum(client_losses) / NUM_CLIENTS
-            avg_client_accuracy = sum(client_accuracies) / NUM_CLIENTS
+        def weighted_average(metrics):
+            accuracies = [num_examples * m["accuracy"] for num_examples, m in metrics]
+            losses = [num_examples * m["loss"] for num_examples, m in metrics]
+            examples = [num_examples for num_examples, _ in metrics]
+            return {"accuracy": sum(accuracies) / sum(examples), "loss": sum(losses) / sum(examples)}
 
-            st.write("### Average Metrics Across All Clients")
-            st.write(f"Average Loss: {avg_client_loss:.4f}, Average Accuracy: {avg_client_accuracy:.4f}")
+        strategy = fl.server.strategy.FedAvg(
+            fraction_fit=1.0,
+            fraction_evaluate=1.0,
+            evaluate_metrics_aggregation_fn=weighted_average,
+        )
 
-            round_losses.append(avg_client_loss)
-            round_accuracies.append(avg_client_accuracy)
+        fl.simulation.start_simulation(
+            client_fn=client_fn,
+            num_clients=NUM_CLIENTS,
+            config=fl.server.ServerConfig(num_rounds=NUM_ROUNDS),
+            strategy=strategy,
+            client_resources={"num_cpus": 1, "num_gpus": 0},
+            ray_init_args={"log_to_driver": False, "num_cpus": 1, "num_gpus": 0}
+        )
 
-            st.write("### Accuracy Over Rounds")
-            plt.plot(range(1, round_num + 1), round_accuracies, marker='o', label="Accuracy")
-            plt.xlabel("Round")
-            plt.ylabel("Accuracy")
-            plt.title("Accuracy Over Rounds")
-            st.pyplot()
-
-            st.write("### Loss Over Rounds")
-            plt.plot(range(1, round_num + 1), round_losses, marker='o', color='red', label="Loss")
-            plt.xlabel("Round")
-            plt.ylabel("Loss")
-            plt.title("Loss Over Rounds")
-            st.pyplot()
-
-            st.success(f"Round {round_num} completed successfully!")
+        st.success(f"Training completed successfully!")
 
     else:
         st.write("Click the 'Start Training' button to start the training process.")
@@ -125,6 +147,7 @@ def main():
 if __name__ == "__main__":
     main()
 
+
 ##ORIGINAL###